Lazy-load SNAC decoder to reduce idle memory
SNAC + torch no longer load at startup — deferred to first Orpheus call via double-checked locking. Startup drops from ~13s to 0.5s, idle RAM reduced by ~200MB. OrpheusEngine constructor no longer takes snac_model; it self-loads on demand.
This commit is contained in:
parent
c53db4b251
commit
4698d8b0d2
@ -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
|
||||
|
||||
@ -1,18 +1,18 @@
|
||||
"""Orpheus TTS via Ollama + SNAC decoder.
|
||||
|
||||
Sends text to Ollama's Orpheus model, parses <custom_token_N> 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")
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user