128 lines
3.7 KiB
Python
128 lines
3.7 KiB
Python
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()
|