diff --git a/stacks/chatterbox-fast/README.md b/stacks/chatterbox-fast/README.md index 8d7b118..7c13fb4 100644 --- a/stacks/chatterbox-fast/README.md +++ b/stacks/chatterbox-fast/README.md @@ -60,6 +60,10 @@ Phase 3 will add `compose.yaml`, `Dockerfile`, `.env.example`. `GET /health` → `{status, sr, device, default_voice, voices_dir}`. +`GET /voices` → `{voices: [stem…], default}` — predefined `*.wav` stems in +`CBF_VOICES_DIR` (`_`-prefixed scratch/A-B files excluded). Clone refs are passed +per-request as an absolute path and aren't listed. + ## Config (env) | var | default | meaning | @@ -68,6 +72,19 @@ Phase 3 will add `compose.yaml`, `Dockerfile`, `.env.example`. | `CBF_VOICES_DIR` | `/refs` | dir of predefined voice wavs | | `CBF_DEFAULT_VOICE` | first wav in dir | default reference wav (path or name) | | `CBF_BIND` / `CBF_PORT` | `0.0.0.0` / `8197` | uvicorn bind | +| `CBF_TF32` | `1` | TF32 matmul/cudnn (free; off with `0`) | +| `CBF_SDPA_FLASH` | `1` | flash + mem-efficient SDPA backend | + +### Perf notes (measured 2026-06-02, turbo on A6000) + +- Model loads in **float32** (not the fp16 older notes assumed). +- **TF32 + SDPA do not move TTFA** (~0.5s): the first-sentence latency is bound by + the sequential AR token decode (T3 Llama, batch-1), not matmul throughput. They + stay on (free, help the larger chunks marginally). +- **bf16 deferred:** the lever that *would* help batch-1 decode, but `from_pretrained()` + has no dtype arg and turbo's fp32 conditioning path + dtype-sensitive vocoder make + a clean cast nontrivial. Not worth the quality risk while ~0.5s TTFA is fine. +- **torch.compile: deferred** (research flags a batch-1 regression). ## Dev / test on irv-ml1 diff --git a/stacks/chatterbox-fast/app.py b/stacks/chatterbox-fast/app.py index 735b3f1..91f79d4 100644 --- a/stacks/chatterbox-fast/app.py +++ b/stacks/chatterbox-fast/app.py @@ -45,11 +45,44 @@ DEFAULT_VOICE_ENV = os.environ.get("CBF_DEFAULT_VOICE") BIND = os.environ.get("CBF_BIND", "0.0.0.0") PORT = int(os.environ.get("CBF_PORT", "8197")) +# Perf levers (plan §4 Phase 2). TF32 + flash/mem-efficient SDPA are low-risk on +# Ampere and free — default ON. Measured 2026-06-02: they do NOT move TTFA, which +# is bound by the sequential AR token decode (T3 Llama at batch-1), not matmul +# throughput. bf16 (the lever that WOULD help batch-1 decode) is DEFERRED: turbo +# loads fp32 and from_pretrained() exposes no dtype arg, so bf16 needs whole-model +# casting incl. the speaker-conditioning path and the dtype-sensitive vocoder — +# real surgery + quality risk for a TTFA gain not currently needed (~0.5s is fine). +PERF_TF32 = os.environ.get("CBF_TF32", "1") == "1" +PERF_SDPA_FLASH = os.environ.get("CBF_SDPA_FLASH", "1") == "1" + # Turbo sampling knobs validated in the spike (plan §2). CFG / exaggeration / # min_p are ignored by turbo (it warns, harmless). WARMUP_TEXT = "Warming up the streaming engine." +def _predefined_wavs() -> list[Path]: + """Predefined voice wavs in VOICES_DIR, excluding `_`-prefixed scratch files + (bench/A-B outputs land as `_*.wav` in the same dir).""" + if not VOICES_DIR.is_dir(): + return [] + return sorted(p for p in VOICES_DIR.glob("*.wav") if not p.name.startswith("_")) + + +def _setup_perf() -> None: + """Apply the safe, low-risk speed levers before model load.""" + if PERF_TF32: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + log.info("perf: TF32 matmul/cudnn enabled") + if PERF_SDPA_FLASH and DEVICE.startswith("cuda"): + try: + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(True) + log.info("perf: flash + mem-efficient SDPA enabled") + except Exception as e: # pragma: no cover - backend-dependent + log.warning("perf: SDPA toggle failed: %s", e) + + # ── model holder ────────────────────────────────────────────────────────── @@ -71,10 +104,12 @@ class Engine: def load(self) -> None: from chatterbox.tts_turbo import ChatterboxTurboTTS + _setup_perf() log.info("loading ChatterboxTurboTTS on %s …", DEVICE) t0 = time.perf_counter() self.model = ChatterboxTurboTTS.from_pretrained(device=DEVICE) self.sr = int(getattr(self.model, "sr", 24000)) + self._log_model_dtype() self.default_voice = self._discover_default_voice() if self.default_voice: self._prepare(self.default_voice) @@ -84,16 +119,23 @@ class Engine: time.perf_counter() - t0, self.sr, self.default_voice, ) + def _log_model_dtype(self) -> None: + for name in ("t3", "s3gen", "model"): + sub = getattr(self.model, name, None) + try: + if sub is not None: + dt = next(sub.parameters()).dtype + log.info("dtype[%s]=%s", name, dt) + except (StopIteration, AttributeError): + pass + def _discover_default_voice(self) -> str | None: if DEFAULT_VOICE_ENV: # Resolve a bare name ("glados_25s") the same way request-time voices # are resolved — append .wav and look under VOICES_DIR. return self.resolve_voice(DEFAULT_VOICE_ENV) - if VOICES_DIR.is_dir(): - wavs = sorted(VOICES_DIR.glob("*.wav")) - if wavs: - return str(wavs[0]) - return None + wavs = _predefined_wavs() + return str(wavs[0]) if wavs else None def resolve_voice(self, voice: str | None) -> str: if not voice: @@ -209,6 +251,15 @@ def health() -> dict: } +@app.get("/voices") +def voices() -> dict: + """Predefined voices = the *.wav stems in CBF_VOICES_DIR. Clone refs are + passed per-request as an absolute path, so they're not listed here.""" + names = [p.stem for p in _predefined_wavs()] + default = Path(engine.default_voice).stem if engine.default_voice else None + return {"voices": names, "default": default} + + @app.post("/tts") def tts(req: TTSRequest) -> StreamingResponse: if engine.model is None: