esh-pfi-infrastructure/stacks/chatterbox-fast/app.py

"""chatterbox-fast — streaming TTS server (Phase 1 MVP).

A lean FastAPI server on the ChatterboxTurboTTS library that streams audio using
the adaptive buffer-ratchet scheduler in ``scheduler.py`` (the meat). Sub-second
time-to-first-audio while keeping turbo's full quality; workload is single-stream
interactive (see docs/design/chatterbox-fast-plan.md).

Endpoints:
  POST /tts     — StreamingResponse of audio chunks (raw PCM s16le default).
  GET  /health  — model/voice readiness.

Config via env (all optional; sane dev defaults):
  CBF_MODEL_DEVICE   cuda | cuda:0 | cpu                  (default: cuda)
  CBF_VOICES_DIR     dir of predefined voice wavs          (default: /refs)
  CBF_DEFAULT_VOICE  default reference wav path/name       (default: first wav in dir)
  CBF_BIND / CBF_PORT  uvicorn bind                        (default: 0.0.0.0:8197)
"""

from __future__ import annotations

import logging
import os
import struct
import threading
import time
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Iterator, Literal

import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from scheduler import ChunkConfig, ChunkResult, stream_chunks

log = logging.getLogger("chatterbox-fast")
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

# ── config ────────────────────────────────────────────────────────────────

DEVICE = os.environ.get("CBF_MODEL_DEVICE", "cuda")
VOICES_DIR = Path(os.environ.get("CBF_VOICES_DIR", "/refs"))
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 ──────────────────────────────────────────────────────────


class Engine:
    """Owns the single ChatterboxTurboTTS instance + a generation lock.

    The workload is single-stream interactive, but prepare_conditionals mutates
    model state, so all generation is serialized under one lock. Concurrent
    multi-voice serving is a Phase-2 concern.
    """

    def __init__(self) -> None:
        self.model = None
        self.sr = 24000
        self.lock = threading.Lock()
        self._current_voice: str | None = None
        self.default_voice: str | None = None

    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)
            self._warm()
        log.info(
            "model ready in %.1fs (sr=%d, default_voice=%s)",
            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)
        wavs = _predefined_wavs()
        return str(wavs[0]) if wavs else None

    def resolve_voice(self, voice: str | None) -> str:
        if not voice:
            if not self.default_voice:
                raise HTTPException(503, "no default voice; set CBF_DEFAULT_VOICE")
            return self.default_voice
        p = Path(voice)
        if p.is_absolute() and p.exists():
            return str(p)
        cand = VOICES_DIR / (voice if voice.endswith(".wav") else f"{voice}.wav")
        if cand.exists():
            return str(cand)
        raise HTTPException(404, f"unknown voice {voice!r}")

    def _prepare(self, voice_path: str, exaggeration: float = 0.5) -> None:
        if voice_path == self._current_voice:
            return
        log.info("prepare_conditionals(%s)", voice_path)
        self.model.prepare_conditionals(voice_path, exaggeration=exaggeration, norm_loudness=True)
        self._current_voice = voice_path

    def _warm(self) -> None:
        with torch.inference_mode():
            self.model.generate(WARMUP_TEXT, repetition_penalty=1.2, top_p=0.95,
                                 temperature=0.8, top_k=1000)
        if DEVICE.startswith("cuda"):
            torch.cuda.synchronize()

    def generate(self, text: str, knobs: "TTSRequest") -> tuple[torch.Tensor, float]:
        """Synthesize ``text`` → (wav tensor [1,T], audio_seconds). CUDA-synced
        so the caller's clock delta is honest gen time."""
        with torch.inference_mode():
            wav = self.model.generate(
                text,
                repetition_penalty=knobs.repetition_penalty,
                top_p=knobs.top_p,
                temperature=knobs.temperature,
                top_k=knobs.top_k,
            )
        if DEVICE.startswith("cuda"):
            torch.cuda.synchronize()
        audio_sec = wav.shape[-1] / self.sr
        return wav, audio_sec


engine = Engine()


@asynccontextmanager
async def lifespan(app: FastAPI):
    engine.load()
    yield


app = FastAPI(title="chatterbox-fast", lifespan=lifespan)


# ── request / audio encoding ──────────────────────────────────────────────


