diff --git a/docker-compose.yml b/docker-compose.yml index 3d0e0b5..f6609d2 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,7 +15,7 @@ services: volumes: # Kokoro ONNX models (read-only) - ./models:/app/models:ro - # HuggingFace cache for SNAC model download + # HuggingFace cache for SNAC model download (lazy-loaded on first Orpheus call) - hf-cache:/home/tts/.cache/huggingface # PipeWire socket for audio playback through host speakers - /run/user/1000/pipewire-0:/run/user/1000/pipewire-0 diff --git a/src/tts_mcp/engines/orpheus.py b/src/tts_mcp/engines/orpheus.py index fb86577..9ec22fe 100644 --- a/src/tts_mcp/engines/orpheus.py +++ b/src/tts_mcp/engines/orpheus.py @@ -1,18 +1,18 @@ """Orpheus TTS via Ollama + SNAC decoder. Sends text to Ollama's Orpheus model, parses responses, -and decodes through SNAC to 24kHz WAV. SNAC runs on CPU (RTX 5070 SM 120 -is not yet supported by PyTorch 2.6). +and decodes through SNAC to 24kHz WAV. SNAC is lazy-loaded on first use +and runs on CPU (RTX 5070 SM 120 not yet supported by PyTorch 2.6). """ import asyncio +import os import re import sys import time import numpy as np import requests -import torch from ..audio import wav_duration, write_wav from ..settings import settings @@ -46,6 +46,8 @@ def _tokens_to_audio(token_strings: list[str], snac_model) -> np.ndarray | None: - codes_1: 2 per frame (positions 1, 4) -> mid - codes_2: 4 per frame (positions 2,3,5,6) -> fine """ + import torch + num_frames = len(token_strings) // 7 if num_frames == 0: return None @@ -71,7 +73,6 @@ def _tokens_to_audio(token_strings: list[str], snac_model) -> np.ndarray | None: def clamp(lst): return [max(0, min(SNAC_CODEBOOK_SIZE - 1, x)) for x in lst] - # SNAC runs on CPU — RTX 5070 SM 120 not supported by PyTorch 2.6 device = "cpu" codes = [ torch.tensor(clamp(codes_0), dtype=torch.long).unsqueeze(0).to(device), @@ -85,17 +86,44 @@ def _tokens_to_audio(token_strings: list[str], snac_model) -> np.ndarray | None: return audio.squeeze().cpu().numpy() +def _load_snac(): + """Load SNAC model on CPU. Called once, lazily.""" + os.environ["CUDA_VISIBLE_DEVICES"] = "" + from snac import SNAC + + return SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to("cpu") + + class OrpheusEngine(TTSEngine): - """Orpheus TTS via Ollama's completions API + SNAC audio decoding.""" + """Orpheus TTS via Ollama's completions API + SNAC audio decoding. + + SNAC is lazy-loaded on first synthesize() call to avoid holding + ~200MB of RAM when the engine isn't being used. + """ name = "orpheus" default_voice = "tara" - def __init__(self, snac_model, ollama_url: str, model_name: str) -> None: - self._snac = snac_model + def __init__(self, ollama_url: str, model_name: str) -> None: + self._snac = None + self._snac_lock = asyncio.Lock() self._ollama_url = ollama_url self._model = model_name + async def _get_snac(self): + """Lazy-load SNAC on first use.""" + if self._snac is not None: + return self._snac + async with self._snac_lock: + if self._snac is not None: + return self._snac + print("Loading SNAC decoder on CPU (first Orpheus call)...", file=sys.stderr) + t0 = time.time() + loop = asyncio.get_running_loop() + self._snac = await loop.run_in_executor(None, _load_snac) + print(f" SNAC ready in {time.time() - t0:.1f}s", file=sys.stderr) + return self._snac + async def synthesize(self, text: str, voice: str | None = None) -> TTSResult: voice = voice or self.default_voice prompt = f"<|audio|>{voice}: {text}<|eot_id|>" @@ -153,9 +181,10 @@ class OrpheusEngine(TTSEngine): f"Response preview: {resp_text[:200]}" ) - # SNAC decode (CPU-bound) + # Lazy-load SNAC, then decode (CPU-bound) + snac = await self._get_snac() audio = await loop.run_in_executor( - None, _tokens_to_audio, token_strings, self._snac + None, _tokens_to_audio, token_strings, snac ) if audio is None: raise RuntimeError("SNAC decoding produced no audio") diff --git a/src/tts_mcp/server.py b/src/tts_mcp/server.py index cd7f6c7..dc93ccf 100644 --- a/src/tts_mcp/server.py +++ b/src/tts_mcp/server.py @@ -1,7 +1,6 @@ """FastMCP 3.0 server — tools, lifespan, and resource definitions.""" import asyncio -import os import sys import time from contextlib import asynccontextmanager @@ -41,24 +40,11 @@ async def app_lifespan(server: FastMCP): ) print(f" Kokoro ready in {time.time() - t0:.1f}s", file=sys.stderr) - # --- SNAC decoder (CPU only — RTX 5070 SM 120 unsupported by PyTorch 2.6) --- - print("Loading SNAC decoder on CPU...", file=sys.stderr) - t0 = time.time() - - def _load_snac(): - os.environ["CUDA_VISIBLE_DEVICES"] = "" - from snac import SNAC - - return SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to("cpu") - - snac_model = await loop.run_in_executor(None, _load_snac) - print(f" SNAC ready in {time.time() - t0:.1f}s", file=sys.stderr) - # --- Build engines --- engines: dict[str, TTSEngine] = { "piper": PiperEngine(settings.piper_host, settings.piper_port), "kokoro": KokoroEngine(kokoro_model), - "orpheus": OrpheusEngine(snac_model, settings.ollama_url, settings.orpheus_model), + "orpheus": OrpheusEngine(settings.ollama_url, settings.orpheus_model), } # Health check all engines at startup