vh/esh-pfi-infrastructure
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects Releases Wiki Activity

revert(chatterbox-fast): drop context-priming (§1.6) — discard-cut leaks context

Browse Source 
Revert the priming feature from d707439. Live A/B caught an audible artifact: the
context-priming discard-cut left part of the throwaway prefix in the output, so a
clause ("...without a trace of sarcasm,") was spoken an extra time.

Root cause is structural: generate() returns one finished waveform with no marker
for where the prefix ends, and the model renders the same prefix with different
timing when followed by content than when generated solo — so the duration-estimate
+ energy-minimum cut is a guess and can leave a sliver (or a whole clause) of prefix
in. A reliable cut would need token-level access (the abandoned native-streaming
arc) or a per-chunk ASR/alignment pass (heavy, still imperfect, eats the latency
budget). Fails the agreed bar: "keep only if it closes the gap without a seam."

Kept from d707439: the .gitignore (build artifacts). NOT re-applied: the bundled
margin_first fix — wiring it would shrink chunk 1 (more joins = worse coherence),
against the operator's priority, and margin=0.8 there is already starvation-safe.

Coherence loss at joins stays an accepted limitation; cold streaming was judged
"really good". Phase 1 + Phase 2 parity/perf untouched. Next: Phase 3 deploy.
This commit is contained in:
Vuong Hoang
2026-06-01 23:26:56 -07:00
parent d707439041
commit 090e70aed5
4 changed files with 22 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
stacks/chatterbox-fast/app.py
+7 -70
View File
@@ -164,8 +164,9 @@ class Engine:
if DEVICE.startswith("cuda"):
torch.cuda.synchronize()
def _raw_wav(self, text: str, knobs: "TTSRequest") -> torch.Tensor:
"""One generate pass → wav tensor [1,T], CUDA-synced for honest timing."""
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,
@@ -176,30 +177,8 @@ class Engine:
)
if DEVICE.startswith("cuda"):
torch.cuda.synchronize()
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
audio_sec = wav.shape[-1] / self.sr
return wav, audio_sec
engine = Engine()
@@ -227,50 +206,12 @@ 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()
@@ -293,8 +234,6 @@ 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
@@ -350,9 +289,7 @@ def tts(req: TTSRequest) -> StreamingResponse:
yield _pcm16(wav)
return
def _gen(text: str, context: str | None) -> tuple[torch.Tensor, float]:
if context:
return engine.generate_primed(text, context, req)
def _gen(text: str) -> tuple[torch.Tensor, float]:
return engine.generate(text, req)
total_audio = 0.0
@@ -368,7 +305,7 @@ def tts(req: TTSRequest) -> StreamingResponse:
def _log_chunk(r: ChunkResult, ttfa_ms: float) -> None:
tag = (" PRIMED" if r.primed else "") + (" STARVED" if r.starved else "")
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)
stacks/chatterbox-fast/bench.py
+3 -8
View File
@@ -33,10 +33,8 @@ DEFAULT_TEXT = (
)
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}
def run(host: str, text: str, out: str, *, oneshot: bool, voice: str | None) -> None:
payload = {"text": text, "format": "pcm", "stream": not oneshot}
if voice:
payload["voice"] = voice
req = urllib.request.Request(
@@ -100,11 +98,8 @@ 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,
prime=args.prime, prime_n=args.prime_n)
run(args.host, args.text, args.out, oneshot=args.oneshot, voice=args.voice)
if __name__ == "__main__":
stacks/chatterbox-fast/scheduler.py
+7 -64
View File
@@ -63,24 +63,6 @@ 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 ─────────────────────────────────────────────────────────
@@ -100,7 +82,6 @@ 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:
@@ -156,11 +137,6 @@ 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,
@@ -168,27 +144,19 @@ 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 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:
if _est_gen_time(candidate, rtf=rtf, sec_per_char=sec_per_char) > budget:
break
chunk.append(remaining[i])
i += 1
@@ -226,10 +194,8 @@ def _ema(old: float, new: float, alpha: float) -> float:
# ── the online loop ───────────────────────────────────────────────────────
# 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]"]
# generate(text) -> (audio_payload, audio_seconds)
GenerateFn = Callable[[str], "tuple[object, float]"]
ClockFn = Callable[[], float]
@@ -261,7 +227,6 @@ def stream_chunks(
buffer_remaining = 0.0
remaining: list[str] = units
index = 0
prev_text: str | None = None
while remaining:
first = index == 0
@@ -271,34 +236,14 @@ def stream_chunks(
remaining = relieve_leader(
remaining, buffer_remaining, rtf=rtf, sec_per_char=sec_per_char
)
# 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
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, prime_context=context,
remaining, buffer_remaining, margin=margin, 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
est_gen = _est_gen_time(chunk_text, rtf=rtf, sec_per_char=sec_per_char)
t0 = clock()
audio, audio_sec = generate(chunk_text, context)
audio, audio_sec = generate(chunk_text)
gen_time = clock() - t0
# Starvation: did the buffer run dry while we generated this chunk?
@@ -328,7 +273,5 @@ 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
+5 -48
View File
@@ -45,19 +45,13 @@ 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 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.
Audio duration is proportional to text length; generation costs
``audio_sec / true_rtf`` of (simulated) wall-clock, advancing the clock.
"""
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
def generate(text: str):
audio_sec = len(text) * sec_per_char
clock.t += audio_sec / true_rtf
return None, audio_sec
return generate
@@ -191,43 +185,6 @@ 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} "