Fix ZeroGPU continuation: bound the lead-in, normalize the tail, shrink the GPU window

Three deployed bugs, all reproduced on the Blackwell GPU and fixed:

1. Near-silence / "only 5s of new audio". A long source (e.g. 100s) filled
SA3's buffer and left only a tiny masked window, which the 8-step distilled
sampler rendered as near-silence (rms 0.001). engine now conditions on at
most MAX_LEAD_SECONDS (30s) of the clip's TAIL and masks [lead, lead+new],
keeping the generated region substantial. The splice still rejoins the tail
onto the full pristine original, so the listener hears the whole clip.
Verified: 100s->120 now yields rms 0.15 (was 0.001).

2. Length/mask. With the whole source in the buffer, "total < source" floored
to source+5s. With a bounded lead, new = total - source actually extends to
the requested finished length (100s->120 generates +20s).

3. Stream abort (BodyStreamBuffer was aborted). enhance/fingerprint/decode ran
INSIDE @spaces.GPU, burning the ZeroGPU window before generation started.
Only the diffusion call is now wrapped by @spaces.GPU; all CPU prep and the
splice run outside it, so the GPU task starts and finishes fast (sub-second).

Also normalize the generated tail by its own peak (not the whole buffer) so a
loud lead transient can't divide the continuation down toward silence. Engine
unit tests updated to the bounded-lead contract.

app.py

@@ -128,50 +128,53 @@ def analyze_on_upload(audio_path):
                 gr.update(interactive=False))
 
 
-@spaces.GPU(duration=240)
+@spaces.GPU(duration=120)
+def _continue_on_gpu(listen_path, total_seconds, vibe):
+    """ONLY the SA3 diffusion call runs inside the GPU window. All CPU work —
+    enhancement, key/tempo analysis, decode, splice — happens OUTSIDE it (in
+    `finish_song`), so the scarce ZeroGPU allocation is spent generating instead
+    of decoding/analyzing audio. That's the stream-abort fix: the GPU task now
+    starts and finishes fast instead of sitting through a slow analysis preamble
+    until the browser aborts the stream.
+
+    Pin the seed to the lab's known-good draw. SA3 is generative: a random seed
+    gives a different (often weaker) continuation every run. seed=7 produced the
+    verified lab_out/sa3 takes, so the app reproduces that instead of re-rolling.
+    """
+    return engine.continue_audio(
+        listen_path, total_seconds=int(total_seconds),
+        prompt=(vibe or "").strip(), seed=7)
+
+
 def finish_song(audio_path, total_seconds, vibe, remaster,
                 progress=gr.Progress()):
