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,202 @@
|
||||
"""GPU-free simulation of the adaptive-chunk scheduler.
|
||||
|
||||
Validates the two acceptance properties from the plan (§6) without a GPU:
|
||||
* NO STARVATION — the stream stays ahead of 1× playback.
|
||||
* RATCHET — chunks grow after the latency-critical first one.
|
||||
|
||||
Run directly (``python test_scheduler.py``) or under pytest.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from scheduler import (
|
||||
ChunkConfig,
|
||||
plan_chunk,
|
||||
protect_first_audio,
|
||||
split_sentences,
|
||||
stream_chunks,
|
||||
)
|
||||
|
||||
# A multi-sentence paragraph with varied lengths. The long sentences carry
|
||||
# commas (as natural prose does), so the clause-split starvation relief has
|
||||
# boundaries to work with.
|
||||
PARAGRAPH = (
|
||||
"The cake is a lie. "
|
||||
"I am being entirely sincere, without a trace of sarcasm, when I say "
|
||||
"that this is the single most important scientific breakthrough in the "
|
||||
"entire history of this facility. "
|
||||
"You will be baked, and then there will be cake. "
|
||||
"It is delicious and moist, assuming you survive the testing protocol, "
|
||||
"which the available data suggests you almost certainly will not. "
|
||||
"Goodbye."
|
||||
)
|
||||
|
||||
|
||||
class FakeClock:
|
||||
"""A clock the fake generator advances, so gen_time reflects simulated work."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.t = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.t
|
||||
|
||||
|
||||
def make_generator(clock: FakeClock, *, true_rtf: float, sec_per_char: float):
|
||||
"""A fake generate() that costs realistic wall-clock and returns audio_sec.
|
||||
|
||||
Audio duration is proportional to text length; generation costs
|
||||
``audio_sec / true_rtf`` of (simulated) wall-clock, advancing the clock.
|
||||
"""
|
||||
|
||||
def generate(text: str):
|
||||
audio_sec = len(text) * sec_per_char
|
||||
clock.t += audio_sec / true_rtf
|
||||
return None, audio_sec
|
||||
|
||||
return generate
|
||||
|
||||
|
||||
def run(true_rtf: float = 3.8, sec_per_char: float = 0.060, cfg: ChunkConfig | None = None):
|
||||
clock = FakeClock()
|
||||
gen = make_generator(clock, true_rtf=true_rtf, sec_per_char=sec_per_char)
|
||||
return list(stream_chunks(PARAGRAPH, generate=gen, clock=clock, cfg=cfg))
|
||||
|
||||
|
||||
# ── splitting ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_split_sentences_count():
|
||||
units = split_sentences(PARAGRAPH)
|
||||
assert units[0] == "The cake is a lie."
|
||||
assert units[-1] == "Goodbye."
|
||||
assert len(units) == 5
|
||||
|
||||
|
||||
def test_split_preserves_inline_tags():
|
||||
units = split_sentences("Hello there. [laugh] That was funny. Bye.")
|
||||
# The tag must not be torn from its sentence nor split on.
|
||||
assert any("[laugh]" in u for u in units)
|
||||
assert len(units) == 3
|
||||
|
||||
|
||||
def test_protect_first_audio_clause_splits_long_opener():
|
||||
cfg = ChunkConfig(max_first_sec=0.5) # force the guard
|
||||
long_open = ["This is a long opener, with a clause, and another clause.", "Short."]
|
||||
out = protect_first_audio(long_open, cfg)
|
||||
assert len(out) > len(long_open)
|
||||
assert out[0] == "This is a long opener,"
|
||||
|
||||
|
||||
def test_protect_first_audio_keeps_unsplittable_opener_whole():
|
||||
cfg = ChunkConfig(max_first_sec=0.1)
|
||||
out = protect_first_audio(["No clause boundaries here at all."], cfg)
|
||||
assert out == ["No clause boundaries here at all."]
|
||||
|
||||
|
||||
# ── chunk planning ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_first_chunk_is_single_unit():
|
||||
# Zero buffer ⇒ exactly the first unit, regardless of margin/rtf.
|
||||
units = split_sentences(PARAGRAPH)
|
||||
chunk, rest = plan_chunk(units, 0.0, margin=0.65, rtf=3.8, sec_per_char=0.06)
|
||||
assert chunk == units[0]
|
||||
assert len(rest) == len(units) - 1
|
||||
|
||||
|
||||
def test_plan_never_returns_empty():
|
||||
chunk, rest = plan_chunk(["Only one."], 0.0, margin=0.8, rtf=3.8, sec_per_char=0.06)
|
||||
assert chunk == "Only one."
|
||||
assert rest == []
|
||||
|
||||
|
||||
# ── the acceptance properties ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_no_starvation_nominal():
|
||||
"""Stream never starves when the RTF prior matches reality."""
|
||||
results = run(true_rtf=3.8)
|
||||
assert all(not r.starved for r in results), [
|
||||
(r.index, r.drained) for r in results if r.starved
|
||||
]
|
||||
|
||||
|
||||
def test_no_starvation_when_rtf_overestimated():
|
||||
"""Prior says 3.8× but the GPU only delivers 3.0× — live correction + margin
|
||||
must still keep the stream fed. This is the property that matters: the
|
||||
scheduler must be robust to an optimistic prior, not just a matching one."""
|
||||
results = run(true_rtf=3.0)
|
||||
assert all(not r.starved for r in results), [
|
||||
(r.index, r.drained) for r in results if r.starved
|
||||
]
|
||||
|
||||
|
||||
def test_no_starvation_slow_gpu():
|
||||
"""Even at 2.5× (well below prior) the margin absorbs it for this text."""
|
||||
results = run(true_rtf=2.5)
|
||||
assert all(not r.starved for r in results), [
|
||||
(r.index, r.drained) for r in results if r.starved
|
||||
]
|
||||
|
||||
|
||||
def test_chunks_ratchet_up():
|
||||
"""After the latency-critical first chunk, chunks grow (buffer-ratchet)."""
|
||||
results = run(true_rtf=3.8)
|
||||
assert len(results) >= 3, "paragraph should not collapse to one chunk"
|
||||
# Chunk 1 is the smallest (single sentence); the second chunk is larger.
|
||||
assert len(results[1].text) > len(results[0].text)
|
||||
# The paragraph consolidates: the last chunk carries multiple sentences.
|
||||
assert results[-1].audio_sec >= results[0].audio_sec
|
||||
|
||||
|
||||
def test_first_chunk_low_latency():
|
||||
"""First chunk is one short sentence ⇒ smallest gen time ⇒ fast first audio."""
|
||||
results = run(true_rtf=3.8)
|
||||
assert results[0].text == "The cake is a lie."
|
||||
# In the sim, chunk 0's gen_time IS the time-to-first-audio, and it must be
|
||||
# the smallest of all chunks (everything after it is ≥ a clause).
|
||||
assert results[0].gen_time == min(r.gen_time for r in results)
|
||||
assert results[0].gen_time < 0.4
|
||||
|
||||
|
||||
def test_unsplittable_long_sentence_flags_starvation():
|
||||
"""KNOWN LIMITATION (surfaced, not hidden): a long COMMA-LESS sentence right
|
||||
after a short opener cannot be clause-split, so we honor the never-split-mid-
|
||||
sentence rule and accept a brief gap — which MUST show up as drained>0 in
|
||||
telemetry so it is measurable, never silent."""
|
||||
clock = FakeClock()
|
||||
gen = make_generator(clock, true_rtf=3.8, sec_per_char=0.060)
|
||||
text = (
|
||||
"Hi. "
|
||||
"I am now going to speak one extremely long sentence with no clause "
|
||||
"boundaries at all so that nothing in here can ever be split apart by "
|
||||
"the scheduler no matter how hard it tries to find a comma."
|
||||
)
|
||||
results = list(stream_chunks(text, generate=gen, clock=clock))
|
||||
assert any(r.starved for r in results), "expected the gap to be flagged"
|
||||
assert any(r.drained > 0 for r in results)
|
||||
|
||||
|
||||
def test_full_text_reconstructed():
|
||||
"""Every sentence is emitted exactly once, in order."""
|
||||
results = run(true_rtf=3.8)
|
||||
joined = " ".join(r.text for r in results)
|
||||
assert joined == " ".join(split_sentences(PARAGRAPH))
|
||||
|
||||
|
||||
def _main():
|
||||
results = run(true_rtf=3.8)
|
||||
print(f"{'idx':>3} {'chars':>5} {'audio_s':>8} {'gen_s':>7} "
|
||||
f"{'est_s':>6} {'buf_before':>10} {'buf_after':>9} {'rtf':>5} {'drain':>6}")
|
||||
for r in results:
|
||||
print(f"{r.index:>3} {len(r.text):>5} {r.audio_sec:>8.2f} {r.gen_time:>7.3f} "
|
||||
f"{r.est_gen:>6.3f} {r.buffer_before:>10.2f} {r.buffer_after:>9.2f} "
|
||||
f"{r.rtf:>5.2f} {r.drained:>6.3f}")
|
||||
starved = [r.index for r in results if r.starved]
|
||||
print(f"\nchunks={len(results)} starved={starved or 'none'} "
|
||||
f"ttfa≈{results[0].gen_time:.3f}s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
_main()
|
||||
Reference in New Issue
Block a user