Harden best-of-N selection: reject dropouts and squashed takes

The selection metric only scored loud-bursts and whole-tail silence, so two
failure modes that sound bad to a listener slipped through as clean:
- a brief mid-tail near-silent dropout (overall RMS stays high and max/median
stays low, so neither existing term fires) - an audible cut-out;
- dynamics/transient collapse - a perfectly tonal but squashed, attack-less
wash that low-flatness and loudness checks all green-light.
Add a sustained-dropout term and a crest-factor (peak/RMS) floor to the score,
raise the candidate pool to 5 and tighten early-accept to 3.5 so best-of-N keeps
searching past a merely-okay draw for a genuinely good one. Generation is
unchanged; only which draw is selected. Adds unit tests for all four failure
modes plus the natural-ending-taper false-positive guard.

engine.py

@@ -88,11 +88,21 @@ DEFAULT_CFG = 1.0        # distilled-model guidance; the prompt still conditions
 
 # Best-of-N: with a random seed each draw differs, so we generate a few and keep
 # the cleanest. Bounded so it never blows the ZeroGPU window.
-DEFAULT_CANDIDATES = 3       # how many draws to consider when no seed is pinned
+DEFAULT_CANDIDATES = 5       # how many draws to consider when no seed is pinned.
+                             # Raised from 3: the quality bar is the BEST draw, not
+                             # the first clean-ish one, and a fast H200 draw is
+                             # cheap. Early-accept + the GPU budget still short-
+                             # circuit when an early draw is already great.
 GPU_BUDGET_SECONDS = 85.0    # stop drawing once this much wall-clock is spent
                              # (the @spaces.GPU window is 120s; leave slack)
-EARLY_ACCEPT_SCORE = 4.0     # a draw this clean is taken immediately, no re-draw
-                             # (heuristic — see _tail_artifact_score; tune live)
+EARLY_ACCEPT_SCORE = 3.5     # a draw this clean is taken immediately, no re-draw.
+                             # Tightened from 4.0 so a merely-okay draw doesn't
+                             # short-circuit the search for a genuinely good one.
+DROPOUT_FLOOR = 0.12         # quietest sustained 0.2 s below this fraction of the
+                             # tail median counts as a mid-tail hole (re-draw it)
+CREST_FLOOR = 4.0            # peak/RMS below this is a squashed, transient-less
+                             # wash (real music here is ~6-8); penalize it
+CREST_SCALE = 1.5            # how hard a collapsed crest is penalized
 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.
@@ -215,18 +225,31 @@ def _tail_artifact_score(tail, sr=SR):
     """Lower is better. A blind, ear-free quality score for a generated tail,
     used to pick the cleanest of several candidate draws.
 
-    It targets the two ways an SA3 draw goes bad:
+    It targets the four ways an SA3 draw goes bad:
       * "sporadic loud random synth noises" — even a FEW short windows far louder
         than the body push the loudest window way above the median. (After the
         whole-buffer peak-normalize, a burst that set the peak crushes the body,
         making the gap larger still.) Sustained dynamics rarely make any single
         50 ms window many times the median, so musical loudness doesn't trip it.
-      * silence collapse — a near-silent tail (the other known failure) is caught
-        by the explicit loudness floor below.
-
-    Score = max(window RMS) / median(window RMS)  +  silence penalty.
+      * silence collapse — a near-silent WHOLE tail is caught by the loudness
+        floor below (it keys off the tail's overall RMS).
+      * mid-tail dropout — a brief near-silent HOLE inside an otherwise healthy
+        tail. This is the gap the first terms miss: overall RMS stays high (so the
+        silence floor never fires) and max/median stays low (so spikiness never
+        fires), yet a listener plainly hears the music cut out for a beat. We
+        detect it as a SUSTAINED quiet stretch — the quietest ~0.2 s envelope
+        falling well below the median.
+      * dynamics/transient collapse — a draw can be perfectly tonal and steady yet
+        sound DULL and lifeless: its transients are smeared, so there's no attack,
+        just a wall of mush. Flatness and loudness checks all read "clean". It
+        shows up as a collapsed crest factor (peak/RMS): real music here sits at
+        crest ~6-8, a squashed draw falls to ~2-3. We penalize a low crest so
+        best-of-N prefers the punchy draw over the mushy one.
+
+    Score = max/median + silence + dropout + dynamics penalties.
     Computed on a mono mix over short (~50 ms) windows. A flat, steady signal
-    scores ~1; isolated loud bursts or a crushed body score high.
+    scores ~1; loud bursts, a crushed body, a mid-tail hole, or a smeared,
+    transient-less wash all score high.
     """
     mono = tail.mean(axis=0) if tail.ndim == 2 else np.asarray(tail)
     mono = np.asarray(mono, dtype=np.float64)
@@ -241,7 +264,31 @@ def _tail_artifact_score(tail, sr=SR):
     spikiness = loudest / median
     overall = float(np.sqrt(np.mean(mono ** 2)) + 1e-12)
     silence_penalty = 0.0 if overall > 0.02 else (0.02 - overall) * 200.0