-    """The whole job inside one GPU window: listen -> continue (SA3) -> splice.
+    """Orchestrate the job: CPU prep -> GPU continuation -> CPU splice.
     Returns (output_wav_path, summary_markdown)."""
     if not audio_path:
         raise gr.Error("Upload a clip (or load the PUSHBACK demo) first.")
 
     total_seconds = int(total_seconds)
-    progress(0.05, desc="Listening to your clip…")
 
-    # analyze + build the conditioning feed from a cleaned copy
+    # --- CPU prep (outside the GPU window) ---
+    progress(0.05, desc="Listening to your clip…")
     listen_path = enhance_to_tempfile(audio_path)
     info = fingerprint(listen_path)
-
     # the pristine original — what the listener hears for the first stretch
     original, sr = librosa.load(audio_path, sr=None, mono=False)
     if remaster:
-        progress(0.12, desc="Remastering your part…")
+        progress(0.15, desc="Remastering your part…")
         original = enhance_audio(original, sr)
 
-    # SA3 is a single call; map its stages onto the bar
-    def _on_stage(stage):
-        marks = {
-            "reading": (0.22, "Reading key, tempo & groove…"),
-            "composing": (0.45, "Composing the continuation…"),
-            "finalizing": (0.82, "Rendering 44.1 kHz stereo…"),
-        }
-        if stage in marks:
-            frac, desc = marks[stage]
-            progress(frac, desc=desc)
-
+    # --- GPU continuation (the ONLY @spaces.GPU call) ---
+    progress(0.35, desc="Composing the continuation…")
     try:
-        # Pin the seed to the lab's known-good draw. SA3 is generative: a
-        # random seed gives a different (and often weaker) continuation every
-        # run. seed=7 is the draw that produced the verified lab_out/sa3 takes,
-        # so the app reproduces that result instead of re-rolling each time.
-        new_tail, source_seconds, SR = engine.continue_audio(
-            listen_path, total_seconds=total_seconds,
-            prompt=(vibe or "").strip(), seed=7, progress=_on_stage)
+        new_tail, source_seconds, SR = _continue_on_gpu(
+            listen_path, total_seconds, vibe)
     except ValueError as e:
-        # e.g. the clip is too long to continue under the 120s cap
+        # e.g. the clip is a full-length track, not a clip to continue
         raise gr.Error(str(e))
 
+    # --- CPU splice + write (outside the GPU window) ---
     progress(0.9, desc="Splicing onto your original…")
     out = stitch.stitch(original, sr, new_tail, source_seconds)
 

engine.py

@@ -9,14 +9,24 @@ audio — true long-form continuation, 44.1 kHz stereo, no multi-pass chaining,
 no energy guards, no re-roll logic.
 
 This module is the whole generation core. It returns ONLY the newly generated
-tail (the model's [source_end, total] region) plus the source length in seconds;
-`stitch.py` joins that tail onto the user's *pristine* original so the real
-recording (and any vocals) plays untouched up to the seam.
+tail plus the source length in seconds; `stitch.py` joins that tail onto the
+user's *pristine* original so the real recording (and any vocals) plays
+untouched up to the seam.
+
+Bounded lead-in (the deployed-bug fix): SA3 only needs a short run-up to know
+where the song is going. We therefore condition on at most MAX_LEAD_SECONDS of
+the clip's TAIL, not the whole clip. Feeding a long clip (e.g. 100 s) into the
+buffer and masking only a few seconds makes the 8-step distilled sampler
+collapse to near-silence in that tiny window — the bug that shipped. A bounded
+lead keeps the masked (generated) region substantial and healthy, and because
+stitch rejoins the tail onto the full pristine original, the listener still
+hears their entire clip before the seam.
 
 Mask convention (verified against the installed library source):
   inpaint_mask = ones(buffer); inpaint_mask[start:end] = 0
-  -> 1 = keep the input audio, 0 = generate. So masking [L_src, L_total] keeps
-  the source in [0, L_src] and generates everything after it.
+  -> 1 = keep the input audio, 0 = generate. We place `lead` seconds of source
+  at the front and mask [lead, lead+new], so SA3 keeps the lead and generates a
+  fresh `new`-second tail that continues from the clip's end.
 """
 import numpy as np
 import torch
@@ -30,6 +40,13 @@ DEFAULT_CFG = 1.0        # distilled-model guidance; the prompt still conditions
                          # at 1.0 (CFG amplification off, conditional path on)
 MAX_TOTAL_SECONDS = 120  # SA3 Small duration cap (sample_size / sample_rate)
 MIN_NEW_SECONDS = 5      # below this a "continuation" isn't worth a GPU call
+MAX_LEAD_SECONDS = 30    # how much of the clip's TAIL to feed SA3 as run-up.
+                         # SA3 generates a healthy continuation from a bounded
+                         # lead-in; keeping a very long source in the buffer and
+                         # masking only a few seconds makes the distilled sampler
+                         # produce near-silence. 30 s is inside the model's
+                         # healthy range (the verified lab takes used ~30 s leads)
+                         # and the splice restores the full clip anyway.
 
 _model = None
 _model_config = None
@@ -95,31 +112,36 @@ MAX_SOURCE_SECONDS = MAX_TOTAL_SECONDS - MIN_NEW_SECONDS
 
 
 def plan_continuation(source_seconds, total_seconds):
