diff --git a/.dockerignore b/.dockerignore index e858388..0931bae 100644 --- a/.dockerignore +++ b/.dockerignore @@ -4,3 +4,4 @@ __pycache__ *.pyc tests/test_custom_path.py temp_models/ +/.venv/ \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index 43f029e..523273a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -7,8 +7,6 @@ COPY ./README.md . COPY ./.python-version . COPY ./pocket_tts ./pocket_tts -RUN uv run pocket-tts serve --help && \ - rm -rf /root/.cache/uv && \ - uv run pocket-tts serve --help +RUN uv run pocket-tts --help -CMD ["uv", "run", "pocket-tts", "serve"] \ No newline at end of file +ENTRYPOINT ["uv", "run", "pocket-tts"] diff --git a/deploy.sh b/deploy.sh new file mode 100755 index 0000000..b57e357 --- /dev/null +++ b/deploy.sh @@ -0,0 +1,6 @@ +set -ex +export COMMIT_SHA=$(git rev-parse --short HEAD) +docker buildx bake -f docker-bake.hcl --push + +docker -H ssh://root@pocket-tts-api.kyutai.org stack deploy \ + --with-registry-auth -c swarm-config.yaml pocket-tts diff --git a/docker-bake.hcl b/docker-bake.hcl new file mode 100644 index 0000000..2a17a21 --- /dev/null +++ b/docker-bake.hcl @@ -0,0 +1,16 @@ +variable "COMMIT_SHA" { + default = "latest" +} + +group "default" { + targets = ["pocket-tts-server"] +} + +target "pocket-tts-server" { + context = "." + dockerfile = "Dockerfile" + platforms = ["linux/arm64", "linux/amd64"] + tags = [ + "rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA}" + ] +} diff --git a/pocket_tts/config/french_24l.yaml b/pocket_tts/config/french_24l.yaml index 311448a..dfd815a 100644 --- a/pocket_tts/config/french_24l.yaml +++ b/pocket_tts/config/french_24l.yaml @@ -4,7 +4,7 @@ weights_path: hf://kyutai/pocket-tts/languages/french_24l/model.safetensors@3959 weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/french_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc remove_semicolons: true -model_recommended_frames_after_eos: 12 +model_recommended_frames_after_eos: 8 flow_lm: insert_bos_before_voice: true diff --git a/pocket_tts/config/german.yaml b/pocket_tts/config/german.yaml index dbda4de..9ac6414 100644 --- a/pocket_tts/config/german.yaml +++ b/pocket_tts/config/german.yaml @@ -3,6 +3,7 @@ weights_path: hf://kyutai/pocket-tts/languages/german/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/german/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc +remove_semicolons: true flow_lm: insert_bos_before_voice: true diff --git a/pocket_tts/config/german_24l.yaml b/pocket_tts/config/german_24l.yaml index 8d0cd25..17fd263 100644 --- a/pocket_tts/config/german_24l.yaml +++ b/pocket_tts/config/german_24l.yaml @@ -3,6 +3,7 @@ weights_path: hf://kyutai/pocket-tts/languages/german_24l/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/german_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc +remove_semicolons: true flow_lm: insert_bos_before_voice: true diff --git a/pocket_tts/default_parameters.py b/pocket_tts/default_parameters.py index 5027016..e99540b 100644 --- a/pocket_tts/default_parameters.py +++ b/pocket_tts/default_parameters.py @@ -7,3 +7,43 @@ DEFAULT_EOS_THRESHOLD = -4.0 DEFAULT_FRAMES_AFTER_EOS = None # TODO: make this dynamic since english_2026-04 supports bigger chunks MAX_TOKEN_PER_CHUNK = 50 + +DEFAULT_TEXT_FOR_LANGUAGE = { + "english": ( + "Hello world. I am Kyutai's Pocket TTS. " + "I'm fast enough to run on small CPUs. " + "I hope you'll like me." + ), + "french": ( + "Bonjour le monde. Je suis le TTS de poche de Kyutai. " + "Je suis assez rapide pour fonctionner sur de petits CPU. " + "J'espère que vous m'aimerez." + ), + "german": ( + "Hallo Welt. Ich bin Kyutais Pocket TTS. " + "Ich bin schnell genug, um auf kleinen CPUs zu laufen. " + "Ich hoffe, du wirst mich mögen." + ), + "portuguese": ( + "Olá mundo. Eu sou o Pocket TTS da Kyutai. " + "Sou rápido o suficiente para rodar em CPUs pequenas. " + "Espero que você goste de mim." + ), + "italian": ( + "Ciao mondo. Sono il Pocket TTS di Kyutai. " + "Sono abbastanza veloce da funzionare su piccole CPU. " + "Spero che ti piacerò." + ), + "spanish": ( + "Hola mundo. Soy el Pocket TTS de Kyutai. " + "Soy lo suficientemente rápido para funcionar en pequeñas CPU. " + "Espero que te guste." + ), +} + + +def get_default_text_for_language(language: str | None) -> str: + for key, text in DEFAULT_TEXT_FOR_LANGUAGE.items(): + if language is not None and key in language: + return text + return DEFAULT_TEXT_FOR_LANGUAGE[DEFAULT_LANGUAGE] diff --git a/pocket_tts/main.py b/pocket_tts/main.py index f63694e..c040425 100644 --- a/pocket_tts/main.py +++ b/pocket_tts/main.py @@ -11,7 +11,7 @@ import typer import uvicorn from fastapi import FastAPI, File, Form, HTTPException, UploadFile from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import FileResponse, StreamingResponse +from fastapi.responses import HTMLResponse, StreamingResponse from typing_extensions import Annotated from pocket_tts.data.audio import stream_audio_chunks @@ -23,6 +23,7 @@ from pocket_tts.default_parameters import ( DEFAULT_NOISE_CLAMP, DEFAULT_TEMPERATURE, MAX_TOKEN_PER_CHUNK, + get_default_text_for_language, ) from pocket_tts.models.tts_model import TTSModel, export_model_state from pocket_tts.utils.logging_utils import enable_logging @@ -58,11 +59,17 @@ web_app.add_middleware( ) -@web_app.get("/") +@web_app.get("/", response_class=HTMLResponse) async def root(): """Serve the frontend.""" static_path = Path(__file__).parent / "static" / "index.html" - return FileResponse(static_path) + content = static_path.read_text() + # Replace the placeholder with the actual default text prompt + print(str(tts_model.origin)) + content = content.replace( + "DEFAULT_TEXT_PROMPT", get_default_text_for_language(str(tts_model.origin)) + ) + return content @web_app.get("/health") @@ -214,9 +221,7 @@ def serve( @cli_app.command() def generate( - text: Annotated[ - str, typer.Option(help="Text to generate") - ] = "Hello world. I am Kyutai's Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.", + text: Annotated[str, typer.Option(help="Text to generate")] = None, voice: Annotated[ str, typer.Option(help="Path to audio conditioning file (voice to clone)") ] = DEFAULT_AUDIO_PROMPT, @@ -268,6 +273,8 @@ def generate( """Generate speech using Kyutai Pocket TTS.""" log_level = logging.ERROR if quiet else logging.INFO with enable_logging("pocket_tts", log_level): + if text is None: + text = get_default_text_for_language(language) if text == "-": # Read text from stdin text = sys.stdin.read() diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index fd46a71..0159486 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -293,6 +293,10 @@ class TTSModel(nn.Module): if config is None and language is None: language = DEFAULT_LANGUAGE if language is not None: + if language == "french": + raise ValueError( + "The french model is not ready yet, please use 'french_24l' instead." + ) config = CONFIGS_DIR / f"{language}.yaml" config = Path(config) if config.suffix not in (".yaml", ".yml"): diff --git a/pocket_tts/static/index.html b/pocket_tts/static/index.html index d05c88b..efcf3e1 100644 --- a/pocket_tts/static/index.html +++ b/pocket_tts/static/index.html @@ -28,7 +28,7 @@ class="w-full border rounded p-2" placeholder="Enter your text here..." rows="4" - >Hello world. I am Kyutai's Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me. + >DEFAULT_TEXT_PROMPT
@@ -314,7 +314,9 @@ formData.append('voice_url', voiceUrl); } - const response = await fetch('/tts', { + const ttsEndpoint = `${window.location.pathname.replace(/\/+$/, '')}/tts`; + + const response = await fetch(ttsEndpoint, { method: 'POST', body: formData }); @@ -392,4 +394,4 @@ }); - \ No newline at end of file + diff --git a/swarm-config.yaml b/swarm-config.yaml index 4cf7300..18c28e5 100644 --- a/swarm-config.yaml +++ b/swarm-config.yaml @@ -1,3 +1,4 @@ + services: traefik: image: traefik:v3.6.7 @@ -32,11 +33,10 @@ services: - node.role == manager backend-english: - image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server-english:${COMMIT_SHA} - build: - context: . + image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA} volumes: - hf-cache:/root/.cache/huggingface + command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "english", "--quantize"] deploy: labels: - "traefik.enable=true" @@ -51,11 +51,10 @@ services: - "traefik.http.routers.server-english.priority=100" backend-german: - image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server-german:${COMMIT_SHA} - build: - context: . + image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA} volumes: - hf-cache:/root/.cache/huggingface + command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "german", "--quantize"] deploy: labels: - "traefik.enable=true" @@ -69,21 +68,86 @@ services: - "traefik.http.services.server-german.loadbalancer.server.port=80" - "traefik.http.routers.server-german.priority=100" + backend-french: + image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA} + volumes: + - hf-cache:/root/.cache/huggingface + command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "french_24l"] + deploy: + labels: + - "traefik.enable=true" + - "traefik.http.routers.server-french.rule=Host(`pocket-tts-api.kyutai.org`) && PathPrefix(`/french`)" + - "traefik.http.routers.server-french.middlewares=strip-api-french" + - "traefik.http.middlewares.strip-api-french.replacepathregex.regex=^/french/(.*)" + - "traefik.http.middlewares.strip-api-french.replacepathregex.replacement=/$$1" + - "traefik.http.routers.server-french.entrypoints=websecure" + - "traefik.http.routers.server-french.tls=true" + - "traefik.http.routers.server-french.tls.certresolver=letsencrypt" + - "traefik.http.services.server-french.loadbalancer.server.port=80" + - "traefik.http.routers.server-french.priority=100" + + backend-italian: + image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA} + volumes: + - hf-cache:/root/.cache/huggingface + command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "italian", "--quantize"] + deploy: + labels: + - "traefik.enable=true" + - "traefik.http.routers.server-italian.rule=Host(`pocket-tts-api.kyutai.org`) && PathPrefix(`/italian`)" + - "traefik.http.routers.server-italian.middlewares=strip-api-italian" + - "traefik.http.middlewares.strip-api-italian.replacepathregex.regex=^/italian/(.*)" + - "traefik.http.middlewares.strip-api-italian.replacepathregex.replacement=/$$1" + - "traefik.http.routers.server-italian.entrypoints=websecure" + - "traefik.http.routers.server-italian.tls=true" + - "traefik.http.routers.server-italian.tls.certresolver=letsencrypt" + - "traefik.http.services.server-italian.loadbalancer.server.port=80" + - "traefik.http.routers.server-italian.priority=100" + + backend-portuguese: + image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA} + volumes: + - hf-cache:/root/.cache/huggingface + command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "portuguese", "--quantize"] + deploy: + labels: + - "traefik.enable=true" + - "traefik.http.routers.server-portuguese.rule=Host(`pocket-tts-api.kyutai.org`) && PathPrefix(`/portuguese`)" + - "traefik.http.routers.server-portuguese.middlewares=strip-api-portuguese" + - "traefik.http.middlewares.strip-api-portuguese.replacepathregex.regex=^/portuguese/(.*)" + - "traefik.http.middlewares.strip-api-portuguese.replacepathregex.replacement=/$$1" + - "traefik.http.routers.server-portuguese.entrypoints=websecure" + - "traefik.http.routers.server-portuguese.tls=true" + - "traefik.http.routers.server-portuguese.tls.certresolver=letsencrypt" + - "traefik.http.services.server-portuguese.loadbalancer.server.port=80" + - "traefik.http.routers.server-portuguese.priority=100" + + backend-spanish: + image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA} + volumes: + - hf-cache:/root/.cache/huggingface + command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "spanish", "--quantize"] + deploy: + labels: + - "traefik.enable=true" + - "traefik.http.routers.server-spanish.rule=Host(`pocket-tts-api.kyutai.org`) && PathPrefix(`/spanish`)" + - "traefik.http.routers.server-spanish.middlewares=strip-api-spanish" + - "traefik.http.middlewares.strip-api-spanish.replacepathregex.regex=^/spanish/(.*)" + - "traefik.http.middlewares.strip-api-spanish.replacepathregex.replacement=/$$1" + - "traefik.http.routers.server-spanish.entrypoints=websecure" + - "traefik.http.routers.server-spanish.tls=true" + - "traefik.http.routers.server-spanish.tls.certresolver=letsencrypt" + - "traefik.http.services.server-spanish.loadbalancer.server.port=80" + - "traefik.http.routers.server-spanish.priority=100" + volumes: letsencrypt: - portainer_data: - prometheus-data: - grafana-data: + hf-cache: networks: - agent_network: - driver: overlay - attachable: true - driver_opts: - encrypted: "true" default: driver: overlay attachable: true driver_opts: - encrypted: "true" \ No newline at end of file + encrypted: "true"