class TTSRequest(BaseModel):
    text: str
    voice: str | None = None
    format: Literal["pcm", "wav"] = "pcm"
    stream: bool = True  # False ⇒ whole-text one-shot (for A/B vs streaming)
    exaggeration: float = 0.5
    temperature: float = 0.8
    top_p: float = 0.95
    top_k: int = 1000
    repetition_penalty: float = 1.2
    # Scheduler overrides (None ⇒ ChunkConfig defaults).
    margin: float | None = Field(default=None)
    margin_first: float | None = Field(default=None)
    rtf_prior: float | None = Field(default=None)


def _pcm16(wav: torch.Tensor) -> bytes:
    a = wav.detach().to(torch.float32).clamp_(-1.0, 1.0).cpu().numpy().reshape(-1)
    return (a * 32767.0).astype("<i2").tobytes()


def _wav_header(sr: int) -> bytes:
    """Streaming WAV header with unknown length (0xFFFFFFFF sizes)."""
    return (
        b"RIFF" + struct.pack("<I", 0xFFFFFFFF) + b"WAVE"
        + b"fmt " + struct.pack("<IHHIIHH", 16, 1, 1, sr, sr * 2, 2, 16)
        + b"data" + struct.pack("<I", 0xFFFFFFFF)
    )


def _chunk_config(req: TTSRequest) -> ChunkConfig:
    cfg = ChunkConfig()
    if req.margin is not None:
        cfg.margin = req.margin
    if req.margin_first is not None:
        cfg.margin_first = req.margin_first
    if req.rtf_prior is not None:
        cfg.rtf_prior = req.rtf_prior
    return cfg


# ── endpoints ─────────────────────────────────────────────────────────────


@app.get("/health")
def health() -> dict:
    return {
        "status": "ok" if engine.model is not None else "loading",
        "sr": engine.sr,
        "device": DEVICE,
        "default_voice": engine.default_voice,
        "voices_dir": str(VOICES_DIR),
    }


@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:
        raise HTTPException(503, "model still loading")
    if not req.text.strip():
        raise HTTPException(400, "empty text")

    voice_path = engine.resolve_voice(req.voice)
    cfg = _chunk_config(req)
    media = "audio/wav" if req.format == "wav" else "application/octet-stream"

    def body() -> Iterator[bytes]:
        # One request holds the lock for its whole stream (single-stream
        # workload); concurrent callers queue rather than corrupt conditionals.
        with engine.lock:
            engine._prepare(voice_path, exaggeration=req.exaggeration)
            if req.format == "wav":
                yield _wav_header(engine.sr)

            t_req = time.perf_counter()
            first_audio_ms: float | None = None

            if not req.stream:
                wav, audio_sec = engine.generate(req.text, req)
                first_audio_ms = (time.perf_counter() - t_req) * 1000
                log.info("oneshot: %.0fms gen, %.2fs audio", first_audio_ms, audio_sec)
                yield _pcm16(wav)
                return

            def _gen(text: str) -> tuple[torch.Tensor, float]:
                return engine.generate(text, req)

            total_audio = 0.0
            for r in stream_chunks(req.text, generate=_gen, clock=time.perf_counter, cfg=cfg):
                if first_audio_ms is None:
                    first_audio_ms = (time.perf_counter() - t_req) * 1000
                total_audio += r.audio_sec
                _log_chunk(r, first_audio_ms)
                yield _pcm16(r.audio)
            log.info("stream done: ttfa=%.0fms total_audio=%.2fs", first_audio_ms or 0, total_audio)

    return StreamingResponse(body(), media_type=media)


def _log_chunk(r: ChunkResult, ttfa_ms: float) -> None:
    tag = " STARVED" if r.starved else ""
    if r.index == 0:
        log.info("chunk 0: ttfa=%.0fms gen=%.0fms audio=%.2fs rtf=%.2f%s",
                 ttfa_ms, r.gen_time * 1000, r.audio_sec, r.rtf, tag)
    else:
        log.info("chunk %d: gen=%.0fms audio=%.2fs buf=%.2f→%.2f rtf=%.2f%s",
                 r.index, r.gen_time * 1000, r.audio_sec,
                 r.buffer_before, r.buffer_after, r.rtf, tag)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host=BIND, port=PORT)