Add the demo (#163)

This commit is contained in:
Gabriel de Marmiesse
2026-04-20 15:55:33 +02:00
committed by GitHub
parent c9ebef4d1a
commit 5304da453b
12 changed files with 169 additions and 29 deletions
+1
View File
@@ -4,3 +4,4 @@ __pycache__
*.pyc *.pyc
tests/test_custom_path.py tests/test_custom_path.py
temp_models/ temp_models/
/.venv/
+2 -4
View File
@@ -7,8 +7,6 @@ COPY ./README.md .
COPY ./.python-version . COPY ./.python-version .
COPY ./pocket_tts ./pocket_tts 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"]
Executable
+6
View File
@@ -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
+16
View File
@@ -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}"
]
}
+1 -1
View File
@@ -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 weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/french_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc
remove_semicolons: true remove_semicolons: true
model_recommended_frames_after_eos: 12 model_recommended_frames_after_eos: 8
flow_lm: flow_lm:
insert_bos_before_voice: true insert_bos_before_voice: true
+1
View File
@@ -3,6 +3,7 @@
weights_path: hf://kyutai/pocket-tts/languages/german/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 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 weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/german/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc
remove_semicolons: true
flow_lm: flow_lm:
insert_bos_before_voice: true insert_bos_before_voice: true
+1
View File
@@ -3,6 +3,7 @@
weights_path: hf://kyutai/pocket-tts/languages/german_24l/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 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 weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/german_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc
remove_semicolons: true
flow_lm: flow_lm:
insert_bos_before_voice: true insert_bos_before_voice: true
+40
View File
@@ -7,3 +7,43 @@ DEFAULT_EOS_THRESHOLD = -4.0
DEFAULT_FRAMES_AFTER_EOS = None DEFAULT_FRAMES_AFTER_EOS = None
# TODO: make this dynamic since english_2026-04 supports bigger chunks # TODO: make this dynamic since english_2026-04 supports bigger chunks
MAX_TOKEN_PER_CHUNK = 50 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]
+13 -6
View File
@@ -11,7 +11,7 @@ import typer
import uvicorn import uvicorn
from fastapi import FastAPI, File, Form, HTTPException, UploadFile from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse from fastapi.responses import HTMLResponse, StreamingResponse
from typing_extensions import Annotated from typing_extensions import Annotated
from pocket_tts.data.audio import stream_audio_chunks from pocket_tts.data.audio import stream_audio_chunks
@@ -23,6 +23,7 @@ from pocket_tts.default_parameters import (
DEFAULT_NOISE_CLAMP, DEFAULT_NOISE_CLAMP,
DEFAULT_TEMPERATURE, DEFAULT_TEMPERATURE,
MAX_TOKEN_PER_CHUNK, MAX_TOKEN_PER_CHUNK,
get_default_text_for_language,
) )
from pocket_tts.models.tts_model import TTSModel, export_model_state from pocket_tts.models.tts_model import TTSModel, export_model_state
from pocket_tts.utils.logging_utils import enable_logging 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(): async def root():
"""Serve the frontend.""" """Serve the frontend."""
static_path = Path(__file__).parent / "static" / "index.html" 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") @web_app.get("/health")
@@ -214,9 +221,7 @@ def serve(
@cli_app.command() @cli_app.command()
def generate( 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[ voice: Annotated[
str, typer.Option(help="Path to audio conditioning file (voice to clone)") str, typer.Option(help="Path to audio conditioning file (voice to clone)")
] = DEFAULT_AUDIO_PROMPT, ] = DEFAULT_AUDIO_PROMPT,
@@ -268,6 +273,8 @@ def generate(
"""Generate speech using Kyutai Pocket TTS.""" """Generate speech using Kyutai Pocket TTS."""
log_level = logging.ERROR if quiet else logging.INFO log_level = logging.ERROR if quiet else logging.INFO
with enable_logging("pocket_tts", log_level): with enable_logging("pocket_tts", log_level):
if text is None:
text = get_default_text_for_language(language)
if text == "-": if text == "-":
# Read text from stdin # Read text from stdin
text = sys.stdin.read() text = sys.stdin.read()
+4
View File
@@ -293,6 +293,10 @@ class TTSModel(nn.Module):
if config is None and language is None: if config is None and language is None:
language = DEFAULT_LANGUAGE language = DEFAULT_LANGUAGE
if language is not None: 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 = CONFIGS_DIR / f"{language}.yaml"
config = Path(config) config = Path(config)
if config.suffix not in (".yaml", ".yml"): if config.suffix not in (".yaml", ".yml"):
+4 -2
View File
@@ -28,7 +28,7 @@
class="w-full border rounded p-2" class="w-full border rounded p-2"
placeholder="Enter your text here..." placeholder="Enter your text here..."
rows="4" 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>
<div class="bg-white p-4 rounded-lg shadow"> <div class="bg-white p-4 rounded-lg shadow">
@@ -314,7 +314,9 @@
formData.append('voice_url', voiceUrl); formData.append('voice_url', voiceUrl);
} }
const response = await fetch('/tts', { const ttsEndpoint = `${window.location.pathname.replace(/\/+$/, '')}/tts`;
const response = await fetch(ttsEndpoint, {
method: 'POST', method: 'POST',
body: formData body: formData
}); });
+78 -14
View File
@@ -1,3 +1,4 @@
services: services:
traefik: traefik:
image: traefik:v3.6.7 image: traefik:v3.6.7
@@ -32,11 +33,10 @@ services:
- node.role == manager - node.role == manager
backend-english: 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: volumes:
- hf-cache:/root/.cache/huggingface - hf-cache:/root/.cache/huggingface
command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "english", "--quantize"]
deploy: deploy:
labels: labels:
- "traefik.enable=true" - "traefik.enable=true"
@@ -51,11 +51,10 @@ services:
- "traefik.http.routers.server-english.priority=100" - "traefik.http.routers.server-english.priority=100"
backend-german: 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: volumes:
- hf-cache:/root/.cache/huggingface - hf-cache:/root/.cache/huggingface
command: ["serve", "--port", "80", "--host", "0.0.0.0", "--language", "german", "--quantize"]
deploy: deploy:
labels: labels:
- "traefik.enable=true" - "traefik.enable=true"
@@ -69,19 +68,84 @@ services:
- "traefik.http.services.server-german.loadbalancer.server.port=80" - "traefik.http.services.server-german.loadbalancer.server.port=80"
- "traefik.http.routers.server-german.priority=100" - "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: volumes:
letsencrypt: letsencrypt:
portainer_data: hf-cache:
prometheus-data:
grafana-data:
networks: networks:
agent_network:
driver: overlay
attachable: true
driver_opts:
encrypted: "true"
default: default:
driver: overlay driver: overlay
attachable: true attachable: true