Precise beat grid: global (period,phase) fit (app.py)

app.py

@@ -15,9 +15,10 @@ import numpy as np
 import torch
 
 from softchart.generate import generate_song, load_hf
+from softchart.grid import debias_to_grid, fit_grid
 from softchart.hf import SoftChartPlanner
 from softchart.rhythm import snap_chart
-from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR, WINDOW
+from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR
 
 GEN_REPO = os.environ.get("SC_GEN", "JacobLinCool/softchart-generator")
 BEAT_REPO = os.environ.get("SC_BEAT", "JacobLinCool/softchart-beat")
@@ -61,27 +62,6 @@ def load_logmel(path):
     return mel, wav
 
 
-def predict_downbeats(beat_model, mel):
-    from scipy.signal import find_peaks
-    from softchart.generate import _autocast
-
-    T = mel.shape[1]
-    L = WINDOW // 4
-    acc = np.zeros(T // 4 + L)
-    cnt = np.zeros(T // 4 + L)
-    for st in range(0, max(T - 1, 1), WINDOW):
-        w = torch.from_numpy(mel[:, st:st + WINDOW].astype(np.float32))
-        if w.shape[1] < WINDOW:
-            w = torch.nn.functional.pad(w, (0, WINDOW - w.shape[1]), value=float(np.log(1e-5)))
-        with torch.no_grad(), _autocast(DEVICE):
-            mem = beat_model.encode(w[None].to(DEVICE))
-            pr = torch.sigmoid(beat_model.beat(mem[:, -L:]).float())[0, :, 1].cpu().numpy()
-        acc[st // 4: st // 4 + L] += pr
-        cnt[st // 4: st // 4 + L] += 1
-    acc /= np.maximum(cnt, 1)
-    return find_peaks(acc, height=0.4, distance=int(0.8 * FPS / 4))[0] * 4 / FPS
-
-
 def auto_plan(mel, bpm, downbeats=None):
     T = mel.shape[1]
     dur = T / FPS
@@ -154,20 +134,37 @@ def group_quantize(times, phase, grid, min_run=3):
     return slots
 
 
-def write_tja(gen, bpm, title, course, level, downbeats=None):
+def write_tja(gen, bpm, title, course, level, downbeats=None, grid_fit=None):
     hits = sorted((h["t"], CHAR[h["type"]]) for h in gen["hits"])
-    times = np.array([t for t, _ in hits]) if hits else np.array([0.0])
     beat = 60.0 / bpm
     grid = beat / (SUB / 4)
-    cands = np.arange(0, beat, grid / 4)
-    phase = float(cands[int(np.argmin([np.mean(np.abs(((times - o) / grid) - np.round((times - o) / grid))) for o in cands]))])
+    bias = 0.0
+    if grid_fit is not None and hits:
+        # authoritative fitted grid: barlines ARE the fitted downbeats.
+        # De-bias the generator's systematic latency (global shift only),
+        # then anchor slot 0 on the last fitted barline at/before the first note.
+        times, bias = debias_to_grid([t for t, _ in hits], grid_fit["phase"], grid)
+        phase = grid_fit["phase"] + float(np.floor((times[0] - grid_fit["phase"]) / (4 * beat))) * 4 * beat
+        q_times = list(times)
+    else:
+        times = np.array([t for t, _ in hits]) if hits else np.array([0.0])
+        cands = np.arange(0, beat, grid / 4)
+        phase = float(cands[int(np.argmin([np.mean(np.abs(((times - o) / grid) - np.round((times - o) / grid))) for o in cands]))])
+        q_times = [t for t, _ in hits]
+    slot_idx = group_quantize(q_times, phase, grid)
+    if slot_idx and min(slot_idx) < 0:
+        # note quantized just before the anchor barline: pull back whole bars
+        # so nothing is dropped (barline alignment is preserved mod SUB)
+        nb = int(np.ceil(-min(slot_idx) / SUB))
+        slot_idx = [s + nb * SUB for s in slot_idx]
+        phase -= nb * SUB * grid
     slots = {}
-    for idx, (t, ch) in zip(group_quantize([t for t, _ in hits], phase, grid), hits):
+    for idx, (t, ch) in zip(slot_idx, hits):
         if idx >= 0 and idx not in slots:
             slots[idx] = ch
     for sp in gen["spans"]:
-        i0 = int(round((sp["t0"] - phase) / grid))
-        i1 = int(round((sp["t1"] - phase) / grid))
+        i0 = int(round((sp["t0"] - bias - phase) / grid))
+        i1 = int(round((sp["t1"] - bias - phase) / grid))
         while i0 in slots:
             i0 += 1
         while i1 in slots or i1 <= i0:
@@ -175,7 +172,9 @@ def write_tja(gen, bpm, title, course, level, downbeats=None):
         if i0 >= 0:
             slots[i0] = CHAR[sp["type"]]
             slots[i1] = "8"
-    if slots:
+    if slots and grid_fit is None:
+        # legacy anchoring (no trusted grid): shift so the first note sits on a
+        # detected downbeat if one is nearby, else on the first barline
         first_t = min(slots) * grid + phase
         anchor_t = None
         if downbeats is not None and len(downbeats):
@@ -237,17 +236,34 @@ def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_pla
         raise gr.Error("Please upload an audio file.")
     M = get_models()
     mel, wav = load_logmel(audio)
-    dbs = None
+    grid = dbs = None
     if use_beat and M["beat"] is not None:
-        dbs = predict_downbeats(M["beat"], mel)
+        # global robust (period, phase) fit over the whole song — much more
+        # precise than per-peak use (each raw peak carries ~±23 ms bin noise)
+        grid = fit_grid(M["beat"], mel, device=DEVICE)
+        if grid is not None:
+            # rigid synthesized barlines when the fit is trustworthy; raw peaks
+            # (plan-block edges only, no anchoring) when it is not
+            dbs = grid["downbeats"] if grid["ok"] else grid["db_peaks"]
+            if not grid["ok"]:
+                grid = None
     if bpm_override and bpm_override > 0:
         bpm = float(bpm_override)
+        if grid is not None and abs(grid["bpm"] - bpm) > 0.5:
+            grid = None  # user disagrees with the fit: don't anchor to it
+    elif grid is not None:
+        bpm = grid["bpm"]  # already integer-snapped when the residual allows
     elif dbs is not None and len(dbs) > 4:
-        bpm = 60.0 / float(np.median(np.diff(dbs)))
+        period = float(np.median(np.diff(dbs)))  # downbeat gap = one 4/4 bar
+        bpm = 240.0 / period if period > 0 else 0.0
+        while bpm >= 210:  # octave guard
+            bpm /= 2.0
+        while 0 < bpm < 70:
+            bpm *= 2.0
     else:
         import librosa
         bpm = float(np.atleast_1d(librosa.beat.beat_track(y=wav, sr=SR)[0])[0])
-    if abs(bpm - round(bpm)) < 0.06:
+    if grid is None and abs(bpm - round(bpm)) < 0.06:
         bpm = float(round(bpm))
 
     plan = None
@@ -262,12 +278,18 @@ def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_pla
                       seed=0, device=DEVICE, plan=plan)
     g = snap_chart(g, bpm)
     title = os.path.splitext(os.path.basename(audio))[0]
-    tja = write_tja(g, bpm, title, course, int(level), dbs)
+    tja = write_tja(g, bpm, title, course, int(level), dbs, grid_fit=grid)
     tja_path = tempfile.mktemp(suffix=f"_{course}.tja")
     with open(tja_path, "w") as f:
         f.write(tja)
     img = render(mel, g, title, course)
-    info = (f"BPM {bpm:.1f} · {len(g['hits'])} notes · {len(g['spans'])} spans"
+    if grid is not None:
+        grid_info = f" · grid: rms {grid['rms_ms']:.1f}ms ({grid['inlier_frac']:.0%} inlier)"
+    elif use_beat and M["beat"] is not None:
+        grid_info = " · grid: unreliable, barline anchoring off"
+    else:
+        grid_info = ""
+    info = (f"BPM {bpm:.1f} · {len(g['hits'])} notes · {len(g['spans'])} spans" + grid_info
             + (f" · plan: {sum(1 for b in plan if b[3]==1)} gaps, {sum(1 for b in plan if b[3]==2)} climax" if plan else ""))
     return img, tja_path, info