first

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