Files
rocm-chatterbox-whisper/main.py
scott 59731084cd
All checks were successful
Build ROCm Image / build (push) Successful in 2m49s
Multi-pass warmup and smaller chunk_size to fix HA timeout
torch.compile with dynamic=True still specializes per shape family on
first call. The warmup was running one text length, leaving real requests
to JIT-compile their own shapes (15-22s for first chunk). HA freezes
because it gets no AudioChunk for 22 seconds.

Fix:
- Run 3 warmup passes (short/medium/long text) so torch.compile builds
  a dynamic shape graph covering the range HA actually sends. Real
  requests then hit a cached compilation and synthesize in 3-8s.
- Reduce default chunk_size from 300 to 120 chars so the first text
  chunk is shorter, producing faster synthesis and earlier first audio.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-05 15:04:46 -04:00

78 lines
2.4 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__)
_WARMUP_TEXTS = [
# Short: covers brief HA notifications (lights on/off, etc.)
"Okay.",
# Medium: covers typical HA announcements
"The front door is open. Please close it.",
# Long: covers longer TTS requests and pre-compiles dynamic shape graph
(
"This is a warmup synthesis to pre-compile neural network kernels "
"for longer text lengths used in Home Assistant announcements and notifications."
),
]
def _warmup(voices: dict) -> None:
from wyoming_voices import resolve_voice
audio_prompt = resolve_voice(None, voices) if voices else None
logger.info(
f"Running {len(_WARMUP_TEXTS)}-pass warmup to pre-compile torch kernels "
"for short, medium, and long text lengths..."
)
for i, text in enumerate(_WARMUP_TEXTS, 1):
try:
engine.synthesize(text=text, audio_prompt_path=audio_prompt)
logger.info(f"Warmup pass {i}/{len(_WARMUP_TEXTS)} complete")
except Exception:
logger.warning(f"Warmup pass {i} failed (non-fatal)", exc_info=True)
logger.info("Warmup complete")
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)
# Run a warmup synthesis before accepting connections so MIOpen benchmarks
# and caches the best convolution algorithms for all layer shapes. Without
# this, the first real HA request triggers benchmarking (hundreds of runs)
# and times out before any audio is returned.
_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())