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:
Ryan Malloy 2026-02-21 13:04:35 -07:00
parent c53db4b251
commit 4698d8b0d2
3 changed files with 40 additions and 25 deletions

View File

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

View File

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

View File

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