first

.dbg/trae-debug-log-swagger-audio-play.ndjson

@@ -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}

fastmcp_server.py

@@ -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()

pyproject.toml

@@ -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

server.py

@@ -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)

start-pocke-tts.sh

@@ -0,0 +1,2 @@
+uv run --with "fastapi,uvicorn,scipy" python3 server.py
+