feat(chatterbox-fast): Phase 1 streaming server — adaptive-chunk scheduler
Build the streaming TTS server MVP per docs/design/chatterbox-fast-plan.md §4. - scheduler.py: adaptive buffer-ratchet chunker (the meat) — GPU-free pure logic. First sentence emitted alone for low TTFA, then chunks ratchet ~3x by packing whole sentences to margin x buffered-audio; drives off measured RTF + sec/char (EMA). relieve_leader() clause-splits a too-big mid-stream sentence to avoid starvation (joins land on commas); a long comma-less sentence is the one honored-but-flagged limitation. - test_scheduler.py: GPU-free simulation, 13 tests — asserts no-starvation (incl. overestimated RTF) and the ratchet. - app.py: FastAPI model holder + POST /tts StreamingResponse (raw PCM s16le default, wav optional, stream/oneshot) + GET /health. - bench.py: client — ground-truth TTFB + real 1x-consumer starvation check. Live test on irv-ml1 (turbo, A6000, GLaDOS voice): streaming TTFB 499ms vs oneshot 5230ms (~10x), stayed ahead of a 1x player (no starvation), ratchet 1.64->4.08->8.60->8.60s audio, measured RTF self-corrected 3.38->4.01. Kill the superseded docs/design/chatterbox-fast.md — its §5 windowed-token streaming was the abandoned native-frame-streaming arc; the adaptive-chunk plan supersedes it. Repoint persistent-memory + README at the canonical plan.
This commit is contained in:
@@ -0,0 +1,270 @@
|
||||
"""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"))
|
||||
|
||||
# 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."
|
||||
|
||||
|
||||
# ── 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
|
||||
|
||||
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.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 _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
|
||||
|
||||
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.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)
|
||||
Reference in New Issue
Block a user