vh
76314624bb
ci / test (push) Has been cancelled
Sub-second streaming TTS on Chatterbox-Turbo via adaptive buffer-ratchet chunking. First audio in ~0.5s (vs ~5s one-shot) with no quality compromise — chunk joins land on natural sentence pauses and the stream converges to one large near-full-context chunk within 2-3 joins. Works because the engine runs faster than realtime; the no-starvation guarantee is proven in a GPU-free simulation (tests/test_scheduler.py). - chatterbox_fast/scheduler.py: the adaptive-chunk scheduler (pure logic, no GPU) - chatterbox_fast/app.py: FastAPI server (POST /tts streaming, /voices, /health) - bench.py: streaming client (ground-truth TTFB + starvation check) - Self-contained Dockerfile (slim base + chatterbox-tts from PyPI) - Three public-domain LibriVox starter voices baked in (see voices/ATTRIBUTION.md) MIT licensed.
344 lines
14 KiB
Python
344 lines
14 KiB
Python
"""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 chatterbox_fast.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
|
|
seed: int = 0 # 0 ⇒ random; a fixed seed repeats a one-shot take (see note below)
|
|
# 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, data_len: int | None = None) -> bytes:
|
|
"""WAV header. data_len=None → streaming (length unknown, 0xFFFFFFFF sizes,
|
|
player reads to EOF); an int → correct RIFF/data sizes for a complete file."""
|
|
data_size = 0xFFFFFFFF if data_len is None else data_len
|
|
riff_size = 0xFFFFFFFF if data_len is None else 36 + data_len
|
|
return (
|
|
b"RIFF" + struct.pack("<I", riff_size) + b"WAVE"
|
|
+ b"fmt " + struct.pack("<IHHIIHH", 16, 1, 1, sr, sr * 2, 2, 16)
|
|
+ b"data" + struct.pack("<I", data_size)
|
|
)
|
|
|
|
|
|
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)
|
|
# Seed once per request (under the lock). One-shot is then reproducible
|
|
# for a fixed seed + params; streaming is NOT — adaptive-chunk boundaries
|
|
# depend on live-measured RTF (wall-clock), so chunk splits vary run to run.
|
|
if req.seed:
|
|
torch.manual_seed(req.seed)
|
|
if DEVICE.startswith("cuda"):
|
|
torch.cuda.manual_seed_all(req.seed)
|
|
t_req = time.perf_counter()
|
|
first_audio_ms: float | None = None
|
|
|
|
if not req.stream:
|
|
# One-shot: the full length is known, so emit a correct-sized WAV
|
|
# header (a buffered consumer wants well-formed sizes).
|
|
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)
|
|
pcm = _pcm16(wav)
|
|
if req.format == "wav":
|
|
yield _wav_header(engine.sr, len(pcm))
|
|
yield pcm
|
|
return
|
|
|
|
# Streaming: length is unknown up front → open-ended WAV header.
|
|
if req.format == "wav":
|
|
yield _wav_header(engine.sr)
|
|
|
|
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)
|
|
|
|
|
|
def main() -> None:
|
|
"""Console entrypoint (`chatterbox-fast`) and Docker CMD."""
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host=BIND, port=PORT)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|