diff --git a/stacks/chatterbox-fast/.gitignore b/stacks/chatterbox-fast/.gitignore new file mode 100644 index 0000000..75c6182 --- /dev/null +++ b/stacks/chatterbox-fast/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.pyc +.pytest_cache/ diff --git a/stacks/chatterbox-fast/app.py b/stacks/chatterbox-fast/app.py index 91f79d4..f476e6b 100644 --- a/stacks/chatterbox-fast/app.py +++ b/stacks/chatterbox-fast/app.py @@ -164,9 +164,8 @@ class Engine: 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.""" + def _raw_wav(self, text: str, knobs: "TTSRequest") -> torch.Tensor: + """One generate pass → wav tensor [1,T], CUDA-synced for honest timing.""" with torch.inference_mode(): wav = self.model.generate( text, @@ -177,8 +176,30 @@ class Engine: ) if DEVICE.startswith("cuda"): torch.cuda.synchronize() - audio_sec = wav.shape[-1] / self.sr - return wav, audio_sec + return wav + + def generate(self, text: str, knobs: "TTSRequest") -> tuple[torch.Tensor, float]: + """Synthesize ``text`` → (wav [1,T], audio_seconds).""" + wav = self._raw_wav(text, knobs) + return wav, wav.shape[-1] / self.sr + + def generate_primed( + self, content: str, context: str, knobs: "TTSRequest" + ) -> tuple[torch.Tensor, float]: + """Context-primed generate (plan §1.6): render ``context + content`` + together so ``content``'s prosody knows what preceded it, then discard the + context audio. The cut snaps to the inter-sentence pause (energy minimum) + near the context's solo duration, with a short fade-in to kill any click. + + Costs two passes (context-solo to locate the cut, then the joint) — the + scheduler budgets for this via prime_cost_factor.""" + ctx_wav = self._raw_wav(context, knobs) + d_ctx = ctx_wav.shape[-1] / self.sr + joint = self._raw_wav(f"{context} {content}", knobs) + cut = _cut_at_pause(joint, self.sr, d_ctx) + kept = joint[..., cut:].clone() + _fade_in(kept, self.sr) + return kept, kept.shape[-1] / self.sr engine = Engine() @@ -206,12 +227,50 @@ class TTSRequest(BaseModel): top_p: float = 0.95 top_k: int = 1000 repetition_penalty: float = 1.2 + # Context-priming at joins (plan §1.6) — opt-in for A/B. + prime: bool = False + prime_first_n: int = 2 # how many early joins to prime when prime=True + # 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 _cut_at_pause( + joint: torch.Tensor, sr: int, approx_sec: float, + window_sec: float = 0.4, frame_sec: float = 0.02, +) -> int: + """Sample index to cut the discarded context off ``joint``. Searches ±window + around the context's solo duration for the lowest-energy 20 ms frame — the + inter-sentence pause — so the seam lands in silence, not mid-phone.""" + audio = joint.reshape(-1) + n = audio.shape[0] + center = int(approx_sec * sr) + w = int(window_sec * sr) + frame = max(1, int(frame_sec * sr)) + lo = max(0, center - w) + hi = min(n - frame, center + w) + if hi <= lo: + return min(center, n) + best_i, best_e = center, float("inf") + i = lo + while i < hi: + e = float((audio[i:i + frame] ** 2).mean()) + if e < best_e: + best_e, best_i = e, i + i += frame + return best_i + + +def _fade_in(wav: torch.Tensor, sr: int, fade_sec: float = 0.005) -> None: + """In-place short fade-in to remove any click at the cut seam.""" + fade = int(fade_sec * sr) + if wav.shape[-1] > fade > 0: + ramp = torch.linspace(0.0, 1.0, fade, device=wav.device, dtype=wav.dtype) + wav[..., :fade] *= ramp + + 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(" ChunkConfig: cfg.margin_first = req.margin_first if req.rtf_prior is not None: cfg.rtf_prior = req.rtf_prior + if req.prime: + cfg.prime_first_n = req.prime_first_n return cfg @@ -289,7 +350,9 @@ def tts(req: TTSRequest) -> StreamingResponse: yield _pcm16(wav) return - def _gen(text: str) -> tuple[torch.Tensor, float]: + def _gen(text: str, context: str | None) -> tuple[torch.Tensor, float]: + if context: + return engine.generate_primed(text, context, req) return engine.generate(text, req) total_audio = 0.0 @@ -305,7 +368,7 @@ def tts(req: TTSRequest) -> StreamingResponse: def _log_chunk(r: ChunkResult, ttfa_ms: float) -> None: - tag = " STARVED" if r.starved else "" + tag = (" PRIMED" if r.primed else "") + (" 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) diff --git a/stacks/chatterbox-fast/bench.py b/stacks/chatterbox-fast/bench.py index 0cc12f2..410c3f5 100644 --- a/stacks/chatterbox-fast/bench.py +++ b/stacks/chatterbox-fast/bench.py @@ -33,8 +33,10 @@ DEFAULT_TEXT = ( ) -def run(host: str, text: str, out: str, *, oneshot: bool, voice: str | None) -> None: - payload = {"text": text, "format": "pcm", "stream": not oneshot} +def run(host: str, text: str, out: str, *, oneshot: bool, voice: str | None, + prime: bool, prime_n: int) -> None: + payload = {"text": text, "format": "pcm", "stream": not oneshot, + "prime": prime, "prime_first_n": prime_n} if voice: payload["voice"] = voice req = urllib.request.Request( @@ -98,8 +100,11 @@ def main() -> None: ap.add_argument("--out", default="/refs/_fast.wav") ap.add_argument("--voice", default=None) ap.add_argument("--oneshot", action="store_true", help="whole-text one-shot baseline") + ap.add_argument("--prime", action="store_true", help="context-prime the early joins") + ap.add_argument("--prime-n", type=int, default=2, help="how many early joins to prime") args = ap.parse_args() - run(args.host, args.text, args.out, oneshot=args.oneshot, voice=args.voice) + run(args.host, args.text, args.out, oneshot=args.oneshot, voice=args.voice, + prime=args.prime, prime_n=args.prime_n) if __name__ == "__main__": diff --git a/stacks/chatterbox-fast/scheduler.py b/stacks/chatterbox-fast/scheduler.py index ea09c35..b4659ce 100644 --- a/stacks/chatterbox-fast/scheduler.py +++ b/stacks/chatterbox-fast/scheduler.py @@ -63,6 +63,24 @@ class ChunkConfig: # granularity — plan §1.1). max_first_sec: float = 2.0 + # Context-priming at joins (plan §1.6). Prime the first N joins (chunks + # 1..N) by prepending the prior sentence as backward prosodic context, then + # discarding its audio. 0 ⇒ off. Priming runs a 2nd "context-solo" generate, + # so a primed chunk costs ~(2·context + content)/rtf. Priming is AFFORDABILITY- + # GATED: a chunk is only primed when that cost fits the buffer; otherwise it + # falls back to a cold (unprimed) generate, so priming can never starve the + # stream. The earliest joins (smallest buffer) thus self-skip until the buffer + # has ratcheted up enough to pay for the extra pass. + prime_first_n: int = 0 + + # Priming headroom: only prime when the buffer is at least this multiple of + # the primed cost, so the 2nd pass doesn't flatten the buffer below the slack + # the NEXT chunk needs to absorb RTF-estimate error. At 1.5, priming fires on + # the early joins for any GPU at/above the rtf_prior floor (3.4 = 3090; A6000 + # ~3.8–4.0), and on a slower-than-fleet GPU it self-skips entirely (degrades + # to cold/unprimed) rather than starving. + prime_buffer_factor: float = 1.5 + # ── result record ───────────────────────────────────────────────────────── @@ -82,6 +100,7 @@ class ChunkResult: 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 + primed: bool = False # context-priming was applied to this chunk @property def starved(self) -> bool: @@ -137,6 +156,11 @@ def _est_gen_time(text: str, *, rtf: float, sec_per_char: float) -> float: return (len(text) * sec_per_char) / rtf +def _est_primed_gen_time(content: str, context: str, *, rtf: float, sec_per_char: float) -> float: + """Primed cost = context-solo pass + joint(context+content) pass.""" + return ((2 * len(context) + len(content)) * sec_per_char) / rtf + + def plan_chunk( remaining: Sequence[str], buffer_remaining: float, @@ -144,19 +168,27 @@ def plan_chunk( margin: float, rtf: float, sec_per_char: float, + prime_context: str | None = None, ) -> 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. + + When ``prime_context`` is set the chunk will be context-primed, so packing + uses the (larger) primed cost estimate to leave room for the 2nd pass. """ 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: + if prime_context is not None: + est = _est_primed_gen_time(candidate, prime_context, rtf=rtf, sec_per_char=sec_per_char) + else: + est = _est_gen_time(candidate, rtf=rtf, sec_per_char=sec_per_char) + if est > budget: break chunk.append(remaining[i]) i += 1 @@ -194,8 +226,10 @@ def _ema(old: float, new: float, alpha: float) -> float: # ── the online loop ─────────────────────────────────────────────────────── -# generate(text) -> (audio_payload, audio_seconds) -GenerateFn = Callable[[str], "tuple[object, float]"] +# generate(text, context) -> (audio_payload, audio_seconds) +# context is the prior sentence to prime backward prosody (discarded by the +# generator), or None for an unprimed chunk. +GenerateFn = Callable[[str, "str | None"], "tuple[object, float]"] ClockFn = Callable[[], float] @@ -227,6 +261,7 @@ def stream_chunks( buffer_remaining = 0.0 remaining: list[str] = units index = 0 + prev_text: str | None = None while remaining: first = index == 0 @@ -236,14 +271,34 @@ def stream_chunks( remaining = relieve_leader( remaining, buffer_remaining, rtf=rtf, sec_per_char=sec_per_char ) - margin = cfg.margin_first if first else cfg.margin + # margin_first tightens the FIRST transition (planning chunk 1 off chunk + # 0's small buffer — highest starvation risk, plan §1.5). + margin = cfg.margin_first if index == 1 else cfg.margin + + # Context-priming (plan §1.6): eligible on chunks 1..N, affordability-gated + # so the 2nd pass can never starve the buffer — fall back to cold otherwise. + context: str | None = None + if 0 < index <= cfg.prime_first_n and prev_text: + prev_units = split_sentences(prev_text) + candidate_ctx = prev_units[-1] if prev_units else prev_text + min_primed = _est_primed_gen_time( + remaining[0], candidate_ctx, rtf=rtf, sec_per_char=sec_per_char + ) + if min_primed * cfg.prime_buffer_factor <= buffer_remaining: + context = candidate_ctx + chunk_text, remaining = plan_chunk( - remaining, buffer_remaining, margin=margin, rtf=rtf, sec_per_char=sec_per_char + remaining, buffer_remaining, margin=margin, rtf=rtf, + sec_per_char=sec_per_char, prime_context=context, ) - est_gen = _est_gen_time(chunk_text, rtf=rtf, sec_per_char=sec_per_char) + if context is not None: + est_gen = _est_primed_gen_time(chunk_text, context, rtf=rtf, sec_per_char=sec_per_char) + else: + est_gen = _est_gen_time(chunk_text, rtf=rtf, sec_per_char=sec_per_char) + primed = context is not None t0 = clock() - audio, audio_sec = generate(chunk_text) + audio, audio_sec = generate(chunk_text, context) gen_time = clock() - t0 # Starvation: did the buffer run dry while we generated this chunk? @@ -273,5 +328,7 @@ def stream_chunks( drained=drained, rtf=rtf, sec_per_char=sec_per_char, + primed=primed, ) + prev_text = chunk_text index += 1 diff --git a/stacks/chatterbox-fast/test_scheduler.py b/stacks/chatterbox-fast/test_scheduler.py index d9f9e4c..dec44e5 100644 --- a/stacks/chatterbox-fast/test_scheduler.py +++ b/stacks/chatterbox-fast/test_scheduler.py @@ -45,13 +45,19 @@ class FakeClock: 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. + Audio duration is proportional to content length; generation costs + ``audio_sec / true_rtf`` of (simulated) wall-clock, advancing the clock. A + primed chunk (context given) costs extra: a context-solo pass plus the + context portion of the joint pass — modelling the ~2× cost the scheduler + must budget for. """ - def generate(text: str): - audio_sec = len(text) * sec_per_char - clock.t += audio_sec / true_rtf + def generate(text: str, context: str | None): + audio_sec = len(text) * sec_per_char # content only (context discarded) + gen_audio = audio_sec + if context: + gen_audio += 2 * len(context) * sec_per_char # ctx-solo + ctx in joint + clock.t += gen_audio / true_rtf return None, audio_sec return generate @@ -185,6 +191,43 @@ def test_full_text_reconstructed(): assert joined == " ".join(split_sentences(PARAGRAPH)) +def test_priming_fires_but_never_on_chunk_zero(): + """With prime_first_n=2, priming is best-effort (affordability-gated): it fires + on at least one early join at fleet RTF, and NEVER on chunk 0 (latency-critical).""" + cfg = ChunkConfig(prime_first_n=2) + results = run(true_rtf=3.8, cfg=cfg) + assert results[0].primed is False + assert any(r.primed for r in results[1:]), "expected at least one primed join" + # Priming only ever lands on chunks 1..N. + assert all(not r.primed for r in results if r.index > cfg.prime_first_n) + + +def test_priming_never_starves_at_fleet_rtf(): + """Priming must keep the no-starvation guarantee for any GPU at/above the + rtf_prior floor (3090 ~3.4, A6000 ~3.8–4.0).""" + cfg = ChunkConfig(prime_first_n=2) + for rtf in (4.0, 3.8, 3.4): + results = run(true_rtf=rtf, cfg=cfg) + assert all(not r.starved for r in results), ( + rtf, [(r.index, r.drained) for r in results if r.starved] + ) + + +def test_priming_self_skips_on_slow_gpu_no_starvation(): + """On a slower-than-fleet GPU (RTF below the prior), priming gracefully + self-skips rather than starving — the guarantee holds, priming just stops.""" + cfg = ChunkConfig(prime_first_n=2) + results = run(true_rtf=3.0, cfg=cfg) + assert all(not r.starved for r in results) + assert not any(r.primed for r in results) # degraded to cold + + +def test_priming_off_by_default(): + """Default config primes nothing (so context is never requested).""" + results = run(true_rtf=3.8) + assert all(not r.primed for r in results) + + def _main(): results = run(true_rtf=3.8) print(f"{'idx':>3} {'chars':>5} {'audio_s':>8} {'gen_s':>7} "