-    return spikiness + silence_penalty
+
+    # mid-tail dropout: smooth the window energies over ~0.2 s and find how far
+    # the quietest SUSTAINED stretch falls below the median. Exclude the final
+    # 0.5 s so a natural ending taper (which stitch fades anyway) isn't punished.
+    # A clean tail's quietest 0.2 s sits ~0.15-0.4x the median -> no penalty; a
+    # real hole drops to <0.1x -> a penalty large enough to lose the early-accept
+    # and force another draw, so best-of-N rolls past the glitch.
+    dropout_penalty = 0.0
+    smooth = np.convolve(energies, np.ones(4) / 4.0, mode="valid")
+    guard = int(0.5 / 0.05)                 # last 0.5 s of windows
+    body = smooth[:-guard] if smooth.size > guard + 4 else smooth
+    if body.size:
+        dropout = float(np.min(body)) / median
+        dropout_penalty = min(8.0, max(0.0, DROPOUT_FLOOR / max(dropout, 1e-3)
+                                       - 1.0) * 2.0)
+
+    # dynamics/transient collapse: crest = peak / RMS. The tail is peak-normalized
+    # to ~1.0, so this is essentially 1/RMS — a squashed, attack-less wash reads
+    # high RMS (low crest); a punchy, dynamic take reads low RMS (high crest).
+    # Penalize only a clearly collapsed crest, so we never punish a naturally
+    # dynamic draw.
+    peak = float(np.abs(mono).max())
+    crest = peak / overall
+    crest_penalty = max(0.0, CREST_FLOOR - crest) * CREST_SCALE
+    return spikiness + silence_penalty + dropout_penalty + crest_penalty
 
 
 def continue_audio(clip_path, total_seconds, prompt="", cfg_scale=DEFAULT_CFG,

test_engine_score.py

@@ -0,0 +1,102 @@
+"""Unit tests for engine._tail_artifact_score — the best-of-N selection metric.
+
+The score is "lower is better"; a draw at or below EARLY_ACCEPT_SCORE is taken
+immediately, anything above forces another draw. These tests pin the four
+failure modes the metric must catch (loud burst, whole-tail silence, mid-tail
+dropout, dynamics/transient collapse) and confirm a clean, dynamic take passes.
+
+Synthetic signals only — no model, no GPU. All are peak-normalized to ~1.0, the
+state the score actually sees (engine peak-normalizes each draw before scoring).
+"""
+import numpy as np
+import types
+import sys
+
+sys.modules.setdefault("stable_audio_tools", types.ModuleType("stable_audio_tools"))
+import engine  # noqa: E402
+
+SR = engine.SR
+DUR = 10.0
+
+
+def _norm(x):
+    p = float(np.abs(x).max())
+    return x / p if p > 0 else x
+
+
+def _stereo(x):
+    return np.stack([x, x]).astype(np.float32)
+
+
+def _t():
+    return np.arange(int(DUR * SR)) / SR
+
+
+def _clean_dynamic():
+    """A continuous tonal bed plus short transients: peak set by the transients,
+    body level steady (no holes), so crest is high (~real music) and spikiness
+    stays moderate — the take the metric should ACCEPT."""
+    t = _t()
+    bed = 0.12 * (np.sin(2 * np.pi * 220 * t)
+                  + 0.6 * np.sin(2 * np.pi * 330 * t)
+                  + 0.4 * np.sin(2 * np.pi * 440 * t))
+    sig = bed.copy()
+    w = int(0.005 * SR)                       # 5 ms transients every 0.5 s
+    for c in range(int(0.4 * SR), len(sig), int(0.5 * SR)):
+        env = np.hanning(2 * w)[:w]
+        sig[c:c + w] += 0.7 * env * np.sin(2 * np.pi * 660 * t[c:c + w])
+    return _norm(sig)
+
+
+def test_clean_dynamic_take_is_accepted():
+    score = engine._tail_artifact_score(_stereo(_clean_dynamic()), SR)
+    assert score <= engine.EARLY_ACCEPT_SCORE, score
+
+
+def test_loud_burst_is_rejected():
+    sig = _clean_dynamic()
+    c = int(5 * SR)
+    sig[c:c + int(0.05 * SR)] *= 20          # a single loud spike
+    score = engine._tail_artifact_score(_stereo(_norm(sig)), SR)
+    assert score > engine.EARLY_ACCEPT_SCORE, score
+
+
+def test_whole_tail_silence_is_rejected():
+    sig = _clean_dynamic() * 0.004           # below the 0.02 RMS floor
+    score = engine._tail_artifact_score(_stereo(sig), SR)
+    assert score > engine.EARLY_ACCEPT_SCORE, score
+
+
+def test_mid_tail_dropout_is_rejected():
+    """A clean take with a 0.5 s near-silent hole in the middle — healthy overall
+    RMS and low spikiness, so only the dropout term can catch it."""
+    sig = _clean_dynamic()
+    sig[int(5.0 * SR):int(5.5 * SR)] *= 0.01
+    score = engine._tail_artifact_score(_stereo(sig), SR)
+    assert score > engine.EARLY_ACCEPT_SCORE, score
+
+
+def test_squashed_transientless_take_is_rejected():
+    """Dense, near-constant amplitude (crest collapses): tonal and steady, so
+    every other term reads clean — only the crest term flags the mush."""
+    t = _t()
+    sig = np.tanh(5 * (np.sin(2 * np.pi * 220 * t)
+                       + 0.8 * np.sin(2 * np.pi * 331 * t)
+                       + 0.7 * np.sin(2 * np.pi * 440 * t)))
+    score = engine._tail_artifact_score(_stereo(_norm(sig)), SR)
+    assert score > engine.EARLY_ACCEPT_SCORE, score
+
+
+def test_natural_ending_taper_is_not_a_dropout():
+    """A clean take that simply fades over its final ~0.6 s must NOT be read as a
+    dropout (stitch fades the end anyway); the back-guard protects it."""
+    sig = _clean_dynamic()
+    tail = sig[-int(0.6 * SR):]
+    sig[-int(0.6 * SR):] = tail * np.linspace(1.0, 0.0, len(tail))
+    score = engine._tail_artifact_score(_stereo(sig), SR)
+    assert score <= engine.EARLY_ACCEPT_SCORE, score
+
+
+def test_too_short_tail_is_avoided():
+    score = engine._tail_artifact_score(_stereo(np.zeros(int(0.01 * SR))), SR)
+    assert score == float("inf")