feat(chatterbox-fast): context-priming at joins (§1.6, opt-in)

Prime early joins by prepending the prior sentence as backward prosodic context,
generating context+content together, then discarding the context audio. The cut
snaps to the inter-sentence pause (energy-minimum search around the context's
solo duration) with a 5ms fade-in to kill any seam click (app: _cut_at_pause /
_fade_in / Engine.generate_primed). Opt-in via request `prime` (default off).

Scheduler: priming is AFFORDABILITY-GATED so it can never starve. A primed chunk
costs ~(2·context + content)/rtf (a 2nd context-solo pass); a chunk is only primed
when buffer ≥ prime_buffer_factor (1.5) × that cost, else it falls back to a cold
generate. Consequences proven in the GPU-free sim (17 tests):
  - fires on early joins for any GPU at/above rtf_prior (3.4 = 3090; A6000 ~3.8-4.0)
  - self-skips (degrades to cold) on a slower-than-fleet GPU rather than starving
  - never primes chunk 0 (latency-critical)
Also fixed a latent Phase-1 bug: margin_first was applied at chunk 0 (budget always
0 there) so it never did anything — now applied at chunk 1 (the first transition).

Live A/B on irv-ml1 (A6000, GLaDOS): TTFB unaffected (445 vs 467ms), no starvation;
priming fired on chunk 2 (gen 1.6s for the doubled pass). On typical text exactly
ONE early join safely primes — priming chunk 2 flattens the buffer so later/larger
chunks no longer clear the safety gate. Samples: ~/chatterbox-ab/_p2_{cold,primed}.wav.

stacks/chatterbox-fast/.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+.pytest_cache/

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("<i2").tobytes()
@@ -234,6 +293,8 @@ def _chunk_config(req: TTSRequest) -> 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)

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__":

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

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} "