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,277 @@
|
||||
"""Adaptive buffer-ratchet chunk scheduler for chatterbox-fast.
|
||||
|
||||
This module is the *meat* of the streaming engine and is deliberately free of
|
||||
any GPU / torch / chatterbox / FastAPI imports so the no-starvation guarantee
|
||||
can be validated as a pure simulation (see ``test_scheduler.py``).
|
||||
|
||||
The design (docs/design/chatterbox-fast-plan.md §1):
|
||||
|
||||
* Chunk 1 = the first sentence alone — generated and emitted immediately so
|
||||
first-audio latency is minimal.
|
||||
* While chunk N plays, generate chunk N+1 by greedily accumulating WHOLE
|
||||
sentences until the next sentence would push estimated gen-time past
|
||||
``margin × audio_buffered_remaining``. Never split mid-sentence — each
|
||||
chunk stays prosodically self-coherent and joins land at natural pauses.
|
||||
* Because Chatterbox runs faster than realtime (RTF > 1), every chunk's
|
||||
playback buys wall-clock for a larger next chunk, so chunks ratchet up
|
||||
~3× and after 2-3 joins the rest of the paragraph is one big near-full-
|
||||
context chunk. Context loss is confined to those 2-3 sentence-boundary
|
||||
joins.
|
||||
* Drive off the *measured* realtime factor and sec-per-char, not constants —
|
||||
track them live and self-correct.
|
||||
|
||||
The scheduler is *online*: chunk N+1's boundary depends on the RTF measured
|
||||
while generating chunk N, so it cannot be precomputed. ``stream_chunks`` runs
|
||||
the loop with ``generate`` and ``clock`` injected, which is what makes the
|
||||
simulation possible.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Iterator, Sequence
|
||||
|
||||
# ── tunables ──────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkConfig:
|
||||
"""Scheduler tunables. Priors are conservative; live measurement corrects."""
|
||||
|
||||
# Realtime factor prior. 3090 ~3.4×, A6000 ~3.8× (plan §2). Start
|
||||
# conservative — under-estimating RTF makes chunks smaller and safer.
|
||||
rtf_prior: float = 3.4
|
||||
|
||||
# Safety fraction of the buffer to spend on the next chunk's generation.
|
||||
# margin < 1 leaves headroom so estimate error doesn't starve the stream.
|
||||
margin: float = 0.8
|
||||
|
||||
# Tighter margin on the FIRST transition (chunk 1 → chunk 2): the buffer is
|
||||
# smallest there, so starvation risk is highest (plan §1.5: ~0.6-0.7).
|
||||
margin_first: float = 0.65
|
||||
|
||||
# Seconds of audio per character of text. ~0.060 ≈ 16.6 chars/sec speech.
|
||||
# Calibrated live per request.
|
||||
sec_per_char_prior: float = 0.060
|
||||
|
||||
# EMA weight on the newest measurement when updating rtf / sec_per_char.
|
||||
ema_alpha: float = 0.4
|
||||
|
||||
# If the first sentence's estimated audio exceeds this, clause-split it to
|
||||
# protect first-audio latency (the ONLY place we split below sentence
|
||||
# granularity — plan §1.1).
|
||||
max_first_sec: float = 2.0
|
||||
|
||||
|
||||
# ── result record ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkResult:
|
||||
"""Telemetry + payload for one emitted chunk."""
|
||||
|
||||
index: int
|
||||
text: str
|
||||
audio: object # opaque payload from generate() (wav tensor, fake, …)
|
||||
audio_sec: float
|
||||
gen_time: float
|
||||
est_gen: float # gen-time the scheduler predicted before generating
|
||||
buffer_before: float # unplayed audio (s) when this chunk's gen started
|
||||
buffer_after: float # unplayed audio (s) once this chunk is emitted
|
||||
drained: float # seconds the buffer ran dry during gen (>0 ⇒ starvation)
|
||||
rtf: float # measured RTF after this chunk
|
||||
sec_per_char: float # measured sec/char after this chunk
|
||||
|
||||
@property
|
||||
def starved(self) -> bool:
|
||||
return self.drained > 1e-9
|
||||
|
||||
|
||||
# ── text splitting ────────────────────────────────────────────────────────
|
||||
|
||||
# Split after sentence-final punctuation when the next non-space looks like a
|
||||
# new sentence start: a capital/digit (optionally behind an open quote) or an
|
||||
# inline tag like "[laugh]". Pragmatic, not perfect: abbreviations ("Mr.",
|
||||
# "e.g.") can over-split — a Phase-2 watch-out, harmless to quality (just an
|
||||
# extra natural-pause join).
|
||||
_SENTENCE_END = re.compile(r'(?<=[.!?])["\')\]]*\s+(?=["\'(]*(?:[A-Z0-9]|\[))')
|
||||
|
||||
# Clause boundaries for the first-sentence latency fallback only.
|
||||
_CLAUSE_END = re.compile(r'(?<=[,;:])\s+')
|
||||
|
||||
|
||||
def split_sentences(text: str) -> list[str]:
|
||||
"""Split into sentence units, preserving punctuation. Whitespace-collapsed."""
|
||||
text = " ".join(text.split())
|
||||
if not text:
|
||||
return []
|
||||
return [s for s in (p.strip() for p in _SENTENCE_END.split(text)) if s]
|
||||
|
||||
|
||||
def _split_clauses(sentence: str) -> list[str]:
|
||||
return [c for c in (p.strip() for p in _CLAUSE_END.split(sentence)) if c]
|
||||
|
||||
|
||||
def protect_first_audio(units: list[str], cfg: ChunkConfig) -> list[str]:
|
||||
"""Clause-split the first unit if it's too long to hit the first-audio target.
|
||||
|
||||
Only the *leading* unit is split; the remainder is untouched. If the first
|
||||
sentence has no clause boundary we accept the latency rather than split
|
||||
mid-clause (quality > latency once we're past the budget).
|
||||
"""
|
||||
if not units:
|
||||
return units
|
||||
if len(units[0]) * cfg.sec_per_char_prior <= cfg.max_first_sec:
|
||||
return units
|
||||
pieces = _split_clauses(units[0])
|
||||
if len(pieces) <= 1:
|
||||
return units # nothing to split on; keep the long first sentence whole
|
||||
return pieces + units[1:]
|
||||
|
||||
|
||||
# ── chunk planning ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _est_gen_time(text: str, *, rtf: float, sec_per_char: float) -> float:
|
||||
return (len(text) * sec_per_char) / rtf
|
||||
|
||||
|
||||
def plan_chunk(
|
||||
remaining: Sequence[str],
|
||||
buffer_remaining: float,
|
||||
*,
|
||||
margin: float,
|
||||
rtf: float,
|
||||
sec_per_char: float,
|
||||
) -> tuple[str, list[str]]:
|
||||
"""Greedily accumulate whole units until the next would blow the budget.
|
||||
|
||||
Always returns at least one unit (never empty, never splits a unit). With
|
||||
``buffer_remaining == 0`` (the first chunk) the budget is 0, so exactly the
|
||||
first unit is taken — which is the latency-critical chunk-1 rule.
|
||||
"""
|
||||
budget = margin * buffer_remaining
|
||||
chunk = [remaining[0]]
|
||||
i = 1
|
||||
while i < len(remaining):
|
||||
candidate = " ".join(chunk + [remaining[i]])
|
||||
if _est_gen_time(candidate, rtf=rtf, sec_per_char=sec_per_char) > budget:
|
||||
break
|
||||
chunk.append(remaining[i])
|
||||
i += 1
|
||||
return " ".join(chunk), list(remaining[i:])
|
||||
|
||||
|
||||
def relieve_leader(
|
||||
remaining: list[str],
|
||||
buffer_remaining: float,
|
||||
*,
|
||||
rtf: float,
|
||||
sec_per_char: float,
|
||||
) -> list[str]:
|
||||
"""Expose a too-big leading sentence's clause boundaries to avoid starvation.
|
||||
|
||||
If generating ``remaining[0]`` alone would overrun the buffer (an audible
|
||||
gap), clause-split it in place so ``plan_chunk`` can pack clause-pieces up to
|
||||
the buffer — joins then land on commas (natural pauses) instead of producing
|
||||
a gap. If the sentence has no clause boundary it's returned unchanged: we do
|
||||
NOT split mid-clause (design rule), and the caller accepts + flags the rare
|
||||
starvation in telemetry.
|
||||
"""
|
||||
leader = remaining[0]
|
||||
if _est_gen_time(leader, rtf=rtf, sec_per_char=sec_per_char) <= buffer_remaining:
|
||||
return remaining
|
||||
pieces = _split_clauses(leader)
|
||||
if len(pieces) <= 1:
|
||||
return remaining
|
||||
return pieces + remaining[1:]
|
||||
|
||||
|
||||
def _ema(old: float, new: float, alpha: float) -> float:
|
||||
return (1 - alpha) * old + alpha * new
|
||||
|
||||
|
||||
# ── the online loop ───────────────────────────────────────────────────────
|
||||
|
||||
# generate(text) -> (audio_payload, audio_seconds)
|
||||
GenerateFn = Callable[[str], "tuple[object, float]"]
|
||||
ClockFn = Callable[[], float]
|
||||
|
||||
|
||||
def stream_chunks(
|
||||
text: str,
|
||||
*,
|
||||
generate: GenerateFn,
|
||||
clock: ClockFn,
|
||||
cfg: ChunkConfig | None = None,
|
||||
) -> Iterator[ChunkResult]:
|
||||
"""Run the adaptive buffer-ratchet schedule, yielding one ChunkResult per chunk.
|
||||
|
||||
``generate`` does the actual synthesis (real GPU or a simulation) and
|
||||
returns its opaque audio payload plus the audio's duration in seconds.
|
||||
``clock`` returns monotonically increasing seconds; ``gen_time`` is measured
|
||||
as the clock delta around ``generate``.
|
||||
|
||||
Buffer model (all in audio-seconds):
|
||||
* Playback begins when chunk 1 arrives, so no drain occurs during chunk 1.
|
||||
* For chunk N≥2, the client plays ``gen_time`` seconds of buffered audio
|
||||
while we generate it, then we add this chunk's audio. If ``gen_time``
|
||||
exceeds the buffer the stream starved (``drained`` > 0).
|
||||
"""
|
||||
cfg = cfg or ChunkConfig()
|
||||
|
||||
units = protect_first_audio(split_sentences(text), cfg)
|
||||
rtf = cfg.rtf_prior
|
||||
sec_per_char = cfg.sec_per_char_prior
|
||||
buffer_remaining = 0.0
|
||||
remaining: list[str] = units
|
||||
index = 0
|
||||
|
||||
while remaining:
|
||||
first = index == 0
|
||||
if not first:
|
||||
# If the next indivisible sentence would starve the buffer, expose
|
||||
# its clause boundaries so plan_chunk can pack to the buffer.
|
||||
remaining = relieve_leader(
|
||||
remaining, buffer_remaining, rtf=rtf, sec_per_char=sec_per_char
|
||||
)
|
||||
margin = cfg.margin_first if first else cfg.margin
|
||||
chunk_text, remaining = plan_chunk(
|
||||
remaining, buffer_remaining, margin=margin, rtf=rtf, sec_per_char=sec_per_char
|
||||
)
|
||||
est_gen = _est_gen_time(chunk_text, rtf=rtf, sec_per_char=sec_per_char)
|
||||
|
||||
t0 = clock()
|
||||
audio, audio_sec = generate(chunk_text)
|
||||
gen_time = clock() - t0
|
||||
|
||||
# Starvation: did the buffer run dry while we generated this chunk?
|
||||
drained = 0.0 if first else max(0.0, gen_time - buffer_remaining)
|
||||
|
||||
# Live self-correction.
|
||||
if gen_time > 0:
|
||||
rtf = _ema(rtf, audio_sec / gen_time, cfg.ema_alpha)
|
||||
if chunk_text:
|
||||
sec_per_char = _ema(sec_per_char, audio_sec / len(chunk_text), cfg.ema_alpha)
|
||||
|
||||
buffer_before = buffer_remaining
|
||||
if first:
|
||||
buffer_remaining = audio_sec
|
||||
else:
|
||||
buffer_remaining = max(0.0, buffer_remaining - gen_time) + audio_sec
|
||||
|
||||
yield ChunkResult(
|
||||
index=index,
|
||||
text=chunk_text,
|
||||
audio=audio,
|
||||
audio_sec=audio_sec,
|
||||
gen_time=gen_time,
|
||||
est_gen=est_gen,
|
||||
buffer_before=buffer_before,
|
||||
buffer_after=buffer_remaining,
|
||||
drained=drained,
|
||||
rtf=rtf,
|
||||
sec_per_char=sec_per_char,
|
||||
)
|
||||
index += 1
|
||||
Reference in New Issue
Block a user