Add the demo (#163)

.dockerignore

@@ -4,3 +4,4 @@ __pycache__
 *.pyc
 tests/test_custom_path.py
 temp_models/
+/.venv/

Dockerfile

@@ -7,8 +7,6 @@ COPY ./README.md .
 COPY ./.python-version .
 COPY ./pocket_tts ./pocket_tts
 
-RUN uv run pocket-tts serve --help && \
+RUN uv run pocket-tts --help
-    rm -rf /root/.cache/uv && \
-    uv run pocket-tts serve --help
 
-CMD ["uv", "run", "pocket-tts", "serve"]
+ENTRYPOINT ["uv", "run", "pocket-tts"]

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

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}"
+  ]
+}

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

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

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

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]

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[
+    text: Annotated[str, typer.Option(help="Text to generate")] = None,
-        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.",
     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()

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"):

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.</textarea>
+            >DEFAULT_TEXT_PROMPT</textarea>
         </div>
 
         <div class="bg-white p-4 rounded-lg shadow">
@@ -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
                 });

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}
+    image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA}
-    build:
-      context: .
     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}
+    image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA}
-    build:
-      context: .
     volumes:
       - hf-cache:/root/.cache/huggingface
+    command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "german", "--quantize"]
     deploy:
       labels:
         - "traefik.enable=true"
@@ -69,19 +68,84 @@ 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:
+  hf-cache:
-  prometheus-data:
-  grafana-data:
 
 networks:
-  agent_network:
-    driver: overlay
-    attachable: true
-    driver_opts:
-      encrypted: "true"
   default:
     driver: overlay
     attachable: true