Files
scott 29b66e24bb
All checks were successful
Build ROCm Image / build (push) Successful in 3m17s
Cache voice conditionals and add FP16 autocast
Voice conditionals (s3tokenizer + voice encoder + mel embeddings) are
expensive to compute but depend only on the reference audio, not the
text. Previously they ran on every synthesis chunk — 3x wasted work for
a 3-chunk request. Now computed once at startup and reused.

Also wrap generate() in torch.amp.autocast(float16) for ~2x speedup on
all model computation (T3 LLM, S3Gen CFM, HiFiGAN vocoder).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-05 20:22:13 -04:00

67 lines
2.0 KiB
Python

import asyncio
import logging
import sys
from functools import partial
from wyoming.server import AsyncServer
import engine
from config import get_wyoming_host, get_wyoming_port, load_config
from wyoming_handler import ChatterboxWyomingHandler
from wyoming_voices import create_wyoming_info, load_voices
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
stream=sys.stdout,
)
logger = logging.getLogger(__name__)
def _warmup(voices: dict) -> None:
"""Pre-compute voice conditionals and populate MIOpen's kernel cache."""
from wyoming_voices import resolve_voice
# Pre-compute conditionals for all discovered voices so the first real
# request doesn't pay the s3tokenizer + voice encoder cost.
for name, path in voices.items():
try:
engine.prepare_voice(path)
except Exception:
logger.warning(f"Failed to prepare voice '{name}' (non-fatal)", exc_info=True)
# Synthesis warmup to populate MIOpen's in-memory kernel cache.
audio_prompt = resolve_voice(None, voices) if voices else None
logger.info("Running warmup synthesis...")
try:
engine.synthesize(text="Warmup.", audio_prompt_path=audio_prompt)
logger.info("Warmup complete")
except Exception:
logger.warning("Warmup synthesis failed (non-fatal)", exc_info=True)
async def main() -> None:
load_config()
logger.info("Loading TTS model...")
if not engine.load_model():
logger.error("Failed to load model, exiting")
sys.exit(1)
voices = load_voices()
wyoming_info = create_wyoming_info(engine.get_sample_rate(), voices)
_warmup(voices)
host = get_wyoming_host()
port = get_wyoming_port()
uri = f"tcp://{host}:{port}"
logger.info(f"Starting Wyoming server on {uri}")
server = AsyncServer.from_uri(uri)
await server.run(partial(ChatterboxWyomingHandler, wyoming_info, voices))
if __name__ == "__main__":
asyncio.run(main())