-    """Pure helper (unit-testable, no model): clamp the request to SA3's limits
-    and return (total_seconds, new_seconds, mask_start, mask_end).
-
-    - the mask runs from where the source ends to the total length: that masked
-      region is what SA3 generates, so mask_end MUST exceed mask_start.
-    - total is capped at MAX_TOTAL_SECONDS and floored so at least
-      MIN_NEW_SECONDS of new audio is generated.
-    - raises ValueError when the source is already so long there's no room to
-      continue under the cap (otherwise the mask would invert and SA3 would
-      silently generate nothing).
+    """Pure helper (unit-testable, no model): turn a (source, requested-total)
+    pair into the SA3 generation buffer and return (lead, new_seconds, buffer).
+
+    - `lead`   : seconds of the clip's TAIL used as run-up context, capped at
+                 MAX_LEAD_SECONDS so a long clip can't drown the masked region.
+    - `new`    : seconds of fresh audio to generate. We extend to the requested
+                 finished length (`total - source`), floored at MIN_NEW_SECONDS
+                 so every call earns its GPU time, and bounded so the buffer
+                 (lead + new) never exceeds SA3's MAX_TOTAL_SECONDS cap.
+    - `buffer` : lead + new, i.e. the full generation buffer. The mask runs
+                 [lead, buffer]; buffer > lead always, so it never inverts.
+
+    Raises ValueError only for a clip longer than MAX_SOURCE_SECONDS — at that
+    point it's a full track, not an unfinished clip to continue.
     """
     source_seconds = float(source_seconds)
     total_seconds = float(total_seconds)
     if source_seconds > MAX_SOURCE_SECONDS:
         raise ValueError(
-            f"clip is {source_seconds:.0f}s — too long to continue under SA3's "
-            f"{MAX_TOTAL_SECONDS:.0f}s cap (need room for at least "
-            f"{MIN_NEW_SECONDS:.0f}s of new audio); trim it under "
-            f"{MAX_SOURCE_SECONDS:.0f}s.")
-    # source <= MAX_SOURCE_SECONDS, so source + MIN_NEW <= MAX_TOTAL: the floor
-    # never pushes total past the cap, and mask_end (total) > mask_start (source).
+            f"clip is {source_seconds:.0f}s — that's a finished-length track, "
+            f"not an unfinished clip. CODA continues clips up to "
+            f"{MAX_SOURCE_SECONDS:.0f}s; trim it shorter and re-upload.")
     total_seconds = min(total_seconds, MAX_TOTAL_SECONDS)
-    total_seconds = max(total_seconds, source_seconds + MIN_NEW_SECONDS)
-    new_seconds = total_seconds - source_seconds
-    return total_seconds, new_seconds, source_seconds, total_seconds
+    lead = min(source_seconds, MAX_LEAD_SECONDS)
+    # extend to the requested finished length; floor at MIN_NEW, and never let
+    # lead + new exceed the buffer cap.
+    new_seconds = max(total_seconds - source_seconds, MIN_NEW_SECONDS)
+    new_seconds = min(new_seconds, MAX_TOTAL_SECONDS - lead)
+    buffer_seconds = lead + new_seconds
+    return lead, new_seconds, buffer_seconds
 
 
 def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG,
@@ -159,19 +181,26 @@ def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG,
     _notify("reading")
     source = _load_source(clip_path)
     source_seconds = source.shape[-1] / SR
+
+    lead, new_seconds, buffer_seconds = plan_continuation(
+        source_seconds, total_seconds)
+
+    # condition on only the TAIL `lead` seconds of the clip. This is the bug fix:
+    # a long source no longer fills the buffer and starves the masked region.
+    lead_samples = min(int(round(lead * SR)), source.shape[-1])
+    lead_audio = source[:, -lead_samples:]
     # the autoencoder runs in the model's dtype (fp16 on CUDA); the conditioning
     # audio must match it or the encoder's conv1d rejects the input dtype.
     model_dtype = next(model.model.parameters()).dtype
-    source = source.to(model_dtype)
+    lead_audio = lead_audio.to(model_dtype)
 
-    total_seconds, new_seconds, mask_start, mask_end = plan_continuation(
-        source_seconds, total_seconds)
+    mask_start, mask_end = lead, buffer_seconds
     prompt = (prompt or "").strip()
 
-    print(f"[coda] continuation: source={source_seconds:.1f}s -> "
-          f"total={total_seconds:.1f}s (+{new_seconds:.1f}s new), "
-          f"mask=[{mask_start:.1f}s, {mask_end:.1f}s], steps={STEPS}, "
-          f"cfg={cfg_scale}, prompt={prompt!r}", flush=True)
+    print(f"[coda] continuation: source={source_seconds:.1f}s, "
+          f"lead={lead:.1f}s -> buffer={buffer_seconds:.1f}s "
+          f"(+{new_seconds:.1f}s new), mask=[{mask_start:.1f}s, {mask_end:.1f}s], "
+          f"steps={STEPS}, cfg={cfg_scale}, prompt={prompt!r}", flush=True)
 
     _notify("composing")
     with torch.no_grad():
@@ -179,10 +208,10 @@ def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG,
             model,
             steps=STEPS,
             cfg_scale=cfg_scale,
-            conditioning=[{"prompt": prompt, "seconds_total": total_seconds}],
+            conditioning=[{"prompt": prompt, "seconds_total": buffer_seconds}],
             sample_size=_sample_size,
             sampler_type=SAMPLER,
-            inpaint_audio=(SR, source),
+            inpaint_audio=(SR, lead_audio),
             inpaint_mask_start_seconds=mask_start,
             inpaint_mask_end_seconds=mask_end,
             seed=seed,
@@ -190,21 +219,24 @@ def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG,
         )
 
     _notify("finalizing")
-    # (b, d, n) -> (d, b*n); peak-normalize like Stability's reference Space
+    # (b, d, n) -> (d, b*n)
     output = rearrange(output, "b d n -> d (b n)")
     audio = output.to(torch.float32).cpu().numpy()
-    peak = float(np.abs(audio).max())
-    if peak > 1e-9:
-        audio = audio / peak
 
     if audio.shape[0] == 1:            # safety: ensure stereo
         audio = np.repeat(audio, 2, axis=0)
 
-    boundary = int(round(source_seconds * SR))
-    end = int(round(total_seconds * SR))
-    end = min(end, audio.shape[-1])
-    new_tail = audio[:, boundary:end]
-    new_tail = np.ascontiguousarray(new_tail.astype(np.float32))
+    # the generated region is [lead, buffer]; slice it out first, THEN normalize
+    # by the tail's OWN peak. Normalizing the whole buffer (as before) let a loud
+    # lead transient divide the tail down toward silence; per-tail normalization
+    # returns the continuation at a healthy standalone level and stitch re-levels
+    # it to the seam.
+    start = int(round(lead * SR))
+    end = min(int(round(buffer_seconds * SR)), audio.shape[-1])
+    new_tail = np.ascontiguousarray(audio[:, start:end].astype(np.float32))
+    peak = float(np.abs(new_tail).max())
+    if peak > 1e-9:
+        new_tail = new_tail / peak
 
     print(f"[coda] generated tail: shape={new_tail.shape} "
           f"({new_tail.shape[-1] / SR:.1f}s), peak after norm "

test_engine_logic.py

@@ -1,69 +1,89 @@
-"""Unit tests for engine.py's pure continuation planning — the source-length ->
-mask-bounds + total-length-cap math. No model, no GPU: `plan_continuation` is a
-pure function, so these run anywhere torch+numpy import."""
+"""Unit tests for engine.py's pure continuation planning — the bounded-lead +
+generation-buffer math. No model, no GPU: `plan_continuation` is a pure
+function, so these run anywhere torch+numpy import.
+
+Contract (post deployed-bug fix): plan_continuation(source, total) returns
+(lead, new_seconds, buffer_seconds) where
+  - lead   = min(source, MAX_LEAD_SECONDS)   — tail run-up fed to SA3
+  - new    = extend to `total`, floored at MIN_NEW, bounded so lead+new <= cap
+  - buffer = lead + new                       — the masked region is [lead, buffer]
+"""
 import pytest
 
 import engine
 
 
-def test_normal_request_maps_to_tail_mask():
-    total, new, mstart, mend = engine.plan_continuation(30, 60)
-    assert total == 60
-    assert new == 30
-    # the mask runs from where the source ends to the total length
-    assert mstart == 30
-    assert mend == 60
+def test_short_source_keeps_whole_clip_as_lead():
+    # source shorter than the lead cap: keep the whole clip, extend to total
+    lead, new, buf = engine.plan_continuation(30, 60)
+    assert lead == 30                       # whole 30s clip is the lead-in
+    assert new == 30                        # generate up to the 60s finished length
+    assert buf == 60                        # buffer = lead + new
 
 
-def test_total_capped_at_120():
-    total, new, mstart, mend = engine.plan_continuation(30, 200)
-    assert total == engine.MAX_TOTAL_SECONDS == 120
-    assert new == 90
-    assert mstart == 30 and mend == 120
+def test_buffer_capped_at_120():
+    lead, new, buf = engine.plan_continuation(30, 200)
+    assert lead == 30
+    assert buf == engine.MAX_TOTAL_SECONDS == 120
+    assert new == 90                        # lead(30) + new(90) == 120 cap
 
 
 def test_min_new_floor_enforced():
     # asking for barely-longer-than-source still generates at least MIN_NEW
-    total, new, mstart, mend = engine.plan_continuation(30, 31)
+    lead, new, buf = engine.plan_continuation(30, 31)
     assert new == engine.MIN_NEW_SECONDS == 5
-    assert total == 35
-    assert mstart == 30 and mend == 35
+    assert lead == 30 and buf == 35
 
 
-def test_mask_always_brackets_the_new_region():
-    for src, req in [(15, 40), (29.5, 60), (50, 90), (10, 120)]:
-        total, new, mstart, mend = engine.plan_continuation(src, req)
-        assert mstart == src                       # mask starts at the seam
-        assert mend == total                       # …and runs to the end
-        assert abs((mend - mstart) - new) < 1e-6   # masked span == new audio
-        assert total <= engine.MAX_TOTAL_SECONDS
+def test_long_source_clamps_lead_and_still_extends():
+    # the deployed bug: a 100s source kept the whole clip and masked ~5s, which
+    # the sampler rendered as near-silence. now the lead is clamped and the
+    # request still extends toward the asked-for total.
+    lead, new, buf = engine.plan_continuation(100, 120)
+    assert lead == engine.MAX_LEAD_SECONDS == 30   # not 100s of kept context
+    assert new == 20                                # 120 - 100 of fresh audio
+    assert buf == 50                                # bounded, healthy buffer
+    assert buf - lead == new                        # masked span == new audio
 
 
-def test_source_at_max_keeps_min_new():
-    # the longest allowed clip still gets exactly MIN_NEW of continuation
-    total, new, mstart, mend = engine.plan_continuation(engine.MAX_SOURCE_SECONDS, 130)
-    assert total == engine.MAX_TOTAL_SECONDS == 120
+def test_request_below_source_still_makes_real_audio():
+    # total < source is contradictory (can't shorten); we still generate a
+    # MIN_NEW tail rather than nothing, and the lead stays bounded.
+    lead, new, buf = engine.plan_continuation(100, 60)
+    assert lead == 30
     assert new == engine.MIN_NEW_SECONDS == 5
-    assert mend > mstart                              # mask never inverts
+    assert buf == 35
+    assert buf > lead                               # mask never inverts
 
 
-def test_overlong_source_raises_not_inverts():
-    # the old bug: source >= cap produced an inverted mask and a silent
-    # no-op continuation. now it must raise instead.
+def test_mask_brackets_the_new_region_across_inputs():
+    for src, req in [(15, 40), (29.5, 60), (50, 90), (10, 120), (100, 110)]:
+        lead, new, buf = engine.plan_continuation(src, req)
+        assert lead == min(src, engine.MAX_LEAD_SECONDS)
+        assert abs((buf - lead) - new) < 1e-6       # masked span == new audio
+        assert buf > lead                           # never inverts
+        assert new >= engine.MIN_NEW_SECONDS - 1e-6
+        assert buf <= engine.MAX_TOTAL_SECONDS
+
+
+def test_normal_request_reaches_requested_total():
+    # for a legitimate request (total > source) the finished length is exactly
+    # source + new == total: it extends to what the user asked for.
+    for src, total in [(15, 40), (29.5, 60), (20, 90)]:
+        _, new, _ = engine.plan_continuation(src, total)
+        assert abs((src + new) - total) < 1e-6
+
+
+def test_overlong_source_raises():
+    # a clip longer than MAX_SOURCE_SECONDS is a finished track, not a clip
     for src in (116, 120, 125, 200):
         with pytest.raises(ValueError):
             engine.plan_continuation(src, 60)
 
 
-def test_no_plan_ever_inverts_the_mask():
-    for src in [1, 15, 29.5, 60, 90, 110, 115]:
-        total, new, mstart, mend = engine.plan_continuation(src, 60)
-        assert mend > mstart        # mask_end strictly after mask_start
-        assert new >= engine.MIN_NEW_SECONDS - 1e-6
-
-
 def test_constants_match_sa3_contract():
     assert engine.SR == 44100
     assert engine.STEPS == 8
     assert engine.SAMPLER == "pingpong"
     assert engine.MAX_TOTAL_SECONDS == 120
+    assert engine.MAX_LEAD_SECONDS == 30