Files
rocm-chatterbox-whisper/main.py
scott 8de67c8bd9
Some checks failed
Build ROCm Image / build (push) Failing after 11s
Switch to ROCm 6.1 + torch 2.5.1 to fix MIOpen workspace=0 slowness
ROCm 7.2 + PyTorch 2.11.0 has a bug where PyTorch passes workspace=0 to
MIOpen convolutions, forcing fallback to the slow GemmFwdRest solver.
This caused s3gen.inference to take 15-22s instead of <5s, making
synthesis 3-4x slower than real-time audio playback.

ROCm 6.1 allocates workspace correctly so MIOpen picks fast GEMM solvers
without needing torch.compile workarounds.

Changes:
- Base image: rocm/dev-ubuntu-22.04:7.2 → 6.1
- torch 2.11.0 → 2.5.1 (rocm6.1 wheel index)
- Add pytorch_triton_rocm==3.1.0
- transformers 5.2.0 → 4.46.3, safetensors 0.5.3 → 0.4.0
- s3tokenizer unpinned → 0.3.0
- resemble-perth==1.0.1 directly (v1.0.1 is pip-installable; drop stub)
- Drop Dockerfile perth_stub steps
- Drop torch.compile and timing patches from engine.py (not needed)
- Drop multi-pass warmup from main.py (torch JIT warmup not needed)
- Drop ROCm 7.2-specific env vars from docker-compose.yml

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-05 17:27:21 -04:00

57 lines
1.6 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:
"""Run one synthesis to populate MIOpen's in-memory kernel cache."""
from wyoming_voices import resolve_voice
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())