Compare commits

...
10 Commits
Author SHA1 Message Date
valenti 381157848b first
Run Tests / pre-commit (push) Has been cancelled
Run Tests / pytest (3.10) (push) Has been cancelled
Run Tests / pytest (3.14) (push) Has been cancelled
2026-09-07 18:37:18 +00:00
Manu Orsini d108410d23 Default the English model's temperature to 0.3 (#223)
## Motivation

Human evaluations consistently prefer the English model at temperature
**0.3** over the current default **0.7**, and the change is free on
every objective axis we measured.
2026-07-16 14:42:04 +02:00
joshwhiton 9514983efd added project LocalVocal.ai to README.md (#200) 2026-06-23 15:12:05 +02:00
Iheuzio 15a6c1817b Add 'seshat-tts' project to README (#197)
Added a new project entry for 'seshat-tts' in the README.
2026-06-03 10:55:08 +02:00
Gabriel de Marmiesse 1522e38b8d Remove outdated issue about quantization 2026-05-27 10:23:12 +02:00
John S. 1fadd1cac5 Add new macOS app links to README (#193) 2026-05-23 18:23:53 +02:00
Gabriel de Marmiesse 2dff8a2d1b Fix pocket-tts generate docs (#182)
They were missing some information related to the new languages and
voices
2026-05-05 10:20:59 +02:00
gabrieldemarmiesse 677126be9e Use only amd64 for the demo 2026-05-04 15:58:38 +02:00
gabrieldemarmiesse 058886528d Bump version to 2.1.0 2026-05-04 15:51:28 +02:00
Gabriel de Marmiesse a882e0680e Fix voice cloning with quantization (#181)
Fix https://github.com/kyutai-labs/pocket-tts/issues/178
2026-05-04 15:50:42 +02:00
16 changed files with 1700 additions and 107 deletions
@@ -0,0 +1,4 @@
{"sessionId": "swagger-audio-play", "runId": "pre-fix", "hypothesisId": "A", "location": "server.py:/tts", "msg": "[DEBUG] TTS request received", "data": {"format": "mp3", "voice_url": "lola", "text_length": 1500}, "ts": 1784807181342}
{"sessionId": "swagger-audio-play", "runId": "pre-fix", "hypothesisId": "B", "location": "server.py:/tts", "msg": "[DEBUG] Audio payload prepared", "data": {"media_type": "audio/mpeg", "sample_rate": 24000, "audio_bytes_len": 767832, "first_16_bytes_hex": "fff384c400000000000000000058696e"}, "ts": 1784807884454}
{"sessionId": "swagger-audio-play", "runId": "pre-fix", "hypothesisId": "C", "location": "server.py:/tts", "msg": "[DEBUG] Response branch selected", "data": {"format": "mp3", "media_type": "audio/mpeg"}, "ts": 1784807884455}
{"sessionId": "swagger-audio-play", "runId": "pre-fix", "hypothesisId": "A", "location": "server.py:/tts", "msg": "[DEBUG] TTS request received", "data": {"format": "mp3", "voice_url": "lola", "text_length": 1290}, "ts": 1784807884459}
+4 -2
View File
@@ -162,7 +162,6 @@ You can check out the [Python API documentation](https://kyutai-labs.github.io/p
At the moment, we do not support (but would love pull requests adding):
- [Adding silence in the text input to generate pauses.](https://github.com/kyutai-labs/pocket-tts/issues/6)
- [Quantization to run the computation in int8.](https://github.com/kyutai-labs/pocket-tts/issues/7)
We tried running this TTS model on the GPU but did not observe a speedup compared to CPU execution,
notably because we use a batch size of 1 and a very small model.
@@ -196,7 +195,8 @@ We don't have official support for this yet, but you can try out one of these co
- [pocket-reader](https://github.com/lukasmwerner/pocket-reader) by @lukasmwerner- Browser screen reader
- [pocket-tts-wyoming](https://github.com/ikidd/pocket-tts-wyoming) by @ikidd - Docker container for pocket-tts using Wyoming protocol, ready for Home Assistant Voice use.
- [Sonorus](https://www.nexusmods.com/hogwartslegacy/mods/2409) by @KevinAHM - Talk to any named character in Hogwarts Legacy with their original voice.
- [Mac pocket-tts](https://github.com/slaughters85j/pocket-tts) by @slaughters85j - Mac Desktop App + macOS Quick Action
- [Native macOS App](https://github.com/slaughters85j/pocket-tts-macos) by @slaughters85j - Native macOS app, Python-free. Runs Pocket-TTS via Core ML, fully on-device. Includes signed and notarized .app releases.
- [Electron macOS App](https://github.com/slaughters85j/pocket-tts) by @slaughters85j - Electron Mac Desktop App + macOS Quick Action
- [pocket-tts-openai_streaming_server](https://github.com/teddybear082/pocket-tts-openai_streaming_server) by @teddybear082 - OpenAI-compatible streaming server, dockerized and with an `.exe` release
- [pocket-tts-unity](https://github.com/lookbe/pocket-tts-unity) by @lookbe - A Unity 6 integration for Pocket-TTS.
- [ComfyUI-Pocket-TTS](https://github.com/ai-joe-git/ComfyUI-Pocket-TTS) by @ai-joe-git Lightweight CPU-based Text-to-Speech for ComfyUI
@@ -208,6 +208,8 @@ We don't have official support for this yet, but you can try out one of these co
- [openclaw-pockettts](https://github.com/dodgyrabbit/openclaw-pockettts) by @dodgyrabbit - A Docker container with the Python implementation but exposed as an OpenAI TTS API for easy integration with OpenClaw.
- [openclaw-pocketts.cpp](https://github.com/dodgyrabbit/openclaw-pockettts.cpp) by @dodgyrabbit - A Docker container with the PocketTTS.cpp version, packaged for easy integration with OpenClaw.
- [tts-audiobook-tool](https://github.com/zeropointnine/tts-audiobook-tool) by @zeropointnine - Multi-model audiobook generator with automatic error detection, 48khz upscaling, synced browser reader, stand-alone server-mode.
- [seshat-tts](https://github.com/scriptriva/seshat-tts) by @scriptriva - Accessibility tool that provides real-time audio synthesis for games and apps. It also features a voice manager capable of cloning voices based on user presets.
- [LocalVocal.ai](https://localvocal.ai) by @joshwhiton - Fully local conversational voice-harness for Macs with Apple Silicon. Includes voice-activity & turn detection, dictation, voice cloning, CLI to talk to Claude, Codex... and more.
## Prohibited use
+1 -1
View File
@@ -9,7 +9,7 @@ group "default" {
target "pocket-tts-server" {
context = "."
dockerfile = "Dockerfile"
platforms = ["linux/arm64", "linux/amd64"]
platforms = ["linux/amd64"]
tags = [
"rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server:${COMMIT_SHA}"
]
+2 -2
View File
@@ -37,14 +37,14 @@ The main class for text-to-speech generation.
#### Class Methods
##### `load_model(language=None, config=None, temp=0.7, lsd_decode_steps=1, noise_clamp=None, eos_threshold=-4.0, quantize=False)`
##### `load_model(language=None, config=None, temp=None, lsd_decode_steps=1, noise_clamp=None, eos_threshold=-4.0, quantize=False)`
Load and return a TTSModel instance with pre-trained weights.
**Parameters:**
- `language` (str | None): Name of built-in language config to load. Supported values: `"english_2026-01"`, `"english_2026-04"`, `"english"`, `"french_24l"`, `"german_24l"`, `"portuguese_24l"`, `"italian_24l"`, `"spanish_24l"`. If both `language` and `config` are omitted, defaults to `"english"`, which is the same model as `"english_2026-04"`. The "24l" variants are larger models that are not distilled yet and are here only as a preview.
- `config` (str | None): Path to model config YAML file. Incompatible with `language`.
- `temp` (float): Sampling temperature for generation (default: 0.7)
- `temp` (float | None): Sampling temperature for generation. `None` uses the model's recommended default from its config file (`default_temperature`; 0.3 for the English model, 0.7 otherwise).
- `lsd_decode_steps` (int): Number of generation steps (default: 1)
- `noise_clamp` (float | None): Maximum value for noise sampling (default: None)
- `eos_threshold` (float): Threshold for end-of-sequence detection (default: -4.0)
+3 -3
View File
@@ -17,15 +17,15 @@ This will generate a WAV file `./tts_output.wav` with the default text and voice
### Core Options
- `--text TEXT`: Text to generate (default: "Hello world! I am Kyutai Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.")
- `--voice VOICE`: Path to audio conditioning file (voice to clone). Defaults to a built-in voice chosen from the language: `giovanni` (it) for italian, `lola` (es) for spanish, `alba` (en) otherwise. Urls and local paths are supported.
- `--voice VOICE`: Path to audio conditioning file (voice to clone). Defaults to a built-in voice chosen from the language: `giovanni` (it) for italian, `lola` (es) for spanish, `juergen` (de) for german, `rafael` (pt) for portuguese, `estelle` (fr) for french, `alba` (en) otherwise. Urls and local paths are supported.
- `--output-path OUTPUT_PATH`: Output path for generated audio (default: "./tts_output.wav")
- `--language LANGUAGE`: Language for the TTS model, one of `'english_2026-01'`, `'english_2026-04'`, `'english'`, `'french_24l'`, `'german_24l'`, `'portuguese_24l'`, `'italian_24l'`, `'spanish_24l'` (default: `english`, which is the same model as `'english_2026-04'`). Incompatible with `--config`. The "24l" variants are bigger models, not distilled yet and here only as preview.
- `--language LANGUAGE`: Language for the TTS model, one of `'english_2026-01'`, `'english_2026-04'`, `'english'`, `'french_24l'`, `'german'`, `'german_24l'`, `'portuguese'`, `'portuguese_24l'`, `'italian'`, `'italian_24l'`, `'spanish'`, `'spanish_24l'` (default: `english`, which is the same model as `'english_2026-04'`). Incompatible with `--config`. The "24l" variants are bigger models, not distilled yet and here only as preview.
### Generation Parameters
- `--config CONFIG_PATH`: Path to custom config.yaml (for loading local model files). Incompatible with `--language`.
- `--lsd-decode-steps LSD_DECODE_STEPS`: Number of generation steps (default: 1)
- `--temperature TEMPERATURE`: Temperature for generation (default: 0.7)
- `--temperature TEMPERATURE`: Temperature for generation (default: the model's recommended value from its config — 0.3 for the English model, 0.7 otherwise)
- `--noise-clamp NOISE_CLAMP`: Noise clamp value (default: None)
- `--eos-threshold EOS_THRESHOLD`: EOS threshold (default: -4.0)
- `--frames-after-eos FRAMES_AFTER_EOS`: Number of frames to generate after EOS (default: None, auto-calculated based on the text length). Each frame is 80ms.
+127
View File
@@ -0,0 +1,127 @@
import base64
import io
import json
import os
import numpy as np
import scipy.io.wavfile as wavfile
from fastmcp import FastMCP
from pocket_tts import TTSModel
from pocket_tts.default_parameters import get_default_voice_for_language
from pocket_tts.utils.utils import _ORIGINS_OF_PREDEFINED_VOICES
DEFAULT_LANGUAGE = os.environ.get("POCKET_TTS_MCP_LANGUAGE", "spanish_24l")
DEFAULT_TEMP = float(os.environ.get("POCKET_TTS_MCP_TEMP", "0.5"))
DEFAULT_LSD_DECODE_STEPS = int(os.environ.get("POCKET_TTS_MCP_LSD_STEPS", "5"))
DEFAULT_EOS_THRESHOLD = float(os.environ.get("POCKET_TTS_MCP_EOS_THRESHOLD", "-3.0"))
DEFAULT_HOST = os.environ.get("POCKET_TTS_MCP_HOST", "0.0.0.0")
DEFAULT_PORT = int(os.environ.get("POCKET_TTS_MCP_PORT", "8008"))
_model: TTSModel | None = None
mcp = FastMCP(
name="Pocket TTS MCP",
instructions=(
"Generate WAV speech audio with Pocket TTS and inspect the built-in voice presets."
),
)
def get_model() -> TTSModel:
global _model
if _model is None:
print("Loading Pocket TTS model for FastMCP server...")
_model = TTSModel.load_model(
language=DEFAULT_LANGUAGE,
temp=DEFAULT_TEMP,
lsd_decode_steps=DEFAULT_LSD_DECODE_STEPS,
eos_threshold=DEFAULT_EOS_THRESHOLD,
)
return _model
def list_predefined_voices() -> list[dict[str, str | bool]]:
default_voice = get_default_voice_for_language(DEFAULT_LANGUAGE)
return [
{
"name": name,
"source": source,
"is_default": name == default_voice,
}
for name, source in sorted(_ORIGINS_OF_PREDEFINED_VOICES.items())
]
def encode_wav_base64(audio_tensor, sample_rate: int) -> tuple[str, int]:
if hasattr(audio_tensor, "detach"):
audio_numpy = audio_tensor.detach().cpu().numpy()
else:
audio_numpy = np.array(audio_tensor)
if audio_numpy.ndim > 1:
audio_numpy = audio_numpy.squeeze()
audio_numpy = np.asarray(audio_numpy, dtype=np.float32)
buffer = io.BytesIO()
wavfile.write(buffer, sample_rate, audio_numpy)
wav_bytes = buffer.getvalue()
return base64.b64encode(wav_bytes).decode("ascii"), len(wav_bytes)
@mcp.tool()
def list_voices() -> dict:
"""List the built-in Pocket TTS voice presets available to the MCP client."""
return {
"language": DEFAULT_LANGUAGE,
"default_voice": get_default_voice_for_language(DEFAULT_LANGUAGE),
"voices": list_predefined_voices(),
}
@mcp.resource("voice://predefined")
def predefined_voices_resource() -> str:
"""Return the predefined voice catalog as JSON."""
return json.dumps(list_voices(), indent=2)
@mcp.tool()
def generate_speech(
text: str,
voice: str | None = None,
frames_after_eos: int | None = None,
copy_state: bool = True,
) -> dict:
"""Generate speech audio and return a base64-encoded WAV payload."""
if not text.strip():
raise ValueError("text must not be empty")
model = get_model()
selected_voice = voice or get_default_voice_for_language(DEFAULT_LANGUAGE)
voice_state = model.get_state_for_audio_prompt(selected_voice)
audio_tensor = model.generate_audio(
model_state=voice_state,
text_to_generate=text,
frames_after_eos=frames_after_eos,
copy_state=copy_state,
)
audio_base64, wav_size_bytes = encode_wav_base64(audio_tensor, model.sample_rate)
return {
"text": text,
"voice": selected_voice,
"language": DEFAULT_LANGUAGE,
"sample_rate": model.sample_rate,
"mime_type": "audio/wav",
"wav_size_bytes": wav_size_bytes,
"audio_base64": audio_base64,
}
def main() -> None:
mcp.run(transport="sse", host=DEFAULT_HOST, port=DEFAULT_PORT)
if __name__ == "__main__":
main()
+2
View File
@@ -2,6 +2,8 @@
weights_path: hf://kyutai/pocket-tts/languages/english/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9
weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/english/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc
# Human evals preferred this model at 0.3 (equal WER/similarity/speech rate).
default_temperature: 0.3
flow_lm:
+2
View File
@@ -2,6 +2,8 @@
weights_path: hf://kyutai/pocket-tts/languages/english_2026-04/model.safetensors@19f95fe2df36e79fbd9f10008595cc4c977a0fcc
weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/english_2026-04/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc
# Human evals preferred this model at 0.3 (equal WER/similarity/speech rate).
default_temperature: 0.3
flow_lm:
+6 -3
View File
@@ -20,7 +20,6 @@ from pocket_tts.default_parameters import (
DEFAULT_FRAMES_AFTER_EOS,
DEFAULT_LSD_DECODE_STEPS,
DEFAULT_NOISE_CLAMP,
DEFAULT_TEMPERATURE,
MAX_TOKEN_PER_CHUNK,
get_default_text_for_language,
get_default_voice_for_language,
@@ -261,8 +260,12 @@ def generate(
int, typer.Option(help="Number of generation steps")
] = DEFAULT_LSD_DECODE_STEPS,
temperature: Annotated[
float, typer.Option(help="Temperature for generation")
] = DEFAULT_TEMPERATURE,
float | None,
typer.Option(
help="Temperature for generation. Defaults to the model's recommended "
"value from its config (0.3 for the English model, 0.7 otherwise)."
),
] = None,
noise_clamp: Annotated[float, typer.Option(help="Noise clamp value")] = DEFAULT_NOISE_CLAMP,
eos_threshold: Annotated[float, typer.Option(help="EOS threshold")] = DEFAULT_EOS_THRESHOLD,
frames_after_eos: Annotated[
+6 -3
View File
@@ -25,7 +25,6 @@ from pocket_tts.default_parameters import (
DEFAULT_LANGUAGE,
DEFAULT_LSD_DECODE_STEPS,
DEFAULT_NOISE_CLAMP,
DEFAULT_TEMPERATURE,
MAX_TOKEN_PER_CHUNK,
)
from pocket_tts.models.flow_lm import FlowLMModel
@@ -234,7 +233,7 @@ class TTSModel(nn.Module):
cls,
language: str | None = None,
config: str | Path | None = None,
temp: float | int = DEFAULT_TEMPERATURE,
temp: float | int | None = None,
lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS,
noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP,
eos_threshold: float = DEFAULT_EOS_THRESHOLD,
@@ -253,7 +252,9 @@ class TTSModel(nn.Module):
If neither `config` nor `language` is provided, defaults to `"english", which is the same model as 'english_2026-04'`.
config: A path to a custom YAML config file saved locally (e.g., `"C://pocket_tts/pocket_tts_config.yaml"`).
temp: Sampling temperature for generation. Higher values produce more
diverse but potentially lower quality output.
diverse but potentially lower quality output. If None, defaults to
the model's recommended value from its config file
(``default_temperature``, e.g. 0.3 for the English model).
lsd_decode_steps: Number of steps for Lagrangian Self Distillation
decoding. More steps can improve quality but increase computation.
noise_clamp: Maximum value for noise sampling. If None, no clamping
@@ -303,6 +304,8 @@ class TTSModel(nn.Module):
raise ValueError("Config should be a path to a YAML file ending with .yaml")
config_path = Path(config)
config = load_config(config_path)
if temp is None:
temp = config.default_temperature
logger.info(f"Loading model from config at {config_path}...")
tts_model = TTSModel._from_pydantic_config_with_weights(
+9 -2
View File
@@ -118,8 +118,15 @@ class StreamingMultiheadAttention(StatefulModule):
self.out_proj = nn.Linear(embed_dim, mult * embed_dim, bias=False)
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
device = self.in_proj.weight.device
dtype = self.in_proj.weight.dtype
weight = self.in_proj.weight
if callable(weight):
# torch.ao dynamic-quantized Linear: weight is a method returning a
# qint8 tensor, but activations stay float32 — that's what the cache holds.
device = weight().device
dtype = torch.float32
else:
device = weight.device
dtype = weight.dtype
return self._cache_backend.init_state(batch_size, sequence_length, device, dtype)
def increment_step(self, state: dict, increment: int = 1):
+1
View File
@@ -116,6 +116,7 @@ class Config(StrictModel):
pad_with_spaces_for_short_inputs: bool = False
remove_semicolons: bool = False
model_recommended_frames_after_eos: int | None = None
default_temperature: float = 0.7
def load_config(yaml_path: str | Path) -> Config:
+4 -1
View File
@@ -2,7 +2,7 @@
name = "pocket-tts"
requires-python = ">= 3.10,<3.15"
description = "Kyutai's pocket-sized text-to-speech!"
version = "2.0.0"
version = "2.1.0"
readme = "README.md"
dependencies = [
"numpy>=2",
@@ -18,9 +18,11 @@ dependencies = [
"uvicorn>=0.13.0",
"python-multipart>=0.0.21",
"scipy>=1.5.0",
"soundfile>=0.12.0",
"einops>=0.4.0",
"huggingface_hub>=0.13.0",
"requests>=2.20.0",
"fastmcp>=2.0.0",
]
[project.optional-dependencies]
@@ -70,6 +72,7 @@ packages = ["pocket_tts"]
[project.scripts]
pocket-tts = "pocket_tts.main:cli_app"
pocket-tts-mcp = "fastmcp_server:main"
[tool.uv]
# System python distribution can have header files missing, so we use only the python
+288
View File
@@ -0,0 +1,288 @@
import io
import re
import traceback
from typing import Literal
import numpy as np
import scipy.io.wavfile as wavfile
from fastapi import FastAPI, Form, HTTPException
from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.responses import RedirectResponse, Response
from pocket_tts import TTSModel
from pocket_tts.models.tts_model import prepare_text_prompt, split_into_best_sentences
app = FastAPI(
title="Pocket TTS High-Quality API",
description="HTTP API for generating MP3 or WAV speech audio with Pocket TTS.",
version="1.0.0",
docs_url=None,
redoc_url="/redoc",
openapi_url="/openapi.json",
)
SWAGGER_UI_VERSION = "5.32.10"
SWAGGER_UI_JS_URL = (
f"https://cdn.jsdelivr.net/npm/swagger-ui-dist@{SWAGGER_UI_VERSION}/swagger-ui-bundle.js"
)
SWAGGER_UI_CSS_URL = (
f"https://cdn.jsdelivr.net/npm/swagger-ui-dist@{SWAGGER_UI_VERSION}/swagger-ui.css"
)
SAFE_MAX_TOKENS_PER_CHUNK = 35
JOIN_SILENCE_MS = 100
MIN_TOKENS_PER_CHUNK = 12
print("Caricamento del modello Pocket TTS in memoria...")
model = TTSModel.load_model(
language="spanish_24l", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0
)
last_generated_audio: dict[str, bytes | str] | None = None
@app.get("/", include_in_schema=False)
async def swagger_ui_redirect():
return RedirectResponse(url="/docs")
@app.get("/docs", include_in_schema=False)
async def custom_swagger_ui():
return get_swagger_ui_html(
openapi_url=app.openapi_url,
title=f"{app.title} - Swagger UI",
swagger_js_url=SWAGGER_UI_JS_URL,
swagger_css_url=SWAGGER_UI_CSS_URL,
swagger_ui_parameters={
"displayRequestDuration": True,
},
)
def build_audio_response(audio_bytes: bytes, media_type: str) -> Response:
return Response(
content=audio_bytes,
media_type=media_type,
headers={"Cache-Control": "no-store"},
)
def normalize_tts_text(text: str) -> str:
def decode_unicode_match(match: re.Match[str]) -> str:
return match.group(0).encode("ascii").decode("unicode_escape")
if "\\u" in text:
text = re.sub(r"\\u[0-9a-fA-F]{4}", decode_unicode_match, text)
text = text.replace("\\n", "\n").replace("\\t", " ")
# Keep heading content while removing markdown wrappers.
text = re.sub(r"\*\*\s*\[([^\]]+)\]\s*\*\*", r"\n\1.\n", text)
text = re.sub(r"\[([^\]]+)\]", r"\1", text)
text = text.replace("**", "").replace("__", "")
# Remove list/heading markdown markers that harm pronunciation.
text = re.sub(r"^[#>*-]+\s*", "", text, flags=re.MULTILINE)
# Normalize whitespace while preserving sentence boundaries.
text = re.sub(r"\r\n?", "\n", text)
text = re.sub(r"\n{2,}", "\n\n", text)
text = re.sub(r"[ \t]+", " ", text)
text = re.sub(r" *\n *", "\n", text)
return text.strip()
def tensor_to_numpy(audio_tensor) -> np.ndarray:
if hasattr(audio_tensor, "detach"):
audio_numpy = audio_tensor.detach().cpu().numpy()
else:
audio_numpy = np.array(audio_tensor)
if audio_numpy.ndim > 1:
audio_numpy = audio_numpy.squeeze()
return np.asarray(audio_numpy, dtype=np.float32)
def count_model_tokens(text: str) -> int:
tokenizer = model.flow_lm.conditioner.tokenizer
prepared_text, _ = prepare_text_prompt(
text,
model.pad_with_spaces_for_short_inputs,
model.remove_semicolons,
)
return len(tokenizer(prepared_text.strip()).tokens[0].tolist())
def split_text_into_safe_chunks(text: str, max_tokens: int) -> list[str]:
tokenizer = model.flow_lm.conditioner.tokenizer
chunks = split_into_best_sentences(
tokenizer,
text,
max_tokens,
model.pad_with_spaces_for_short_inputs,
remove_semicolons=model.remove_semicolons,
)
refined_chunks: list[str] = []
for chunk in chunks:
chunk = chunk.strip()
token_count = count_model_tokens(chunk)
if token_count <= max_tokens:
refined_chunks.append(chunk)
continue
words = chunk.split()
current_words: list[str] = []
for word in words:
candidate_words = current_words + [word]
candidate_text = " ".join(candidate_words)
candidate_tokens = count_model_tokens(candidate_text)
if current_words and candidate_tokens > max_tokens:
refined_chunks.append(" ".join(current_words))
current_words = [word]
else:
current_words = candidate_words
if current_words:
refined_chunks.append(" ".join(current_words))
merged_chunks: list[str] = []
for chunk in refined_chunks:
chunk = chunk.strip()
if not chunk:
continue
chunk_tokens = count_model_tokens(chunk)
if not merged_chunks:
merged_chunks.append(chunk)
continue
previous_chunk = merged_chunks[-1]
previous_tokens = count_model_tokens(previous_chunk)
combined = f"{previous_chunk} {chunk}".strip()
combined_tokens = count_model_tokens(combined)
# Avoid tiny fragments when the merge still fits comfortably.
if chunk_tokens < MIN_TOKENS_PER_CHUNK and combined_tokens <= max_tokens:
merged_chunks[-1] = combined
continue
# Also merge if the previous chunk is tiny and the combined text still fits.
if previous_tokens < MIN_TOKENS_PER_CHUNK and combined_tokens <= max_tokens:
merged_chunks[-1] = combined
continue
merged_chunks.append(chunk)
for chunk in merged_chunks:
token_count = count_model_tokens(chunk)
if token_count > max_tokens:
raise ValueError(
f"Chunk exceeds hard token limit ({token_count} > {max_tokens}): {chunk[:120]!r}"
)
return merged_chunks
def generate_audio_from_chunks(voice_state: dict, text: str, sample_rate: int) -> np.ndarray:
chunks = split_text_into_safe_chunks(text, SAFE_MAX_TOKENS_PER_CHUNK)
audio_segments: list[np.ndarray] = []
silence = np.zeros(int(sample_rate * JOIN_SILENCE_MS / 1000), dtype=np.float32)
for index, chunk in enumerate(chunks):
audio_tensor = model.generate_audio(
voice_state,
chunk,
max_tokens=SAFE_MAX_TOKENS_PER_CHUNK,
copy_state=True,
)
audio_segments.append(tensor_to_numpy(audio_tensor))
if index < len(chunks) - 1:
audio_segments.append(silence.copy())
if not audio_segments:
return np.zeros(0, dtype=np.float32)
return np.concatenate(audio_segments)
@app.get("/v1/audio/speech", include_in_schema=False)
@app.get("/tts", include_in_schema=False)
async def get_last_generated_audio():
if last_generated_audio is None:
raise HTTPException(status_code=404, detail="No generated audio available yet")
return build_audio_response(
audio_bytes=last_generated_audio["audio_bytes"],
media_type=last_generated_audio["media_type"],
)
@app.post("/v1/audio/speech")
@app.post(
"/tts",
tags=["tts"],
summary="Generate speech audio",
description=(
"Generate speech from text and return the result as an MP3 file by default, "
"or as WAV when `format=wav`. "
"The same handler is available on `/v1/audio/speech` and `/tts`."
),
response_class=Response,
responses={
200: {
"description": "Generated MP3 or WAV audio",
"content": {"audio/mpeg": {}, "audio/wav": {}},
},
500: {"description": "Audio generation failed"},
},
)
async def generate_tts(
text: str = Form(...),
voice_url: str = Form("lola"),
format: Literal["mp3", "wav"] = Form("mp3"),
):
global last_generated_audio
try:
normalized_text = normalize_tts_text(text)
# 1. Recupera lo stato del prompt vocale (dict)
voice_state = model.get_state_for_audio_prompt(voice_url)
sample_rate = model.sample_rate if hasattr(model, "sample_rate") else 24000
audio_numpy = generate_audio_from_chunks(voice_state, normalized_text, sample_rate)
buffer = io.BytesIO()
if format == "wav":
wavfile.write(buffer, sample_rate, audio_numpy)
media_type = "audio/wav"
else:
try:
import soundfile as sf
except ImportError as exc:
raise HTTPException(
status_code=500,
detail=(
"MP3 output requires the optional `soundfile` dependency. "
"Install it with `uv add soundfile`."
),
) from exc
sf.write(buffer, audio_numpy, sample_rate, format="MP3")
media_type = "audio/mpeg"
buffer.seek(0)
audio_bytes = buffer.read()
last_generated_audio = {
"audio_bytes": audio_bytes,
"media_type": media_type,
}
return build_audio_response(audio_bytes=audio_bytes, media_type=media_type)
except Exception as e:
print("\n--- ERRORE DURANTE LA GENERAZIONE ---")
traceback.print_exc()
print("-------------------------------------\n")
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8009)
+2
View File
@@ -0,0 +1,2 @@
uv run --with "fastapi,uvicorn,scipy" python3 server.py
Generated
+1239 -90
View File
File diff suppressed because it is too large Load Diff