Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 JacobLinCool
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,33 @@
 ---
-title: Softchart
-emoji: 📚
-colorFrom: blue
-colorTo: red
+title: SoftChart
+emoji: 🥁
+colorFrom: red
+colorTo: yellow
 sdk: gradio
-sdk_version: 6.19.0
-python_version: '3.13'
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
+license: mit
+models:
+  - JacobLinCool/softchart-generator
+  - JacobLinCool/softchart-beat
+  - JacobLinCool/softchart-planner
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🥁 SoftChart
+
+Generate a **Taiko no Tatsujin** chart from any song. An 8M-parameter,
+plan-conditioned encoder–decoder model turns audio into a playable `.tja`,
+with an optional learned planner (song-level structure) and beat-anchoring
+model (barline alignment). All models are MIT-licensed and loaded from the
+Hugging Face Hub via `from_pretrained`.
+
+- **Generator** ([softchart-generator](https://huggingface.co/JacobLinCool/softchart-generator)) — audio → chart events, with difficulty / density / plan control
+- **Planner** ([softchart-planner](https://huggingface.co/JacobLinCool/softchart-planner)) — audio → song-level plan (density envelope, breathing gaps, climax)
+- **Beat** ([softchart-beat](https://huggingface.co/JacobLinCool/softchart-beat)) — beat/downbeat detection for barline anchoring
+
+Upload a song, pick a difficulty, and get a preview plus a `.tja` playable in
+TJAPlayer3 / OpenTaiko.
+
+> Models were trained on charts derived from a commercial game; released for
+> research and personal use. Code and weights are MIT-licensed.

app.py

@@ -0,0 +1,304 @@
+"""SoftChart — Gradio demo for Hugging Face Spaces.
+
+Generate a Taiko no Tatsujin chart from any audio file, using the full system:
+  - SoftChartGenerator  (main model, plan-conditioned)
+  - SoftChartPlanner     (auto song-level planning)
+  - SoftChartBeat        (beat/downbeat for barline anchoring)
+All models load from the Hub via from_pretrained. MIT licensed.
+"""
+
+import os
+import tempfile
+
+import gradio as gr
+import numpy as np
+import torch
+
+from softchart.generate import generate_song, load_hf
+from softchart.hf import SoftChartPlanner
+from softchart.rhythm import snap_chart
+from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR, WINDOW
+
+GEN_REPO = os.environ.get("SC_GEN", "JacobLinCool/softchart-generator")
+BEAT_REPO = os.environ.get("SC_BEAT", "JacobLinCool/softchart-beat")
+PLAN_REPO = os.environ.get("SC_PLAN", "JacobLinCool/softchart-planner")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+COURSE_DENS = {"easy": 1, "normal": 2, "hard": 4, "oni": 7}
+DEFAULT_LEVEL = {"easy": 3, "normal": 5, "hard": 7, "oni": 9}
+CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4",
+        "roll": "5", "roll_big": "6", "balloon": "7"}
+STYLE = {"don": ("#e8453c", 55), "ka": ("#4aa8c0", 55),
+         "don_big": ("#e8453c", 150), "ka_big": ("#4aa8c0", 150),
+         "roll": ("#e8b23a", 80), "roll_big": ("#e8b23a", 150), "balloon": ("#d47fda", 100)}
+SUB = 96
+
+_MODELS = {}
+
+
+def get_models():
+    if not _MODELS:
+        _MODELS["gen"] = load_hf(GEN_REPO, device=DEVICE)
+        try:
+            _MODELS["beat"] = load_hf(BEAT_REPO, device=DEVICE)
+        except Exception:
+            _MODELS["beat"] = None
+        try:
+            _MODELS["plan"] = SoftChartPlanner.from_pretrained(PLAN_REPO).to(DEVICE).eval()
+        except Exception:
+            _MODELS["plan"] = None
+    return _MODELS
+
+
+def load_logmel(path):
+    import librosa
+
+    wav, _ = librosa.load(path, sr=SR, mono=True)
+    fb = librosa.filters.mel(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=20.0, fmax=SR / 2)
+    spec = torch.stft(torch.from_numpy(wav), N_FFT, hop_length=HOP,
+                      window=torch.hann_window(N_FFT), center=True, return_complex=True)
+    mel = np.log(fb @ spec.abs().pow(2).numpy() + 1e-5).astype(np.float32)
+    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
+    flux = np.concatenate([[0], np.maximum(0, np.diff(mel, axis=1)).sum(0)])
+    beat = 60.0 / bpm
+    edges = (list(downbeats[::4]) + [dur]) if (downbeats is not None and len(downbeats) >= 2) \
+        else list(np.arange(0, dur, 4 * beat)) + [dur]
+    vals = [float(flux[int(a * FPS):int(b * FPS)].mean()) if int(b * FPS) > int(a * FPS) else 0.0
+            for a, b in zip(edges, edges[1:])]
+    if not vals:
+        return None
+    vals = np.array(vals)
+    lo, hi = np.percentile(vals, 15), np.percentile(vals, 92)
+    peak = int(np.argmax(vals))
+    plan = []
+    for i, (a, b) in enumerate(zip(edges, edges[1:])):
+        frac = (vals[i] - lo) / max(hi - lo, 1e-6)
+        d8 = int(np.clip(round(frac * 7), 0, 7))
+        fl = 1 if (vals[i] <= lo and 0 < i < len(vals) - 1) else (2 if i == peak and vals[i] >= hi else 0)
+        plan.append([round(a, 3), round(b, 3), d8, fl])
+    return plan
+
+
+def learned_plan(planner, mel, course, bpm, downbeats=None):
+    dur = mel.shape[1] / FPS
+    beat = 60.0 / bpm
+    edges = (list(downbeats[::4]) + [dur]) if (downbeats is not None and len(downbeats) >= 2) \
+        else list(np.arange(0, dur, 4 * beat)) + [dur]
+    feats, spans = [], []
+    for a, b in zip(edges, edges[1:]):
+        seg = mel[:, int(a * FPS):int(b * FPS)]
+        if seg.shape[1] < 2:
+            continue
+        fx = np.maximum(0, np.diff(seg, axis=1)).sum(0)
+        feats.append(np.concatenate([seg.mean(1), seg.std(1), [fx.mean(), fx.std(), fx.max()]]))
+        spans.append((round(float(a), 3), round(float(b), 3)))
+    if not feats:
+        return None
+    cid = {"easy": 0, "normal": 1, "hard": 2, "oni": 3}[course]
+    x = torch.tensor(np.array(feats), dtype=torch.float32)[None].to(DEVICE)
+    with torch.no_grad():
+        pd, pf = planner(x, torch.tensor([cid], device=DEVICE))
+    d8 = pd[0].argmax(-1).cpu().numpy()
+    fl = pf[0].argmax(-1).cpu().numpy()
+    return [[a, b, int(d), int(f)] for (a, b), d, f in zip(spans, d8, fl)]
+
+
+def group_quantize(times, phase, grid, min_run=3):
+    n = len(times)
+    slots = [0] * n
+    i = 0
+    while i < n:
+        j = i
+        while j + 1 < n:
+            ioi = times[j + 1] - times[j]
+            ref = (times[j] - times[i]) / (j - i) if j > i else ioi
+            if 0.02 < ioi < 1.2 and abs(ioi - ref) < 0.22 * max(ref, 1e-6):
+                j += 1
+            else:
+                break
+        if j - i + 1 >= min_run:
+            k = max(1, int(round((times[j] - times[i]) / (j - i) / grid)))
+            anchor = int(round((times[i] - phase) / grid))
+            for m in range(j - i + 1):
+                slots[i + m] = anchor + m * k
+        else:
+            for m in range(i, j + 1):
+                slots[m] = int(round((times[m] - phase) / grid))
+        i = j + 1
+    return slots
+
+
+def write_tja(gen, bpm, title, course, level, downbeats=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]))])
+    slots = {}
+    for idx, (t, ch) in zip(group_quantize([t for t, _ in hits], phase, grid), 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))
+        while i0 in slots:
+            i0 += 1
+        while i1 in slots or i1 <= i0:
+            i1 += 1
+        if i0 >= 0:
+            slots[i0] = CHAR[sp["type"]]
+            slots[i1] = "8"
+    if slots:
+        first_t = min(slots) * grid + phase
+        anchor_t = None
+        if downbeats is not None and len(downbeats):
+            near = downbeats[downbeats <= first_t + 0.12]
+            if len(near) and first_t - near[-1] < 4 * beat:
+                anchor_t = near[-1]
+        shift = int(round((anchor_t - phase) / grid)) if anchor_t is not None else min(slots)
+        if shift:
+            slots = {k - shift: v for k, v in slots.items()}
+            phase += shift * grid
+    n_meas = (max(slots) // SUB + 1) if slots else 1
+    lines = ["".join(slots.get(m * SUB + k, "0") for k in range(SUB)) + "," for m in range(n_meas)]
+    balloons = [10] * sum(1 for s in gen["spans"] if s["type"] == "balloon")
+    return "\n".join([
+        f"TITLE:{title} (SoftChart)", f"BPM:{bpm:g}", "WAVE:song.ogg",
+        f"OFFSET:{-phase:.3f}", f"COURSE:{'Oni' if course == 'oni' else course.capitalize()}",
+        f"LEVEL:{level}", f"BALLOON:{','.join(map(str, balloons))}" if balloons else "BALLOON:",
+        "", "#START", *lines, "#END"]) + "\n"
+
+
+def render(mel, gen, title, course):
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    dur = mel.shape[1] / FPS
+    t0, t1 = dur * 0.33, min(dur * 0.33 + 26, dur)
+    fig = plt.figure(figsize=(13, 4))
+    gs = fig.add_gridspec(2, 1, height_ratios=[2.3, 1], hspace=0.1)
+    ax0 = fig.add_subplot(gs[0])
+    ax0.imshow(mel[:, int(t0 * FPS):int(t1 * FPS)], aspect="auto", origin="lower",
+               cmap="magma", extent=[t0, t1, 0, N_MELS])
+    ax0.set_xticks([])
+    ax0.set_ylabel("mel")
+    ax0.set_title(f"{title} — {course}  |  red=don blue=ka big=large yellow=roll purple=balloon")
+    ax1 = fig.add_subplot(gs[1])
+    for h in gen["hits"]:
+        if t0 <= h["t"] <= t1:
+            c, sz = STYLE.get(h["type"], ("#999", 40))
+            ax1.scatter(h["t"], 0, c=c, s=sz, edgecolors="k",
+                        linewidths=1.2 if "big" in h["type"] else 0.5, zorder=3)
+    for s in gen["spans"]:
+        if s["t1"] >= t0 and s["t0"] <= t1:
+            ax1.plot([max(s["t0"], t0), min(s["t1"], t1)], [0, 0],
+                     c=STYLE.get(s["type"], ("#e8b23a",))[0], lw=6, alpha=0.45, zorder=2)
+    ax1.set_xlim(t0, t1)
+    ax1.set_ylim(-1, 1)
+    ax1.set_yticks([])
+    ax1.set_xlabel("time (s)")
+    ax1.grid(axis="x", alpha=0.25)
+    out = tempfile.mktemp(suffix=".png")
+    fig.savefig(out, dpi=130, bbox_inches="tight")
+    plt.close(fig)
+    return out
+
+
+def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_planner, sampling):
+    if audio is None:
+        raise gr.Error("Please upload an audio file.")
+    M = get_models()
+    mel, wav = load_logmel(audio)
+    dbs = None
+    if use_beat and M["beat"] is not None:
+        dbs = predict_downbeats(M["beat"], mel)
+    if bpm_override and bpm_override > 0:
+        bpm = float(bpm_override)
+    elif dbs is not None and len(dbs) > 4:
+        bpm = 60.0 / float(np.median(np.diff(dbs)))
+    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:
+        bpm = float(round(bpm))
+
+    plan = None
+    if getattr(M["gen"], "_has_plan", False):
+        if use_planner and M["plan"] is not None:
+            plan = learned_plan(M["plan"], mel, course, bpm, dbs)
+        elif auto_plan_on:
+            plan = auto_plan(mel, bpm, dbs)
+
+    g = generate_song(M["gen"], mel, course, level=int(level),
+                      density_bucket=COURSE_DENS[course], greedy=not sampling,
+                      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_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"
+            + (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
+
+
+with gr.Blocks(title="SoftChart — AI Taiko chart generator") as demo:
+    gr.Markdown(
+        "# 🥁 SoftChart\n"
+        "Generate a **Taiko no Tatsujin** chart from any song. 8M-param plan-conditioned "
+        "model + learned planner + beat-anchoring — all MIT-licensed, loaded from the Hub.\n\n"
+        "Outputs a preview and a playable `.tja` (for TJAPlayer3 / OpenTaiko). "
+        "*Models trained on charts derived from a commercial game; for research/personal use.*"
+    )
+    with gr.Row():
+        with gr.Column():
+            audio = gr.Audio(type="filepath", label="Song (any format)")
+            course = gr.Dropdown(["easy", "normal", "hard", "oni"], value="oni", label="Difficulty (course)")
+            level = gr.Slider(1, 10, value=9, step=1, label="Level (1–10 stars)")
+            bpm = gr.Number(label="BPM (0 = auto-detect)", value=0)
+            with gr.Row():
+                auto_plan_on = gr.Checkbox(label="Auto-plan (heuristic structure)", value=True)
+                use_planner = gr.Checkbox(label="Learned planner", value=True)
+            with gr.Row():
+                use_beat = gr.Checkbox(label="Beat-anchor barlines", value=True)
+                sampling = gr.Checkbox(label="Sampling (diverse) vs greedy (best)", value=False)
+            btn = gr.Button("Generate chart", variant="primary")
+        with gr.Column():
+            out_img = gr.Image(label="Preview (mel + chart)")
+            out_info = gr.Textbox(label="Result", interactive=False)
+            out_tja = gr.File(label="Download .tja")
+    btn.click(generate, [audio, course, level, bpm, auto_plan_on, use_beat, use_planner, sampling],
+              [out_img, out_tja, out_info])
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,10 @@
+torch>=2.2
+numpy<2.1
+scipy
+librosa>=0.10
+numba>=0.60
+soundfile
+matplotlib
+gradio>=4.44
+huggingface_hub>=0.24
+safetensors

softchart/baseline.py

@@ -0,0 +1,43 @@
+"""Onset-detection baseline: spectral-flux peak picking on the cached log-mel
++ band-energy don/ka heuristic.
+
+This is the "chart generation is just onset detection" strawman the model must beat.
+Operates on the same cached mel features as the model, so no raw audio is needed.
+"""
+
+import numpy as np
+from scipy.signal import find_peaks
+
+from .vocab import FPS
+
+
+def baseline_chart(mel, target_nps):
+    """mel: (n_mels, T) log-mel (natural log). Returns list of (time_sec, class)."""
+    m = mel.astype(np.float32)
+    # spectral flux onset envelope
+    flux = np.maximum(0.0, np.diff(m, axis=1)).sum(axis=0)
+    flux = np.concatenate([[0.0], flux])
+    # smooth lightly
+    k = np.hanning(5)
+    k /= k.sum()
+    env = np.convolve(flux, k, mode="same")
+
+    dur = m.shape[1] / FPS
+    want = max(1, int(round(target_nps * dur)))
+    peaks, props = find_peaks(env, distance=int(0.09 * FPS), height=0.0)
+    if len(peaks) == 0:
+        return []
+    if len(peaks) > want:
+        strongest = np.argsort(props["peak_heights"])[-want:]
+        peaks = np.sort(peaks[strongest])
+
+    # don/ka: low vs high mel-band energy at the onset (linear power domain)
+    p = np.exp(m)
+    low = p[:40].sum(axis=0)   # ~ <500 Hz
+    high = p[64:].sum(axis=0)  # upper bands
+    events = []
+    for fr in peaks:
+        fr2 = min(fr + 1, m.shape[1] - 1)
+        ratio = high[fr2] / (low[fr2] + 1e-9)
+        events.append((float(fr / FPS), "ka" if ratio > 0.35 else "don"))
+    return events

softchart/data.py

@@ -0,0 +1,350 @@
+"""PyTorch dataset: (mel window, token sequence) pairs from the preprocessed cache."""
+
+import json
+import os
+import random
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from .vocab import FPS, MAX_TGT, NOTE_CLASSES, VOCAB, WINDOW, encode_window
+
+G_CHUNK = 172  # ~2s per whole-song summary chunk
+G_LEN = 100  # fixed number of summary chunks (covers ~200s songs)
+
+
+def song_summary(mel):
+    """(n_mels, T) log-mel -> (G_LEN, n_mels) coarse whole-song summary."""
+    n_mels, T = mel.shape
+    n = min(G_LEN, max(1, T // G_CHUNK))
+    out = np.full((G_LEN, n_mels), np.log(1e-5), dtype=np.float32)
+    for k in range(n):
+        out[k] = mel[:, k * G_CHUNK : (k + 1) * G_CHUNK].astype(np.float32).mean(axis=1)
+    return out
+
+
+class ChartWindowDataset(Dataset):
+    """One item = one (song, course) chart; a random window is cropped per access."""
+
+    def __init__(
+        self,
+        cache_dir,
+        song_ids,
+        train=True,
+        cond_drop=0.15,
+        spec_augment=True,
+        windows_per_chart=8,
+        use_ctx=False,
+        aux=False,
+        global_ctx=False,
+        importance_sampling=False,
+        tempo_aug=False,
+        sibling=False,
+        style=False,
+        beat_head=False,
+        sync_token=False,
+        plan=False,
+        mask_infill=0.0,
+        complexity=False,
+    ):
+        self.cache = cache_dir
+        self.train = train
+        self.cond_drop = cond_drop if train else 0.0
+        self.spec_augment = spec_augment and train
+        self.use_ctx = use_ctx
+        self.aux = aux
+        self.global_ctx = global_ctx
+        self.importance = importance_sampling and train
+        self.tempo_aug = tempo_aug and train
+        self.sibling = sibling
+        self.beat_head = beat_head
+        self.sync_token = sync_token
+        self.mask_infill = mask_infill if train else 0.0
+        self.complexity = complexity
+        self.plans = None
+        if plan:
+            with open(os.path.join(cache_dir, "plans.json")) as f:
+                self.plans = json.load(f)
+        self.styles = None
+        if style:
+            with open(os.path.join(cache_dir, "styles.json")) as f:
+                self.styles = json.load(f)["styles"]
+        with open(os.path.join(cache_dir, "index.json")) as f:
+            index = {e["id"]: e for e in json.load(f)}
+        self.items = []  # (sid, course, level, n_frames)
+        for sid in song_ids:
+            e = index.get(sid)
+            if e is None:
+                continue
+            for c, info in e["courses"].items():
+                self.items.append((sid, c, info["level"], e["n_frames"]))
+        # virtual epoch length: several windows per chart
+        self.mult = windows_per_chart if train else 1
+        self._rng = random.Random(int(os.environ.get('SC_DATA_SEED', 1234)))
+
+    def __len__(self):
+        return len(self.items) * self.mult
+
+    def _load(self, sid):
+        mel = np.load(os.path.join(self.cache, f"{sid}.mel.npy"), mmap_mode="r")
+        notes = np.load(os.path.join(self.cache, f"{sid}.notes.npz"))
+        return mel, notes
+
+    def _beats(self, sid):
+        try:
+            z = np.load(os.path.join(self.cache, f"{sid}.beats.npz"))
+            return z["beats"], z["downbeats"]
+        except Exception:
+            return None, None
+
+    def __getitem__(self, i):
+        sid, course, level, n_frames = self.items[i % len(self.items)]
+        mel, notes = self._load(sid)
+        times = notes[f"{course}_t"]
+        classes = notes[f"{course}_c"]
+
+        # rhythm-preserving tempo augmentation: crop a longer/shorter source
+        # window and resample it to WINDOW (mel-domain time stretch); note
+        # times scale by the same factor, so all rhythmic relations survive.
+        rng0 = random.Random(i * 31337 + (self._rng.randint(0, 1 << 30) if self.train else 0))
+        r = 1.0
+        if self.tempo_aug and rng0.random() < 0.5:
+            r = rng0.uniform(0.9, 1.111)
+        src_len = int(round(WINDOW * r))
+
+        max_start = max(0, n_frames - src_len)
+        rng = random.Random(i * 7919 + (self._rng.randint(0, 1 << 30) if self.train else 0))
+        if self.train and max_start > 0:
+            if self.importance:
+                # importance sampling: bias window starts toward note-dense /
+                # informative regions (uniform floor keeps sparse regions seen)
+                cands = np.arange(0, max_start + 1, int(FPS))  # 1s stride
+                lo = np.searchsorted(times, cands / FPS)
+                hi = np.searchsorted(times, cands / FPS + WINDOW / FPS)
+                w = (hi - lo).astype(np.float64) + 3.0
+                start = int(rng.choices(cands.tolist(), weights=w.tolist())[0])
+                start = min(max_start, start + rng.randint(0, int(FPS)))  # sub-second jitter
+            else:
+                start = rng.randint(0, max_start)
+        else:  # deterministic pseudo-random window per item (not always the intro)
+            start = (i * 2654435761) % (max_start + 1) if max_start > 0 else 0
+
+        t0 = start / FPS
+        t1 = (start + src_len) / FPS
+        sel = (times >= t0) & (times < t1)
+        frames = np.round((times[sel] - t0) * FPS / r).astype(np.int64)
+        frames = np.clip(frames, 0, WINDOW - 1)
+        pairs = list(zip(frames.tolist(), classes[sel].tolist()))
+
+        sib = None
+        if self.sibling:  # skeleton hint from the nearest easier course
+            order = ["easy", "normal", "hard", "oni", "ura"]
+            sib_pairs = []
+            if not (self.train and rng0.random() < 0.2):  # p=0.2: train without hint
+                for c2 in reversed(order[: order.index(course)]):
+                    if f"{c2}_t" in notes:
+                        t2, c2c = notes[f"{c2}_t"], notes[f"{c2}_c"]
+                        s2 = (t2 >= t0) & (t2 < t1) & (c2c <= 3)  # hits only
+                        f2 = np.clip(np.round((t2[s2] - t0) * FPS / r), 0, WINDOW - 1)
+                        ev = list(zip(f2.astype(np.int64).tolist(), c2c[s2].tolist()))
+                        if len(ev) > 12:  # even subsample to 12 slots
+                            idx2 = np.linspace(0, len(ev) - 1, 12).astype(int)
+                            ev = [ev[k] for k in idx2]
+                        sib_pairs = ev
+                        break
+            sib = sib_pairs
+
+        ctx = None
+        if self.use_ctx:
+            from .generate import CTX_LEN, HIT_CLASSES
+
+            hit_ids = [NOTE_CLASSES.index(c) for c in HIT_CLASSES]
+            prev = (times >= t0 - 3.0) & (times < t0)
+            tail = [int(c) for c in classes[prev] if int(c) in hit_ids][-CTX_LEN:]
+            if self.train and rng.random() < 0.1:
+                tail = []  # simulate missing context (first window)
+            elif self.train:
+                tail = [c if rng.random() > 0.1 else rng.choice(hit_ids) for c in tail]
+            ctx = [-1] * (CTX_LEN - len(tail)) + tail
+
+        style = None
+        if self.styles is not None:
+            style = self.styles.get(f"{sid}:{course}", -1)
+
+        sync_band = None
+        beats_w = db_w = None
+        if self.sync_token or self.beat_head:
+            src_sid = sid.split("x")[0] if "x" in sid else sid
+            bts, dbs = self._beats(src_sid)
+            if bts is not None:
+                # window-local, tempo-scaled beat grid
+                bsel = (bts >= t0) & (bts < t1)
+                beats_w = (bts[bsel] - t0) / r
+                dsel = (dbs >= t0) & (dbs < t1)
+                db_w = (dbs[dsel] - t0) / r
+        if self.sync_token:
+            sync_band = lhl_band(
+                np.array([f / FPS for f, c in pairs if NOTE_CLASSES[c] != "end"]),
+                beats_w, db_w)
+
+        plan_slice = None
+        if self.plans is not None:
+            blocks = self.plans.get(f"{sid}:{course}") or []
+            # blocks overlapping this window; times already in source timescale,
+            # window is [t0, t1) in the same scale (r rescales note frames only,
+            # block membership is decided in source time)
+            plan_slice = [(b[2], b[3]) for b in blocks if b[1] > t0 and b[0] < t1]
+
+        cplx = None
+        if self.complexity:
+            from .vocab import complexity_band
+            cplx = complexity_band(pairs)
+
+        seq, prefix_len = encode_window(
+            VOCAB, course, level, pairs, cond_drop=self.cond_drop, rng=rng,
+            ctx_types=ctx, sib_pairs=sib, style=style, sync_band=sync_band,
+            plan_slice=plan_slice, complexity=cplx,
+        )
+        in_seq = None
+        if self.mask_infill > 0 and rng.random() < self.mask_infill:
+            # skeleton->color curriculum: decoder INPUT sees MASK where note
+            # types were; gold targets keep the true types
+            note_ids = set(VOCAB.note.values())
+            in_seq = [VOCAB.mask if (i >= prefix_len and tok in note_ids) else tok
+                      for i, tok in enumerate(seq)]
+        if len(seq) > MAX_TGT:  # truncate overly dense windows, keep EOS
+            seq = seq[: MAX_TGT - 1] + [VOCAB.eos]
+            if VOCAB.is_time(seq[-2]):  # do not end on a dangling TIME token
+                seq = seq[:-2] + [VOCAB.eos]
+
+        x = mel[:, start : start + src_len].astype(np.float32)
+        if x.shape[1] < src_len:
+            x = np.pad(x, ((0, 0), (0, src_len - x.shape[1])), constant_values=np.log(1e-5))
+        x = torch.from_numpy(x)
+        if src_len != WINDOW:  # mel-domain time stretch to the fixed window size
+            x = torch.nn.functional.interpolate(
+                x[None], size=WINDOW, mode="linear", align_corners=False
+            )[0]
+        if self.tempo_aug and rng0.random() < 0.3:  # mild pitch shift (mel-bin roll)
+            k = rng0.randint(-4, 4)
+            if k:
+                x = torch.roll(x, k, dims=0)
+                if k > 0:
+                    x[:k] = x.min()
+                else:
+                    x[k:] = x.min()
+        if self.spec_augment:
+            for _ in range(2):  # freq masks
+                w = rng.randint(0, 12)
+                f0 = rng.randint(0, x.shape[0] - w) if w else 0
+                if w:
+                    x[f0 : f0 + w] = x.mean()
+            for _ in range(2):  # time masks
+                w = rng.randint(0, 24)
+                f0 = rng.randint(0, x.shape[1] - w) if w else 0
+                if w:
+                    x[:, f0 : f0 + w] = x.mean()
+            x = x + (rng.random() * 1.38 - 0.69)  # gain +-6 dB in log space
+
+        aux_t = torch.zeros(0)
+        if self.aux:  # onset heatmap at encoder resolution (WINDOW // 4)
+            aux_t = torch.zeros(WINDOW // 4)
+            for f, c in pairs:
+                if NOTE_CLASSES[c] != "end":
+                    aux_t[min(f // 4, WINDOW // 4 - 1)] = 1.0
+            aux_t = torch.max(aux_t, 0.5 * torch.roll(aux_t, 1))
+            aux_t = torch.max(aux_t, 0.5 * torch.roll(aux_t, -1))
+
+        beat_t = torch.zeros(0)
+        if self.beat_head:
+            beat_t = torch.zeros(WINDOW // 4, 2)
+            for arr, ch in ((beats_w, 0), (db_w, 1)):
+                if arr is None:
+                    continue
+                for bt in arr:
+                    k = int(bt * FPS) // 4
+                    if 0 <= k < WINDOW // 4:
+                        beat_t[k, ch] = 1.0
+            for ch in range(2):
+                col = beat_t[:, ch]
+                beat_t[:, ch] = torch.max(col, 0.5 * torch.roll(col, 1))
+                beat_t[:, ch] = torch.max(beat_t[:, ch], 0.5 * torch.roll(col, -1))
+
+        gsum = torch.zeros(0)
+        posb = 0
+        if self.global_ctx:
+            gsum = torch.from_numpy(song_summary(np.asarray(mel)))
+            posb = min(15, int(16 * start / max(n_frames, 1)))
+
+        in_t = torch.tensor(in_seq if in_seq is not None else seq, dtype=torch.long)
+        return (x, torch.tensor(seq, dtype=torch.long), prefix_len, aux_t, gsum, posb,
+                beat_t, in_t)
+
+
+def lhl_band(note_times, beats, downbeats, band_width=0.6):
+    """Longuet-Higgins & Lee (1984) style syncopation, simplified to the
+    eighth-note grid, normalized per bar, bucketed into 6 bands."""
+    if beats is None or len(beats) < 4 or len(note_times) < 2:
+        return -1
+    beat_len = float(np.median(np.diff(beats))) if len(beats) > 1 else 0.5
+    grid = []  # (time, metrical level)
+    dbset = set(np.round(downbeats, 3)) if downbeats is not None else set()
+    for b in beats:
+        lvl = 3 if round(float(b), 3) in dbset else 2
+        grid.append((float(b), lvl))
+        grid.append((float(b) + beat_len / 2, 1))  # eighth positions
+    grid.sort()
+    tol = 0.04
+    score = 0.0
+    for i, (gt_pos, lvl) in enumerate(grid[:-1]):
+        has_note = np.any(np.abs(note_times - gt_pos) < tol)
+        nxt_pos, nxt_lvl = grid[i + 1]
+        nxt_note = np.any(np.abs(note_times - nxt_pos) < tol)
+        if has_note and not nxt_note and nxt_lvl > lvl:
+            score += nxt_lvl - lvl
+    n_bars = max(1.0, len([g for g in grid if g[1] == 3]))
+    return min(5, int(score / n_bars / band_width))
+
+
+def collate(batch):
+    xs, seqs, plens, auxs, gsums, posbs, beats, inseqs = zip(*batch)
+    x = torch.stack(xs)
+    maxlen = max(len(s) for s in seqs)
+    tgt = torch.full((len(seqs), maxlen), VOCAB.pad, dtype=torch.long)
+    loss_mask = torch.zeros((len(seqs), maxlen), dtype=torch.bool)
+    for i, (s, pl) in enumerate(zip(seqs, plens)):
+        tgt[i, : len(s)] = s
+        loss_mask[i, pl : len(s)] = True  # loss on events + EOS only
+    aux = torch.stack(auxs) if auxs[0].numel() else None
+    gsum = torch.stack(gsums) if gsums[0].numel() else None
+    posb = torch.tensor(posbs, dtype=torch.long)
+    beat = torch.stack(beats) if beats[0].numel() else None
+    in_tgt = torch.full((len(inseqs), maxlen), VOCAB.pad, dtype=torch.long)
+    for i, sq in enumerate(inseqs):
+        in_tgt[i, : len(sq)] = sq
+    return x, tgt, loss_mask, aux, gsum, posb, beat, in_tgt
+
+
+def load_split_ids(cache_dir, val_songs=40, seed=42, use_aug=False):
+    """Song-level split: train / val (carved from train) / test (provided).
+
+    Augmented variants (speed-aug entries with "aug": true) are never used for
+    val, and variants whose source song is in val are excluded from train
+    (no leakage). The base split is stable regardless of augmentation.
+    """
+    with open(os.path.join(cache_dir, "index.json")) as f:
+        index = json.load(f)
+    base_train = sorted(e["id"] for e in index
+                        if e["split"] == "train" and not e.get("aug"))
+    test_ids = sorted(e["id"] for e in index if e["split"] == "test")
+    rng = random.Random(seed)
+    rng.shuffle(base_train)
+    val_ids = base_train[:val_songs]
+    train_ids = base_train[val_songs:]
+    if use_aug:
+        val_set = set(val_ids)
+        train_ids = train_ids + sorted(
+            e["id"] for e in index
+            if e.get("aug") and e.get("src") not in val_set)
+    return train_ids, val_ids, test_ids

softchart/evaluate.py

@@ -0,0 +1,151 @@
+"""Objective evaluation metrics for generated charts."""
+
+import numpy as np
+
+HIT_CLASSES = ("don", "ka", "don_big", "ka_big")
+
+
+def match_onsets(ref_times, est_times, tol=0.05):
+    """Greedy bipartite matching (mir_eval style). Returns list of (ri, ei)."""
+    matches = []
+    ri, ei = 0, 0
+    used_r = set()
+    # standard approach: for each est, match nearest unmatched ref within tol
+    ref = np.asarray(ref_times)
+    est = np.asarray(est_times)
+    order = np.argsort(est)
+    ref_order = np.argsort(ref)
+    ref_sorted = ref[ref_order]
+    taken = np.zeros(len(ref), dtype=bool)
+    for e_i in order:
+        t = est[e_i]
+        lo = np.searchsorted(ref_sorted, t - tol)
+        hi = np.searchsorted(ref_sorted, t + tol)
+        best, best_d = -1, tol + 1
+        for k in range(lo, hi):
+            rid = ref_order[k]
+            if taken[rid]:
+                continue
+            d = abs(ref_sorted[k] - t)
+            if d < best_d:
+                best, best_d = rid, d
+        if best >= 0:
+            taken[best] = True
+            matches.append((best, e_i))
+    return matches
+
+
+def onset_prf(ref_times, est_times, tol=0.05):
+    if len(ref_times) == 0 and len(est_times) == 0:
+        return 1.0, 1.0, 1.0
+    if len(ref_times) == 0 or len(est_times) == 0:
+        return 0.0, 0.0, 0.0
+    m = len(match_onsets(ref_times, est_times, tol))
+    p = m / len(est_times)
+    r = m / len(ref_times)
+    f = 2 * p * r / (p + r) if p + r > 0 else 0.0
+    return p, r, f
+
+
+def type_accuracy(ref, est, tol=0.05):
+    """ref/est: list of (time, class). Accuracy + per-class stats on matched onsets."""
+    if not ref or not est:
+        return None
+    rt = [t for t, _ in ref]
+    et = [t for t, _ in est]
+    matches = match_onsets(rt, et, tol)
+    if not matches:
+        return None
+    correct = sum(1 for ri, ei in matches if ref[ri][1] == est[ei][1])
+    per_class = {}
+    for ri, ei in matches:
+        c = ref[ri][1]
+        d = per_class.setdefault(c, [0, 0])
+        d[1] += 1
+        if est[ei][1] == c:
+            d[0] += 1
+    # binary don/ka accuracy (big variants folded in)
+    fold = lambda c: "don" if "don" in c else ("ka" if "ka" in c else c)
+    correct2 = sum(1 for ri, ei in matches if fold(ref[ri][1]) == fold(est[ei][1]))
+    return {
+        "acc": correct / len(matches),
+        "acc_donka": correct2 / len(matches),
+        "n_matched": len(matches),
+        "per_class": {c: v[0] / v[1] for c, v in per_class.items()},
+    }
+
+
+def density(events, span=None):
+    times = [t for t, c in events if c in HIT_CLASSES]
+    if len(times) < 2:
+        return 0.0
+    span = span or (max(times) - min(times))
+    return len(times) / max(span, 1e-6) if span > 5 else 0.0
+
+
+def ioi_hist(events, bpm, bins=None):
+    """Inter-onset intervals in beat units, snapped to musical fractions."""
+    if bins is None:
+        bins = [0.125, 1 / 6, 0.25, 1 / 3, 0.5, 2 / 3, 0.75, 1.0, 1.5, 2.0, 3.0, 4.0]
+    times = sorted(t for t, c in events if c in HIT_CLASSES)
+    if len(times) < 3 or not bpm or bpm <= 0:
+        return None
+    ioi = np.diff(times) * bpm / 60.0  # beats
+    hist = np.zeros(len(bins) + 1)
+    for x in ioi:
+        d = [abs(x - b) / b for b in bins]
+        j = int(np.argmin(d))
+        if d[j] < 0.2:
+            hist[j] += 1
+        else:
+            hist[-1] += 1  # off-grid
+    s = hist.sum()
+    return hist / s if s > 0 else None
+
+
+def js_divergence(p, q, eps=1e-9):
+    p = np.asarray(p) + eps
+    q = np.asarray(q) + eps
+    p, q = p / p.sum(), q / q.sum()
+    m = (p + q) / 2
+    kl = lambda a, b: float(np.sum(a * np.log(a / b)))
+    return 0.5 * kl(p, m) + 0.5 * kl(q, m)
+
+
+def ngram_dist(events, n=2):
+    """Distribution over note-class n-grams (hits only)."""
+    seq = [c for _, c in sorted(events) if c in HIT_CLASSES]
+    if len(seq) < n + 1:
+        return {}
+    counts = {}
+    for i in range(len(seq) - n + 1):
+        g = tuple(seq[i : i + n])
+        counts[g] = counts.get(g, 0) + 1
+    tot = sum(counts.values())
+    return {g: c / tot for g, c in counts.items()}
+
+
+def ngram_js(ref_events, est_events, n=2):
+    pr = ngram_dist(ref_events, n)
+    pe = ngram_dist(est_events, n)
+    if not pr or not pe:
+        return None
+    keys = sorted(set(pr) | set(pe))
+    return js_divergence([pr.get(k, 0) for k in keys], [pe.get(k, 0) for k in keys])
+
+
+def distinct_n(events, n=3):
+    seq = [c for _, c in sorted(events) if c in HIT_CLASSES]
+    if len(seq) < n:
+        return None
+    grams = [tuple(seq[i : i + n]) for i in range(len(seq) - n + 1)]
+    return len(set(grams)) / len(grams)
+
+
+def spearman(x, y):
+    from scipy.stats import spearmanr
+
+    if len(x) < 2:
+        return None
+    r = spearmanr(x, y).statistic
+    return None if np.isnan(r) else float(r)

softchart/generate.py

@@ -0,0 +1,404 @@
+"""Inference: full-song chart generation with constrained autoregressive decoding.
+
+Features:
+- structural constraints via logit masking: TIME/NOTE alternation, strictly
+  increasing times, minimum playable gap (4 frames ~= 46 ms);
+- classifier-free guidance (cfg_w > 1): contrasts conditional vs
+  condition-UNK logits to strengthen condition adherence;
+- previous-window pattern context (v2 models): two-pass generation, pass 2
+  conditions each window on the tail of the previous window's pass-1 output.
+"""
+
+import numpy as np
+import torch
+
+from .model import ChartModel
+from .vocab import FPS, MAX_TGT, N_LEVELS, NOTE_CLASSES, VOCAB, WINDOW
+
+HIT_CLASSES = ("don", "ka", "don_big", "ka_big")
+SPAN_CLASSES = ("roll", "roll_big", "balloon")
+
+
+def _dev_type(device):
+    return device.split(":")[0] if isinstance(device, str) else device.type
+
+
+def _autocast(device):
+    """bf16 autocast on CUDA; no-op elsewhere (MPS/CPU run fp32)."""
+    import contextlib
+
+    if _dev_type(device) == "cuda":
+        return torch.autocast("cuda", dtype=torch.bfloat16)
+    return contextlib.nullcontext()
+CTX_LEN = 12  # fixed-length previous-window context (note-type tokens)
+MIN_GAP_FRAMES = 4  # ~46 ms minimum inter-note gap enforced during decoding
+LATTICE_RATIOS = (1 / 3, 0.5, 2 / 3, 1.0, 4 / 3, 1.5, 2.0, 3.0)  # allowed IOI ratios
+# GT span-length p99 per class (train stats); generated spans beyond this are
+# truncated — a visual audit found generated balloons up to ~10s vs GT p50 ~1s
+SPAN_MAX = {"roll": 3.5, "roll_big": 4.0, "balloon": 6.5}
+
+
+def load_model(ckpt_path, device="cuda"):
+    ck = torch.load(ckpt_path, map_location=device)
+    sd = ck["model"]
+    vocab_size = sd["tok_emb.weight"].shape[0]
+    aux = any(k.startswith("aux") for k in sd)
+    gctx = any(k.startswith("gsum_proj") for k in sd)
+    a = ck.get("args", {}) or {}
+    model = ChartModel(
+        d_model=a.get("d_model", 512), enc_layers=a.get("enc_layers", 6),
+        dec_layers=a.get("dec_layers", 6), ffn=a.get("ffn", 2048),
+        vocab_size=vocab_size, aux=aux, global_ctx=gctx,
+        func_time=a.get("func_time", False),
+    ).to(device)
+    if any(k.startswith("ptr") for k in sd):
+        model.enable_ptr()
+    if any(k.startswith("beat") for k in sd):
+        model.enable_beat()
+    model.to(device)  # newly enabled heads default to CPU
+    model.load_state_dict(sd)
+    model.eval()
+    # capability detection: prefer explicit training args (vocab-size inference
+    # wrongly added prefix slots to models that never trained with them)
+    if a:
+        model._has_ctx = bool(a.get("ctx", False))
+        model._has_sib = bool(a.get("sibling", False))
+        model._has_style = bool(a.get("style", False))
+        model._has_sync = bool(a.get("sync_token", False))
+        model._has_plan = bool(a.get("plan", False))
+    else:  # legacy fallback
+        model._has_ctx = vocab_size > VOCAB.sep
+        model._has_sib = vocab_size > VOCAB.sib
+        model._has_style = False
+        model._has_sync = False
+        model._has_plan = False
+    model._has_gctx = gctx
+    return model
+
+
+def load_hf(repo_or_dir, device="cuda"):
+    """Load a Hugging Face SoftChartGenerator (safetensors) and attach the
+    capability flags that generate_song() reads, so HF models are drop-in."""
+    from .hf import SoftChartGenerator
+
+    hf = SoftChartGenerator.from_pretrained(repo_or_dir).to(device).eval()
+    m = hf.net
+    caps = hf.capabilities or {}
+    m._has_ctx = bool(caps.get("ctx", False))
+    m._has_sib = bool(caps.get("sibling", False))
+    m._has_style = bool(caps.get("style", False))
+    m._has_sync = bool(caps.get("sync_token", False))
+    m._has_plan = bool(caps.get("plan", False))
+    m._has_gctx = m.gsum_proj is not None
+    return m
+
+
+def build_prefix(course=None, level=None, density_bucket=None, ctx_types=None,
+                 uncond=False, sib_pairs=None, style=None, sync_band=None,
+                 plan_slice=None):
+    v = VOCAB
+    seq = [v.bos]
+    if ctx_types is not None:
+        ctx = list(ctx_types)[-CTX_LEN:]
+        ctx = [v.unk_cond] * (CTX_LEN - len(ctx)) + [v.note[NOTE_CLASSES[c]] if isinstance(c, int) else v.note[c] for c in ctx]
+        seq += ctx + [v.sep]
+    if sib_pairs is not None:  # fixed 12-event sibling segment (unk-padded)
+        from .vocab import SIB_EVENTS
+
+        seq.append(v.sib)
+        pairs = list(sib_pairs)[:SIB_EVENTS]
+        for f, c in pairs:
+            cls = c if isinstance(c, int) else NOTE_CLASSES.index(c)
+            seq += [v.time(int(f)), v.note[NOTE_CLASSES[cls]]]
+        seq += [v.unk_cond] * (2 * (SIB_EVENTS - len(pairs)))
+    if uncond:
+        seq += [v.unk_cond] * 3
+    else:
+        seq += [
+            v.course[course] if course else v.unk_cond,
+            v.level[max(1, min(N_LEVELS, level))] if level else v.unk_cond,
+            v.dens[density_bucket] if density_bucket is not None else v.unk_cond,
+        ]
+    if style is not None:  # -1 = style-capable model, no specific style requested
+        seq.append(v.style[style] if style >= 0 else v.unk_cond)
+    if sync_band is not None:
+        seq.append(v.sync[sync_band] if sync_band >= 0 else v.unk_cond)
+    if plan_slice is not None:  # fixed PLAN_SLOTS blocks, unk-padded
+        from .vocab import PLAN_SLOTS
+
+        seq.append(v.plan)
+        blocks = list(plan_slice)[:PLAN_SLOTS]
+        for d8, fl in blocks:
+            seq += [v.pdens[min(7, max(0, int(d8)))], v.pflag[min(2, max(0, int(fl)))]]
+        seq += [v.unk_cond] * (2 * (PLAN_SLOTS - len(blocks)))
+    return seq
+
+
+@torch.no_grad()
+def decode_windows(model, mels, prefixes, device="cuda", temperature=1.0, top_p=0.95,
+                   greedy=False, seed=0, cfg_w=0.0, gsum=None, pos_buckets=None,
+                   lattice=False, n_cond=3):
+    """mels: (B, n_mels, WINDOW); prefixes: list of B equal-length token lists.
+    Returns list of lists of (frame, note_class)."""
+    v = VOCAB
+    B = mels.shape[0]
+    plen = len(prefixes[0])
+    assert all(len(p) == plen for p in prefixes)
+    use_cfg = cfg_w and cfg_w > 0 and not greedy
+
+    gen_device = "cpu" if _dev_type(device) == "mps" else device
+    gen = torch.Generator(device=gen_device)
+    gen.manual_seed(seed)
+    with _autocast(device):
+        memory = model.encode(
+            mels.to(device),
+            gsum=gsum.to(device) if gsum is not None else None,
+            pos_bucket=pos_buckets.to(device) if pos_buckets is not None else None,
+        )
+
+    seqs = torch.tensor(prefixes, dtype=torch.long, device=device)
+    if use_cfg:  # rows B..2B: same ctx, conditions replaced by UNK
+        unc = seqs.clone()
+        unc[:, -n_cond:] = v.unk_cond
+        seqs = torch.cat([seqs, unc], dim=0)
+        memory = torch.cat([memory, memory], dim=0)
+
+    done = torch.zeros(B, dtype=torch.bool, device=device)
+    last_time = torch.full((B,), -MIN_GAP_FRAMES, dtype=torch.long, device=device)
+    expect_note = torch.zeros(B, dtype=torch.bool, device=device)
+    note_ids = torch.tensor(sorted(v.note.values()), device=device)
+    # lattice state: previous inter-onset interval (frames); -1 = unknown
+    prev_ioi = torch.full((B,), -1, dtype=torch.long, device=device)
+    ratios = torch.tensor(LATTICE_RATIOS, device=device)
+
+    for _ in range(MAX_TGT - plen):
+        with _autocast(device):
+            logits_all = model.decode(seqs, memory)[:, -1].float()
+        if use_cfg:
+            logits = logits_all[B:] + cfg_w * (logits_all[:B] - logits_all[B:])
+        else:
+            logits = logits_all
+
+        mask = torch.full_like(logits, float("-inf"))
+        note_row = torch.full((logits.shape[1],), float("-inf"), device=device)
+        note_row[note_ids] = 0.0
+        mask[expect_note] = note_row
+        idx = (~expect_note).nonzero(as_tuple=True)[0]
+        if len(idx):
+            time_pos = torch.arange(WINDOW, device=device).unsqueeze(0)
+            sub = torch.full((len(idx), logits.shape[1]), float("-inf"), device=device)
+            sub[:, v.eos] = 0.0
+            ok = time_pos >= (last_time[idx] + MIN_GAP_FRAMES).unsqueeze(1)
+            tmask = torch.where(
+                ok, torch.zeros_like(sub[:, :WINDOW]),
+                torch.full_like(sub[:, :WINDOW], float("-inf")),
+            )
+            if lattice:
+                # soft rhythmic-lattice constraint: once an IOI is established,
+                # penalize next onsets whose IOI ratio is not a musical fraction
+                # (kills "between two subdivisions" notes at the source)
+                pi = prev_ioi[idx]
+                active = (pi >= MIN_GAP_FRAMES) & (pi <= 86)  # sub-second IOIs only
+                if active.any():
+                    delta = (time_pos - last_time[idx].unsqueeze(1)).float()  # (n, W)
+                    rel = delta / pi.unsqueeze(1).clamp(min=1).float()
+                    err = (rel.unsqueeze(-1) / ratios - 1.0).abs().min(-1).values
+                    bad = (err > 0.13) & (rel < 4.0) & (delta > 0)
+                    pen = torch.where(bad & active.unsqueeze(1),
+                                      torch.full_like(delta, -6.0),
+                                      torch.zeros_like(delta))
+                    tmask = tmask + pen
+            sub[:, v.time0 : v.time0 + WINDOW] = tmask
+            mask[idx] = sub
+        logits = logits + mask
+
+        if greedy:
+            nxt = logits.argmax(-1)
+        else:
+            probs = torch.softmax(logits / temperature, dim=-1)
+            sp, si = torch.sort(probs, descending=True, dim=-1)
+            cum = torch.cumsum(sp, dim=-1)
+            keep = cum - sp < top_p
+            keep[:, 0] = True
+            sp = sp * keep
+            sp = sp / sp.sum(-1, keepdim=True)
+            if _dev_type(device) == "mps":
+                pick = torch.multinomial(sp.cpu(), 1, generator=gen).squeeze(1).to(device)
+            else:
+                pick = torch.multinomial(sp, 1, generator=gen).squeeze(1)
+            nxt = si[torch.arange(B, device=device), pick]
+
+        nxt = torch.where(done, torch.full_like(nxt, v.pad), nxt)
+        step_tok = torch.cat([nxt, nxt], dim=0) if use_cfg else nxt
+        seqs = torch.cat([seqs, step_tok.unsqueeze(1)], dim=1)
+        is_time = (nxt >= v.time0) & (nxt < v.time0 + WINDOW)
+        new_time = nxt - v.time0
+        upd = is_time & (last_time >= 0)
+        prev_ioi = torch.where(upd, new_time - last_time, prev_ioi)
+        last_time = torch.where(is_time, new_time, last_time)
+        expect_note = is_time
+        done = done | (nxt == v.eos)
+        if done.all():
+            break
+
+    out = []
+    for b in range(B):
+        toks = seqs[b, plen:].tolist()
+        events = []
+        cur_t = None
+        for t in toks:
+            if t == v.eos or t == v.pad:
+                break
+            if v.time0 <= t < v.time0 + WINDOW:
+                cur_t = t - v.time0
+            elif t in v.id2note and cur_t is not None:
+                events.append((cur_t, v.id2note[t]))
+        out.append(events)
+    return out
+
+
+def _decode_pass(model, wins, starts, course, level, density_bucket, ctx_lists,
+                 device, greedy, temperature, top_p, seed, cfg_w, batch_windows,
+                 gsum=None, T=None, lattice=False, sib_default=False, style=None,
+                 sync_band=None, plan_blocks=None, plan_default=False):
+    all_events = []
+    per_window = []
+    for i in range(0, len(wins), batch_windows):
+        chunk = torch.stack(wins[i : i + batch_windows])
+        prefixes = [
+            build_prefix(course, level, density_bucket,
+                         ctx_types=(ctx_lists[i + j] if ctx_lists is not None else None),
+                         sib_pairs=([] if sib_default else None), style=style,
+                         sync_band=sync_band,
+                         plan_slice=(
+                             [(b[2], b[3]) for b in plan_blocks
+                              if b[1] > starts[i + j] / FPS
+                              and b[0] < (starts[i + j] + WINDOW) / FPS]
+                             if plan_blocks is not None
+                             else ([] if plan_default else None)))
+            for j in range(chunk.shape[0])
+        ]
+        g = pb = None
+        if gsum is not None:
+            g = gsum.unsqueeze(0).expand(chunk.shape[0], -1, -1)
+            pb = torch.tensor(
+                [min(15, int(16 * starts[i + j] / max(T, 1))) for j in range(chunk.shape[0])],
+                dtype=torch.long)
+        evs = decode_windows(model, chunk, prefixes, device=device, greedy=greedy,
+                             temperature=temperature, top_p=top_p, seed=seed + i,
+                             cfg_w=cfg_w, gsum=g, pos_buckets=pb, lattice=lattice,
+                             n_cond=4 if style is not None else 3)
+        for j, events in enumerate(evs):
+            per_window.append(events)
+            t_off = starts[i + j] / FPS
+            for f, cls in events:
+                all_events.append((t_off + f / FPS, cls))
+    return all_events, per_window
+
+
+@torch.no_grad()
+def generate_song(model, mel, course, level=None, density_bucket=None, device="cuda",
+                  greedy=False, temperature=1.0, top_p=0.95, seed=0, batch_windows=8,
+                  cfg_w=0.0, use_ctx=None, lattice=False, style=None, sync_band=None,
+                  plan=None):
+    """mel: (n_mels, T). Returns dict with 'hits' and 'spans' in seconds."""
+    if isinstance(mel, np.ndarray):
+        mel = torch.from_numpy(mel.astype(np.float32))
+    T = mel.shape[1]
+    starts = list(range(0, max(T - 1, 1), WINDOW))
+    wins = []
+    for s in starts:
+        w = mel[:, s : s + WINDOW]
+        if w.shape[1] < WINDOW:
+            w = torch.nn.functional.pad(w, (0, WINDOW - w.shape[1]), value=float(np.log(1e-5)))
+        wins.append(w)
+
+    if use_ctx is None:
+        use_ctx = getattr(model, "_has_ctx", False)
+    sib_default = getattr(model, "_has_sib", False)
+    style_val = None
+    if getattr(model, "_has_style", False):
+        style_val = style if style is not None else -1
+    sync_val = None
+    if getattr(model, "_has_sync", False):
+        sync_val = sync_band if sync_band is not None else -1
+    plan_blocks = plan if getattr(model, "_has_plan", False) else None
+    plan_default = getattr(model, "_has_plan", False) and plan is None
+    gsum = None
+    if getattr(model, "_has_gctx", False):
+        from .data import song_summary
+
+        gsum = torch.from_numpy(song_summary(mel.numpy()))
+
+    all_events, per_window = _decode_pass(
+        model, wins, starts, course, level, density_bucket,
+        None if not use_ctx else [[] for _ in wins],  # pass 1: empty ctx
+        device, greedy, temperature, top_p, seed, cfg_w, batch_windows,
+        gsum=gsum, T=T, lattice=lattice, sib_default=sib_default, style=style_val,
+        sync_band=sync_val, plan_blocks=plan_blocks, plan_default=plan_default)
+
+    if use_ctx and len(wins) > 1:
+        # pass 2: condition each window on the tail of the previous window's pass-1 output
+        ctx_lists = [[]]
+        for w_ev in per_window[:-1]:
+            tail = [NOTE_CLASSES.index(c) for _, c in w_ev if c in HIT_CLASSES][-CTX_LEN:]
+            ctx_lists.append(tail)
+        all_events, per_window = _decode_pass(
+            model, wins, starts, course, level, density_bucket, ctx_lists,
+            device, greedy, temperature, top_p, seed, cfg_w, batch_windows,
+            gsum=gsum, T=T, lattice=lattice, sib_default=sib_default, style=style_val,
+            sync_band=sync_val, plan_blocks=plan_blocks, plan_default=plan_default)
+
+    if True:
+        # rescue decoding: decoding occasionally EOSes a whole window early; if a
+        # window is near-empty while its audio is musically active, redo it with
+        # the opposite mode (sampling was empty -> greedy; greedy was empty ->
+        # sampled retry, since a greedy redo would reproduce the same output)
+        flux = [float(np.maximum(0, np.diff(w.numpy(), axis=1)).sum()) for w in wins]
+        med = float(np.median(flux)) if flux else 0.0
+        med_hits = float(np.median([len(ev) for ev in per_window])) if per_window else 0.0
+        retry = [i for i, ev in enumerate(per_window)
+                 if (len(ev) < 4 and flux[i] > 0.3 * med)
+                 or (len(ev) < 0.4 * med_hits and flux[i] > 0.7 * med)]
+        if retry:
+            r_evs, _ = _decode_pass(
+                model, [wins[i] for i in retry], [starts[i] for i in retry],
+                course, level, density_bucket, None, device, not greedy, 0.9,
+                top_p, seed + 7, 0.0, batch_windows, gsum=gsum, T=T,
+                lattice=lattice, sib_default=sib_default, style=style_val,
+                sync_band=sync_val, plan_blocks=plan_blocks, plan_default=plan_default)
+            all_events = [e for i, w_ev in enumerate(per_window) if i not in retry
+                          for e in [(starts[i] / FPS + f / FPS, c) for f, c in w_ev]]
+            all_events += r_evs
+
+    all_events.sort(key=lambda e: e[0])
+    hits, spans = [], []
+    open_span = None
+    last_hit_t = -1.0
+    for t, cls in all_events:
+        if cls in HIT_CLASSES:
+            if open_span is not None:
+                # a span whose end never arrived would swallow every following
+                # hit (visual audit: 20 s of silence) — force-close at SPAN_MAX
+                if t - open_span[0] > SPAN_MAX.get(open_span[1], 6.5):
+                    spans.append({"t0": round(open_span[0], 4),
+                                  "t1": round(open_span[0] + SPAN_MAX.get(open_span[1], 6.5), 4),
+                                  "type": open_span[1]})
+                    open_span = None
+                else:
+                    continue  # no hits inside an open span
+            if t - last_hit_t < 0.025:
+                continue
+            hits.append({"t": round(t, 4), "type": cls})
+            last_hit_t = t
+        elif cls in SPAN_CLASSES:
+            if open_span is None:
+                open_span = (t, cls)
+        elif cls == "end":
+            if open_span is not None and t - open_span[0] > 0.05:
+                t1_span = min(t, open_span[0] + SPAN_MAX.get(open_span[1], 6.5))
+                spans.append({"t0": round(open_span[0], 4), "t1": round(t1_span, 4),
+                              "type": open_span[1]})
+                open_span = None
+    return {"hits": hits, "spans": spans, "course": course, "level": level,
+            "density_bucket": density_bucket}

softchart/hf.py

@@ -0,0 +1,91 @@
+"""Hugging Face Hub-compatible wrappers (safetensors + from_pretrained/push_to_hub).
+
+Three model types, all <=16M params, MIT-licensed:
+  SoftChartGenerator  audio log-mel  -> chart event tokens (the main model; pr12)
+  SoftChartBeat       audio log-mel  -> beat/downbeat activations (barline anchor)
+  SoftChartPlanner    block features -> per-block plan tokens (auto-planning)
+
+Usage:
+    from softchart.hf import SoftChartGenerator
+    gen = SoftChartGenerator.from_pretrained("JacobLinCool/softchart-generator")
+"""
+
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+
+from .model import ChartModel, sinusoidal
+
+_CARD = "See https://github.com/JacobLinCool/SoftChart — MIT licensed."
+
+
+class SoftChartGenerator(
+    nn.Module,
+    PyTorchModelHubMixin,
+    repo_url="https://github.com/JacobLinCool/SoftChart",
+    license="mit",
+    tags=["taiko", "rhythm-game", "chart-generation", "music", "audio-to-symbolic"],
+):
+    """Encoder-decoder chart generator. Config is stored as config.json and the
+    full architecture (condition tokens, heads) is reconstructed on load."""
+
+    def __init__(self, d_model=256, enc_layers=4, dec_layers=4, ffn=1024,
+                 vocab_size=1802, aux=True, global_ctx=False, func_time=False,
+                 ptr=False, beat=False, capabilities=None):
+        super().__init__()
+        self.net = ChartModel(
+            d_model=d_model, enc_layers=enc_layers, dec_layers=dec_layers, ffn=ffn,
+            vocab_size=vocab_size, aux=aux, global_ctx=global_ctx, func_time=func_time,
+            ptr=ptr, beat=beat,
+        )
+        # capabilities: which condition axes this checkpoint was trained with,
+        # so inference knows which prefix tokens to emit
+        self.capabilities = capabilities or {}
+
+    # delegate the generation interface to the inner net
+    def encode(self, *a, **k):
+        return self.net.encode(*a, **k)
+
+    def decode(self, *a, **k):
+        return self.net.decode(*a, **k)
+
+    def forward(self, *a, **k):
+        return self.net(*a, **k)
+
+    @property
+    def ptr(self):
+        return self.net.ptr
+
+    @property
+    def beat(self):
+        return self.net.beat
+
+    @property
+    def tok_emb(self):
+        return self.net.tok_emb
+
+
+class SoftChartPlanner(
+    nn.Module,
+    PyTorchModelHubMixin,
+    repo_url="https://github.com/JacobLinCool/SoftChart",
+    license="mit",
+    tags=["taiko", "rhythm-game", "planning", "music"],
+):
+    """Song-level plan generator: block audio summaries -> (density, flag) tokens."""
+
+    def __init__(self, feat_dim=261, d=192, layers=4, n_dens=8, n_flag=3):
+        super().__init__()
+        self.inp = nn.Linear(feat_dim, d)
+        self.course = nn.Embedding(5, d)
+        self.register_buffer("pos", sinusoidal(128, d), persistent=False)
+        enc = nn.TransformerEncoderLayer(d, 6, d * 4, 0.1, activation="gelu",
+                                         batch_first=True, norm_first=True)
+        self.enc = nn.TransformerEncoder(enc, layers, nn.LayerNorm(d))
+        self.h_dens = nn.Linear(d, n_dens)
+        self.h_flag = nn.Linear(d, n_flag)
+
+    def forward(self, x, cid):
+        h = self.inp(x) + self.course(cid)[:, None] + self.pos[: x.shape[1]]
+        h = self.enc(h)
+        return self.h_dens(h), self.h_flag(h)

softchart/model.py

@@ -0,0 +1,173 @@
+"""Encoder-decoder Transformer: log-mel windows -> chart token sequences."""
+
+import math
+
+import torch
+import torch.nn as nn
+
+from .vocab import N_MELS, VOCAB, WINDOW
+
+
+def sinusoidal(length, dim):
+    pos = torch.arange(length, dtype=torch.float32)[:, None]
+    i = torch.arange(dim // 2, dtype=torch.float32)[None]
+    angle = pos / torch.pow(10000.0, 2 * i / dim)
+    emb = torch.zeros(length, dim)
+    emb[:, 0::2] = torch.sin(angle)
+    emb[:, 1::2] = torch.cos(angle)
+    return emb
+
+
+class ChartModel(nn.Module):
+    def __init__(
+        self,
+        d_model=512,
+        nhead=8,
+        enc_layers=6,
+        dec_layers=6,
+        ffn=2048,
+        dropout=0.1,
+        vocab_size=None,
+        max_tgt=None,
+        aux=False,
+        global_ctx=False,
+        func_time=False,
+        ptr=False,
+        beat=False,
+    ):
+        super().__init__()
+        from .vocab import MAX_TGT
+
+        vocab_size = vocab_size or VOCAB.size
+        max_tgt = max_tgt or MAX_TGT
+        self.d_model = d_model
+
+        # conv frontend: (B, n_mels, T) -> (B, T/4, d)
+        self.frontend = nn.Sequential(
+            nn.Conv1d(N_MELS, d_model, 3, stride=2, padding=1),
+            nn.GELU(),
+            nn.Conv1d(d_model, d_model, 3, stride=2, padding=1),
+            nn.GELU(),
+        )
+        enc_len = WINDOW // 4
+        self.register_buffer("enc_pos", sinusoidal(enc_len, d_model), persistent=False)
+        self.register_buffer("dec_pos", sinusoidal(max_tgt, d_model), persistent=False)
+
+        enc_layer = nn.TransformerEncoderLayer(
+            d_model, nhead, ffn, dropout, activation="gelu", batch_first=True, norm_first=True
+        )
+        self.encoder = nn.TransformerEncoder(enc_layer, enc_layers, nn.LayerNorm(d_model))
+        dec_layer = nn.TransformerDecoderLayer(
+            d_model, nhead, ffn, dropout, activation="gelu", batch_first=True, norm_first=True
+        )
+        self.decoder = nn.TransformerDecoder(dec_layer, dec_layers, nn.LayerNorm(d_model))
+
+        self.tok_emb = nn.Embedding(vocab_size, d_model, padding_idx=VOCAB.pad)
+        nn.init.normal_(self.tok_emb.weight, std=0.02)  # keep tied logits well-scaled
+        self.out = nn.Linear(d_model, vocab_size, bias=False)
+        self.out.weight = self.tok_emb.weight  # weight tying
+        self.dropout = nn.Dropout(dropout)
+        # auxiliary per-position onset-heatmap head on the encoder (v2)
+        self.aux = nn.Linear(d_model, 1) if aux else None
+        # pointer alignment head: each generated note must point to its audio
+        # frame (differentiable provenance; explainability that trains alignment)
+        self.ptr = nn.Linear(d_model, d_model) if ptr else None
+        # beat/downbeat head: explicit metrical percept on the encoder
+        self.beat = nn.Linear(d_model, 2) if beat else None
+        # song-level context: coarse whole-song summary + window position (v4).
+        # Lets the model see beyond the window, so intentional gaps (a breath
+        # before a strong section) are informed decisions, not failures.
+        # v2 fix (after the v4 negative result): summary chunks get DEDICATED
+        # positional embeddings instead of reusing window positions.
+        if global_ctx:
+            self.gsum_proj = nn.Linear(N_MELS, d_model)
+            self.seg_emb = nn.Parameter(torch.zeros(2, d_model))
+            self.pos_emb = nn.Embedding(16, d_model)  # window position bucket
+            self.gpos = nn.Parameter(torch.randn(128, d_model) * 0.02)
+        else:
+            self.gsum_proj = None
+        # functional time embeddings: the 1728 TIME tokens share a sinusoidal
+        # basis + small projection instead of free embeddings (fewer params,
+        # neighbouring times get similar representations)
+        if func_time:
+            self.register_buffer("time_basis", sinusoidal(WINDOW, 64), persistent=False)
+            self.time_proj = nn.Linear(64, d_model)
+            # match the 0.02-std scale of tok_emb rows (basis row norm ~ sqrt(32))
+            nn.init.normal_(self.time_proj.weight, std=0.004)
+            nn.init.zeros_(self.time_proj.bias)
+        else:
+            self.time_proj = None
+
+    def encode(self, mel, gsum=None, pos_bucket=None):
+        # mel: (B, n_mels, T); gsum: (B, G, n_mels) whole-song summary chunks;
+        # pos_bucket: (B,) window-position bucket in [0, 16)
+        h = self.frontend(mel).transpose(1, 2)  # (B, T/4, d)
+        h = h + self.enc_pos[: h.shape[1]]
+        if self.gsum_proj is not None and gsum is not None:
+            g = self.gsum_proj(gsum) + self.seg_emb[1] + self.gpos[: gsum.shape[1]]
+            h = h + self.seg_emb[0]
+            p = self.pos_emb(pos_bucket).unsqueeze(1)  # (B, 1, d)
+            h = torch.cat([p, g, h], dim=1)
+        return self.encoder(self.dropout(h))
+
+    def emb_matrix(self):
+        """Full embedding matrix; TIME rows come from the functional basis."""
+        if self.time_proj is None:
+            return self.tok_emb.weight
+        E = self.tok_emb.weight
+        t = self.time_proj(self.time_basis)  # (WINDOW, d)
+        return torch.cat([E[: VOCAB.time0], t, E[VOCAB.time0 + WINDOW :]], dim=0)
+
+    def decode(self, tgt_in, memory):
+        # tgt_in: (B, L) token ids
+        L = tgt_in.shape[1]
+        E = self.emb_matrix()
+        h = nn.functional.embedding(tgt_in, E) * math.sqrt(self.d_model) + self.dec_pos[:L]
+        mask = nn.Transformer.generate_square_subsequent_mask(L, device=tgt_in.device)
+        pad_mask = tgt_in == VOCAB.pad
+        h = self.decoder(
+            self.dropout(h),
+            memory,
+            tgt_mask=mask,
+            tgt_key_padding_mask=pad_mask,
+            tgt_is_causal=True,
+        )
+        return h @ E.T  # tied output projection (functional rows included)
+
+    def enable_ptr(self):
+        self.ptr = nn.Linear(self.d_model, self.d_model)
+
+    def enable_beat(self):
+        self.beat = nn.Linear(self.d_model, 2)
+
+    def forward(self, mel, tgt, return_aux=False, gsum=None, pos_bucket=None,
+                return_extras=False, in_tgt=None):
+        # in_tgt: optional corrupted decoder input (e.g. MASKed types for the
+        # skeleton->color infill curriculum); gold targets stay = tgt
+        memory = self.encode(mel, gsum=gsum, pos_bucket=pos_bucket)
+        dec_in = (in_tgt if in_tgt is not None else tgt)[:, :-1]
+        L = dec_in.shape[1]
+        E = self.emb_matrix()
+        import math as _m
+        h = nn.functional.embedding(dec_in, E) * _m.sqrt(self.d_model) + self.dec_pos[:L]
+        dm = nn.Transformer.generate_square_subsequent_mask(L, device=tgt.device)
+        h = self.decoder(self.dropout(h), memory, tgt_mask=dm,
+                         tgt_key_padding_mask=dec_in == VOCAB.pad, tgt_is_causal=True)
+        logits = h @ E.T
+        if not (return_aux or return_extras):
+            return logits
+        outs = [logits]
+        aux_mem = memory[:, -(WINDOW // 4):]
+        outs.append(self.aux(aux_mem).squeeze(-1) if self.aux is not None else None)
+        if return_extras:
+            ptr_logits = None
+            if self.ptr is not None:
+                q = self.ptr(h)  # (B, L, d)
+                ptr_logits = q @ aux_mem.transpose(1, 2) / _m.sqrt(self.d_model)
+            beat_logits = self.beat(aux_mem) if self.beat is not None else None
+            outs += [ptr_logits, beat_logits]
+        return tuple(outs)
+
+
+def count_params(m):
+    return sum(p.numel() for p in m.parameters() if p.requires_grad)

softchart/preprocess.py

@@ -0,0 +1,163 @@
+"""Convert the HF dataset into per-song cache files: log-mel (.npy) + notes (.npz).
+
+Usage (on the training machine):
+    python -m softchart.preprocess --out /workspace/softchart/cache [--split train]
+"""
+
+import argparse
+import json
+import os
+from collections import Counter
+
+import numpy as np
+
+from .vocab import HOP, N_FFT, N_MELS, NOTE_CLASSES, NOTE_TYPE_MAP, SR, COURSES
+
+
+def to_mono_22k(audio):
+    """Handle both datasets<=3 dict audio and datasets>=4 torchcodec AudioDecoder."""
+    import torch
+
+    if isinstance(audio, dict):
+        arr = np.asarray(audio["array"], dtype=np.float32)
+        if arr.ndim == 2:
+            arr = arr.mean(axis=0)
+        sr = audio["sampling_rate"]
+    else:  # torchcodec AudioDecoder
+        samples = audio.get_all_samples()
+        arr = samples.data.to(torch.float32).mean(dim=0).numpy()
+        sr = samples.sample_rate
+    if sr != SR:
+        import librosa
+
+        arr = librosa.resample(arr, orig_sr=sr, target_sr=SR, res_type="soxr_hq")
+    return torch.from_numpy(np.ascontiguousarray(arr))[None]  # (1, T)
+
+
+_mel_fb = None
+_window = None
+
+
+def logmel(wav):
+    """(1, T) float32 -> (n_mels, frames) float16 log-mel (torch.stft + librosa fb)."""
+    global _mel_fb, _window
+    import torch
+
+    if _mel_fb is None:
+        import librosa
+
+        fb = librosa.filters.mel(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=20.0, fmax=SR / 2)
+        _mel_fb = torch.from_numpy(fb)  # (n_mels, n_fft//2+1)
+        _window = torch.hann_window(N_FFT)
+    with torch.no_grad():
+        spec = torch.stft(
+            wav, N_FFT, hop_length=HOP, window=_window, center=True, return_complex=True
+        )[0]
+        power = spec.abs().pow(2)  # (freq, frames)
+        m = torch.log(_mel_fb @ power + 1e-5)
+    return m.numpy().astype(np.float16)
+
+
+def extract_notes(course_struct):
+    """course struct -> (times, classes, level) or None."""
+    if course_struct is None or not course_struct.get("segments"):
+        return None
+    times, classes, bpms = [], [], []
+    unknown = Counter()
+    for seg in course_struct["segments"]:
+        for n in seg.get("notes") or []:
+            nt = n["note_type"]
+            canon = NOTE_TYPE_MAP.get(nt) or NOTE_TYPE_MAP.get(nt.lower())
+            if canon is None:
+                unknown[nt] += 1
+                continue
+            times.append(n["timestamp"])
+            classes.append(NOTE_CLASSES.index(canon))
+            if n.get("bpm"):
+                bpms.append(n["bpm"])
+    if not times:
+        return None
+    order = np.argsort(times, kind="stable")
+    return (
+        np.asarray(times, np.float64)[order],
+        np.asarray(classes, np.int8)[order],
+        int(course_struct.get("level") or 0),
+        float(np.median(bpms)) if bpms else 0.0,
+        unknown,
+    )
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default="cache")
+    ap.add_argument("--splits", nargs="+", default=["train", "test"])
+    args = ap.parse_args()
+
+    from datasets import Audio, load_dataset
+
+    ds_raw = load_dataset("JacobLinCool/taiko-1000-parsed")
+    # decode directly at the target sample rate (ffmpeg-side resampling);
+    # ds_raw kept as fallback for files torchcodec fails to resample
+    ds = ds_raw.cast_column("audio", Audio(sampling_rate=SR))
+    os.makedirs(args.out, exist_ok=True)
+    index = []
+    unknown_total = Counter()
+
+    for split in args.splits:
+        d = ds[split]
+        for i in range(len(d)):
+            row = d[i]
+            sid = f"{split}_{i:05d}"
+            mel_path = os.path.join(args.out, f"{sid}.mel.npy")
+            notes_path = os.path.join(args.out, f"{sid}.notes.npz")
+            entry = {
+                "id": sid,
+                "split": split,
+                "title": (row.get("metadata") or {}).get("TITLE", ""),
+                "genre": (row.get("metadata") or {}).get("GENRE", ""),
+                "courses": {},
+            }
+            arrays = {}
+            for c in COURSES:
+                ex = extract_notes(row.get(c))
+                if ex is None:
+                    continue
+                times, classes, level, bpm, unknown = ex
+                unknown_total.update(unknown)
+                arrays[f"{c}_t"] = times
+                arrays[f"{c}_c"] = classes
+                entry["courses"][c] = {
+                    "level": level,
+                    "n_notes": int(len(times)),
+                    "bpm": bpm,
+                }
+            if not arrays:
+                print(f"skip {sid}: no charts")
+                continue
+            if not (os.path.exists(mel_path) and os.path.exists(notes_path)):
+                try:
+                    mel = logmel(to_mono_22k(row["audio"]))
+                except Exception:
+                    try:  # fallback: decode at native rate, resample with librosa
+                        mel = logmel(to_mono_22k(ds_raw[split][i]["audio"]))
+                        print(f"{sid}: used raw-rate fallback", flush=True)
+                    except Exception as e:
+                        print(f"skip {sid}: audio decode failed ({e})", flush=True)
+                        continue
+                np.save(mel_path, mel)
+                np.savez(notes_path, **arrays)
+                entry["n_frames"] = int(mel.shape[1])
+            else:
+                entry["n_frames"] = int(np.load(mel_path, mmap_mode="r").shape[1])
+            index.append(entry)
+            if i % 50 == 0:
+                print(f"{split} {i}/{len(d)}", flush=True)
+
+    with open(os.path.join(args.out, "index.json"), "w") as f:
+        json.dump(index, f)
+    print("unknown note types:", dict(unknown_total))
+    print(f"done: {len(index)} songs")
+
+
+if __name__ == "__main__":
+    main()

softchart/proxies.py

@@ -0,0 +1,238 @@
+"""Automated proxies for playability, groove, and pattern structure.
+
+Every proxy is computable without human input, from (events, bpm[, mel]).
+Together with rhythm.py metrics they form the scorecard that closes the
+evaluate -> diagnose -> improve loop.
+
+Proxy design notes:
+- playability = physical feasibility on the drum (min gaps, stream length,
+  big-note crowding) + learnable structure (motif reuse, compressibility)
+- groove = alignment with the metrical grid (strong-beat rate, accent placement)
+  + audio grounding (onset-envelope energy at note times)
+Reference values come from ground-truth charts of the same course (see scorecard).
+"""
+
+import zlib
+
+import numpy as np
+
+HIT_CLASSES = ("don", "ka", "don_big", "ka_big")
+BIG_CLASSES = ("don_big", "ka_big")
+
+
+def _hits(events):
+    return sorted((t, c) for t, c in events if c in HIT_CLASSES)
+
+
+def min_gap_violation_rate(events, min_gap=0.045):
+    """Fraction of consecutive hits closer than a physically playable gap."""
+    h = _hits(events)
+    if len(h) < 2:
+        return 0.0
+    ioi = np.diff([t for t, _ in h])
+    return float(np.mean(ioi < min_gap))
+
+
+def stream_p99_nps(events, win=2.0):
+    """99th percentile of local density (notes/sec in rolling windows) — stamina proxy."""
+    t = np.array([x for x, _ in _hits(events)])
+    if len(t) < 4:
+        return 0.0
+    counts = [np.sum((t >= x) & (t < x + win)) for x in np.arange(t[0], t[-1], 0.5)]
+    return float(np.percentile(counts, 99) / win) if counts else 0.0
+
+
+def big_crowding_rate(events, bpm):
+    """Fraction of big notes followed by another hit within 0.45 beat.
+
+    Big notes are played with both hands; crowding them against the next note
+    is physically awkward and rarely done in official charts.
+    """
+    h = _hits(events)
+    bigs = [(i, t) for i, (t, c) in enumerate(h) if c in BIG_CLASSES]
+    if not bigs or not bpm or bpm <= 0:
+        return None
+    thresh = 0.45 * 60.0 / bpm
+    crowded = sum(1 for i, t in bigs if i + 1 < len(h) and h[i + 1][0] - t < thresh)
+    return crowded / len(bigs)
+
+
+def big_note_rate(events):
+    h = _hits(events)
+    if not h:
+        return 0.0
+    return sum(1 for _, c in h if c in BIG_CLASSES) / len(h)
+
+
+def _symbol_stream(events, bpm):
+    """Quantized (ioi-class, note-type) symbol stream for structure metrics."""
+    h = _hits(events)
+    if len(h) < 3 or not bpm or bpm <= 0:
+        return []
+    beat = 60.0 / bpm
+    syms = []
+    for i in range(1, len(h)):
+        ioi = h[i][0] - h[i - 1][0]
+        frac = ioi / beat
+        # classify IOI into musical classes
+        classes = [0.25, 1 / 3, 0.5, 2 / 3, 0.75, 1.0, 1.5, 2.0]
+        j = int(np.argmin([abs(frac - c) / c for c in classes]))
+        ioi_cls = j if abs(frac - classes[j]) / classes[j] < 0.25 else len(classes)
+        syms.append((ioi_cls, h[i][1]))
+    return syms
+
+
+def motif_reuse(events, bpm, n=4):
+    """Fraction of n-grams (rhythm+type) that occur at least twice in the chart.
+
+    Human charts repeat and transform motifs; random note salads do not.
+    """
+    syms = _symbol_stream(events, bpm)
+    if len(syms) < n + 2:
+        return None
+    grams = [tuple(syms[i : i + n]) for i in range(len(syms) - n + 1)]
+    from collections import Counter
+
+    counts = Counter(grams)
+    return float(sum(1 for g in grams if counts[g] >= 2) / len(grams))
+
+
+def compression_ratio(events, bpm):
+    """zlib-compressed size / raw size of the symbol stream. Lower = more structured."""
+    syms = _symbol_stream(events, bpm)
+    if len(syms) < 16:
+        return None
+    raw = bytes(b for s in syms for b in (s[0], HIT_CLASSES.index(s[1])))
+    return float(len(zlib.compress(raw, 9)) / len(raw))
+
+
+def strong_beat_rate(events, bpm, phase=None, tol=0.03):
+    """Fraction of hits on integer beats (downbeat proxy needs measures; beat is enough)."""
+    from .rhythm import estimate_phase
+
+    t = np.array([x for x, _ in _hits(events)])
+    if len(t) < 5 or not bpm or bpm <= 0:
+        return None
+    beat = 60.0 / bpm
+    if phase is None:
+        phase = estimate_phase(t, beat)
+    rel = (t - phase) / beat
+    return float(np.mean(np.abs(rel - np.round(rel)) * beat < tol))
+
+
+def accent_on_strong(events, bpm, phase=None, tol=0.03):
+    """Fraction of BIG notes that fall on integer beats — accent placement proxy."""
+    from .rhythm import estimate_phase
+
+    h = _hits(events)
+    tb = np.array([t for t, c in h if c in BIG_CLASSES])
+    if len(tb) < 3 or not bpm or bpm <= 0:
+        return None
+    beat = 60.0 / bpm
+    if phase is None:
+        phase = estimate_phase(np.array([t for t, _ in h]), beat)
+    rel = (tb - phase) / beat
+    return float(np.mean(np.abs(rel - np.round(rel)) * beat < tol))
+
+
+def audio_energy_alignment(events, mel, fps):
+    """Mean z-scored spectral-flux strength at hit times (GT-free audio grounding)."""
+    t = np.array([x for x, _ in _hits(events)])
+    if len(t) < 5:
+        return None
+    flux = np.maximum(0.0, np.diff(mel.astype(np.float32), axis=1)).sum(axis=0)
+    flux = (flux - flux.mean()) / (flux.std() + 1e-9)
+    fr = np.clip(np.round(t * fps).astype(int), 1, len(flux) - 1)
+    return float(np.mean(np.maximum.reduce([flux[fr - 1], flux[fr], flux[np.minimum(fr + 1, len(flux) - 1)]])))
+
+
+def _pair_spans(events):
+    """Pair span-start/end events into (t0, t1) tuples (for GT event streams)."""
+    out = []
+    open_t = None
+    for t, c in sorted(events):
+        if c in ("roll", "roll_big", "balloon") and open_t is None:
+            open_t = t
+        elif c == "end" and open_t is not None:
+            out.append((open_t, t))
+            open_t = None
+    return out
+
+
+def span_stats(events, spans=None):
+    """Span-length distribution + rate. Added after a visual audit found
+    generated balloons ~10s vs GT p50 ~1s (a metric blind spot)."""
+    if spans is None:
+        pairs = _pair_spans(events)
+    else:
+        pairs = [(s["t0"], s["t1"]) if isinstance(s, dict) else tuple(s) for s in spans]
+    h = _hits(events)
+    dur_min = max((h[-1][0] - h[0][0]) / 60.0, 1e-6) if len(h) > 1 else None
+    if not pairs or dur_min is None:
+        return {"span_p90_len": None, "span_per_min": 0.0 if dur_min else None}
+    lens = np.array([b - a for a, b in pairs])
+    return {"span_p90_len": float(np.percentile(lens, 90)),
+            "span_per_min": float(len(pairs) / dur_min)}
+
+
+def compute_proxies(events, bpm, mel=None, fps=None, spans=None):
+    """All proxies for one chart. Returns dict (None where not computable)."""
+    from .rhythm import grid_consistency, tuplet_evenness
+
+    t = [x for x, _ in _hits(events)]
+    out = {
+        "nps": len(t) / max(t[-1] - t[0], 1e-6) if len(t) > 1 else 0.0,
+        "min_gap_violation": min_gap_violation_rate(events),
+        "stream_p99_nps": stream_p99_nps(events),
+        "big_rate": big_note_rate(events),
+        "big_crowding": big_crowding_rate(events, bpm),
+        "motif_reuse": motif_reuse(events, bpm),
+        "compression": compression_ratio(events, bpm),
+        "strong_beat": strong_beat_rate(events, bpm),
+        "accent_on_strong": accent_on_strong(events, bpm),
+    }
+    g = grid_consistency(t, bpm)
+    if g:
+        out["offgrid_rate"] = 1.0 - g["on_grid_frac"]
+        out["grid_dev_ms"] = g["mean_dev_ms"]
+    tp = tuplet_evenness(t, bpm)
+    if tp:
+        out["tuplet_cv"] = tp["cv_mean"]
+    if mel is not None and fps:
+        out["energy_align"] = audio_energy_alignment(events, mel, fps)
+    out.update(span_stats(events, spans))
+    return out
+
+
+# direction: +1 higher is better, -1 lower is better, 0 match GT reference
+PROXY_DIRECTION = {
+    "nps": 0, "min_gap_violation": -1, "stream_p99_nps": 0, "big_rate": 0,
+    "big_crowding": -1, "motif_reuse": 0, "compression": 0, "strong_beat": 0,
+    "accent_on_strong": 0, "offgrid_rate": -1, "grid_dev_ms": -1,
+    "tuplet_cv": -1, "energy_align": 1, "span_p90_len": 0, "span_per_min": 0,
+}
+
+
+def score_against_reference(proxy_row, ref_stats):
+    """Normalized deviation of each proxy vs the GT reference distribution.
+
+    ref_stats: {proxy: {"p25": .., "p50": .., "p75": ..}} per course.
+    Returns {proxy: score} where score = robust z of |dev| (0 = at GT median;
+    for directional proxies, only penalize the bad direction).
+    """
+    scores = {}
+    for k, v in proxy_row.items():
+        if v is None or k not in ref_stats:
+            continue
+        r = ref_stats[k]
+        iqr = max(r["p75"] - r["p25"], 1e-6)
+        d = PROXY_DIRECTION.get(k, 0)
+        dev = (v - r["p50"]) / iqr
+        if d == -1:
+            dev = max(0.0, dev)  # only worse-than-GT (higher) penalized
+        elif d == 1:
+            dev = max(0.0, -dev)
+        else:
+            dev = abs(dev)
+        scores[k] = float(dev)
+    return scores

softchart/rhythm.py

@@ -0,0 +1,185 @@
+"""Rhythmic-structure utilities: adaptive rational-grid snapping and
+tuplet-evenness / grid-consistency metrics.
+
+The model emits note times on an 11.6 ms frame grid; musically, notes live at
+rational subdivisions of the beat (halves, triplets, quintuplets, ...). These
+utilities (a) snap generated notes onto a locally-selected rational grid, and
+(b) measure how "musically exact" a chart's timing is.
+"""
+
+import numpy as np
+
+# subdivisions considered per beat: binary, triple, quint, sept, and compounds
+SUBDIVS = (1, 2, 3, 4, 5, 6, 7, 8, 12, 16)
+COMPLEXITY_PENALTY = 0.002  # seconds of tolerated extra error per log2(subdiv)
+
+
+def estimate_phase(times, beat, resolution=64):
+    """Grid phase (sec) minimizing distance of notes to a fine subdivision grid."""
+    times = np.asarray(times)
+    if len(times) == 0:
+        return 0.0
+    fine = beat / 24
+    cands = np.arange(0.0, beat, beat / resolution)
+    errs = []
+    for o in cands:
+        d = (times - o) / fine
+        errs.append(np.mean(np.abs(d - np.round(d))) * fine)
+    return float(cands[int(np.argmin(errs))])
+
+
+def snap_times(times, bpm, phase=None, local_beats=16, conf_gate_ms=18.0):
+    """Snap each note to the best rational subdivision of its local beat.
+
+    Improvements over naive global snapping (validated by the closed loop):
+    - phase is re-estimated *locally* every `local_beats` beats, so tempo drift
+      and mid-song BPM changes do not accumulate into cross-beat unevenness;
+    - confidence gating: beats whose best-fit residual exceeds `conf_gate_ms`
+      are left unsnapped (better raw than snapped-to-the-wrong-grid).
+    Returns (snapped_times, info dict).
+    """
+    times = np.asarray(sorted(times), dtype=np.float64)
+    if len(times) == 0 or not bpm or bpm <= 0:
+        return times, {"phase": 0.0, "subdiv_hist": {}, "gated_frac": 0.0}
+    beat = 60.0 / bpm
+    if phase is None:
+        phase = estimate_phase(times, beat)
+
+    snapped = times.copy()
+    subdiv_hist = {}
+    n_gated = 0
+    seg_len = local_beats * beat
+    seg_idx = np.floor((times - phase) / seg_len).astype(int)
+    for s in np.unique(seg_idx):
+        seg_sel = seg_idx == s
+        seg_t = times[seg_sel]
+        # local phase refinement within this segment (bounded to +-1/8 beat)
+        local_phase = phase
+        if len(seg_t) >= 6:
+            offs = np.arange(-beat / 8, beat / 8, beat / 64)
+            fine = beat / 24
+            errs = [np.mean(np.abs(((seg_t - phase - o) / fine)
+                                   - np.round((seg_t - phase - o) / fine)) * fine)
+                    for o in offs]
+            local_phase = phase + float(offs[int(np.argmin(errs))])
+
+        beat_idx = np.floor((seg_t - local_phase) / beat).astype(int)
+        seg_snapped = seg_t.copy()
+        for k in np.unique(beat_idx):
+            sel = beat_idx == k
+            t = seg_t[sel]
+            t0 = local_phase + k * beat
+            rel = (t - t0) / beat  # in [0, 1)
+            best_d, best_err, best_raw = None, np.inf, np.inf
+            for d in SUBDIVS:
+                err = np.mean(np.abs(rel - np.round(rel * d) / d)) * beat
+                score = err + COMPLEXITY_PENALTY * np.log2(d)
+                if score < best_err:
+                    best_err, best_d, best_raw = score, d, err
+            if best_raw * 1000 > conf_gate_ms:  # low confidence: keep raw times
+                n_gated += int(sel.sum())
+                continue
+            seg_snapped[sel] = t0 + np.round(rel * best_d) / best_d * beat
+            subdiv_hist[int(best_d)] = subdiv_hist.get(int(best_d), 0) + int(sel.sum())
+        snapped[seg_sel] = seg_snapped
+    return snapped, {"phase": phase, "subdiv_hist": subdiv_hist,
+                     "gated_frac": n_gated / len(times)}
+
+
+def snap_chart(gen, bpm, min_gap=0.04):
+    """Apply rational snapping to a generate_song() output dict (hits only).
+
+    Post-snap min-gap collapse: snapping can push two notes into adjacent fine
+    slots (< playable gap); merge them instead of emitting unplayable pairs.
+    """
+    if not gen["hits"]:
+        return gen
+    times = [h["t"] for h in gen["hits"]]
+    snapped, info = snap_times(times, bpm)
+    out = dict(gen)
+    out["hits"] = []
+    last = -1.0
+    for h, t in zip(gen["hits"], snapped):
+        if t - last < min_gap:  # collapse collisions created by snapping
+            continue
+        out["hits"].append({**h, "t": round(float(t), 4)})
+        last = t
+    out["snap_info"] = info
+    return out
+
+
+def grid_consistency(times, bpm, tol=0.015):
+    """Fraction of notes within tol of the best rational grid + mean deviation.
+
+    Uses the same per-beat adaptive subdivision selection as snap_times.
+    """
+    times = np.asarray(sorted(times), dtype=np.float64)
+    if len(times) < 5 or not bpm or bpm <= 0:
+        return None
+    snapped, _ = snap_times(times, bpm)
+    dev = np.abs(snapped - times)
+    return {
+        "on_grid_frac": float(np.mean(dev <= tol)),
+        "mean_dev_ms": float(np.mean(dev) * 1000),
+        "p95_dev_ms": float(np.percentile(dev, 95) * 1000),
+    }
+
+
+def tuplet_groups(times, bpm, min_len=3):
+    """Find runs of >=min_len consecutive IOIs that form one subdivision class.
+
+    Returns list of (subdiv_class, ioi_array). subdiv_class = round(beat/median_ioi),
+    e.g. 2=8ths, 3=8th-triplets, 4=16ths, 5=quintuplets, 6=16th-triplets, 7=septuplets.
+    """
+    times = np.asarray(sorted(times), dtype=np.float64)
+    if len(times) < min_len + 1 or not bpm or bpm <= 0:
+        return []
+    beat = 60.0 / bpm
+    ioi = np.diff(times)
+    groups = []
+    i = 0
+    while i < len(ioi):
+        j = i
+        while j + 1 < len(ioi) and 0.75 < ioi[j + 1] / max(ioi[i:j + 1].mean(), 1e-9) < 1.33:
+            j += 1
+        run = ioi[i : j + 1]
+        if len(run) >= min_len:
+            med = np.median(run)
+            if med < 0.95 * beat:  # subdivision runs only
+                cls = int(round(beat / med))
+                if 2 <= cls <= 16:
+                    groups.append((cls, run))
+        i = j + 1
+    return groups
+
+
+def tuplet_evenness(times, bpm, min_len=3):
+    """Within-group spacing consistency. Returns per-class and overall stats.
+
+    CV = std(IOI)/mean(IOI) within a group (0 = perfectly even);
+    max_dev_ms = worst absolute deviation from the group's mean IOI.
+    """
+    groups = tuplet_groups(times, bpm, min_len)
+    if not groups:
+        return None
+    per_class = {}
+    cvs, maxdevs = [], []
+    for cls, run in groups:
+        cv = float(np.std(run) / np.mean(run))
+        md = float(np.max(np.abs(run - np.mean(run))) * 1000)
+        cvs.append(cv)
+        maxdevs.append(md)
+        d = per_class.setdefault(cls, {"n": 0, "cv": [], "max_dev_ms": []})
+        d["n"] += 1
+        d["cv"].append(cv)
+        d["max_dev_ms"].append(md)
+    return {
+        "n_groups": len(groups),
+        "cv_mean": float(np.mean(cvs)),
+        "max_dev_ms_mean": float(np.mean(maxdevs)),
+        "per_class": {
+            str(c): {"n": v["n"], "cv": float(np.mean(v["cv"])),
+                     "max_dev_ms": float(np.mean(v["max_dev_ms"]))}
+            for c, v in sorted(per_class.items())
+        },
+    }

softchart/train.py

@@ -0,0 +1,406 @@
+"""Training loop.
+
+    python -m softchart.train --cache /workspace/softchart/cache --out runs/base
+"""
+
+import argparse
+import json
+import math
+import os
+import time
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from .data import ChartWindowDataset, collate, load_split_ids
+from .model import ChartModel, count_params
+from .vocab import VOCAB
+
+
+# class-balanced weights for rare note types (~1/sqrt(train frequency), capped)
+TYPE_WEIGHTS = {"don": 1.0, "ka": 1.23, "don_big": 3.19, "ka_big": 5.29,
+                "roll": 7.18, "roll_big": 8.0, "balloon": 8.0, "end": 5.0}
+
+
+def build_class_weights(device, cap=8.0):
+    w = torch.ones(VOCAB.size, device=device)
+    for name, weight in TYPE_WEIGHTS.items():
+        w[VOCAB.note[name]] = min(weight, cap)
+    return w
+
+
+def loss_fn(logits, tgt, loss_mask, label_smoothing=0.1, class_weight=None):
+    # logits: (B, L-1, V) predicting tgt[:, 1:]
+    gold = tgt[:, 1:]
+    mask = loss_mask[:, 1:]
+    ls = F.cross_entropy(
+        logits.reshape(-1, logits.shape[-1]),
+        gold.reshape(-1),
+        reduction="none",
+        label_smoothing=label_smoothing,
+        weight=class_weight,
+    ).reshape(gold.shape)
+    return (ls * mask).sum() / mask.sum().clamp(min=1)
+
+
+def corrupt_events(tgt, loss_mask, rng_gen, hard=False):
+    """Synthetic negative for DPO. hard=True keeps corruptions NEAR the snap
+    tolerance (small jitter, few swaps) so negatives resemble the policy's own
+    mistakes rather than being trivially separable (fixes dpo_loss->0 collapse)."""
+    from .vocab import WINDOW
+
+    rej = tgt.clone()
+    ev = loss_mask.clone()
+    t0 = VOCAB.time0
+    is_time = (rej >= t0) & (rej < t0 + WINDOW) & ev
+    if hard:  # +-1..3 frames (~12-35ms, straddles the snap grid) on 60% of notes
+        jit = torch.randint(-3, 4, rej.shape, device=rej.device, generator=rng_gen)
+        jit = torch.where(jit == 0, torch.ones_like(jit), jit)
+        do_t = torch.rand(rej.shape, device=rej.device, generator=rng_gen) < 0.6
+        jit = torch.where(do_t, jit, torch.zeros_like(jit))
+        swap_p = 0.12
+    else:
+        jit = torch.randint(-5, 6, rej.shape, device=rej.device, generator=rng_gen)
+        jit = torch.where(jit.abs() < 2, jit.sign() * 2, jit)
+        swap_p = 0.25
+    rej = torch.where(is_time, (rej + jit).clamp(t0, t0 + WINDOW - 1), rej)
+    hit0 = VOCAB.note["don"]
+    is_hit = (rej >= hit0) & (rej < hit0 + 4) & ev
+    swap = torch.randint(0, 4, rej.shape, device=rej.device, generator=rng_gen) + hit0
+    do = torch.rand(rej.shape, device=rej.device, generator=rng_gen) < swap_p
+    rej = torch.where(is_hit & do, swap, rej)
+    return rej
+
+
+def seq_logprob(model, mel, tgt, mask):
+    import torch.nn.functional as F
+
+    logits = model(mel, tgt)
+    lp = F.log_softmax(logits.float(), dim=-1)
+    gold = tgt[:, 1:]
+    g = lp.gather(-1, gold.unsqueeze(-1)).squeeze(-1)
+    return (g * mask[:, 1:]).sum(1)
+
+
+@torch.no_grad()
+def validate(model, loader, device):
+    model.eval()
+    tot, n = 0.0, 0
+    for mel, tgt, mask, _aux, gsum, posb, _beat, _in in loader:
+        mel, tgt, mask = mel.to(device), tgt.to(device), mask.to(device)
+        gsum = gsum.to(device) if gsum is not None else None
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            logits = model(mel, tgt, gsum=gsum, pos_bucket=posb.to(device))
+            l = loss_fn(logits, tgt, mask, label_smoothing=0.0)
+        tot += l.item() * mel.shape[0]
+        n += mel.shape[0]
+    model.train()
+    return tot / max(n, 1)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--cache", required=True)
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--steps", type=int, default=60000)
+    ap.add_argument("--batch", type=int, default=32)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--warmup", type=int, default=1000)
+    ap.add_argument("--cond-drop", type=float, default=0.15)
+    ap.add_argument("--no-augment", action="store_true")
+    ap.add_argument("--balanced", action="store_true", help="class-balanced type loss")
+    ap.add_argument("--balance-cap", type=float, default=8.0,
+                    help="cap for class weights (v3: 2.5 after v2 over-corrected)")
+    ap.add_argument("--ctx", action="store_true", help="prev-window context tokens")
+    ap.add_argument("--aux", action="store_true", help="onset-heatmap auxiliary head")
+    ap.add_argument("--global-ctx", action="store_true",
+                    help="whole-song summary + window-position context (v4)")
+    ap.add_argument("--importance-sampling", action="store_true",
+                    help="bias window sampling toward informative regions (v4)")
+    ap.add_argument("--d-model", type=int, default=512)
+    ap.add_argument("--enc-layers", type=int, default=6)
+    ap.add_argument("--dec-layers", type=int, default=6)
+    ap.add_argument("--ffn", type=int, default=2048)
+    ap.add_argument("--func-time", action="store_true",
+                    help="functional (sinusoidal-basis) TIME token embeddings")
+    ap.add_argument("--distill", default=None,
+                    help="teacher checkpoint for knowledge distillation")
+    ap.add_argument("--use-aug", action="store_true",
+                    help="include waveform speed-augmented variants in training")
+    ap.add_argument("--init-from", default=None,
+                    help="initialize matching weights (decoder/tok_emb) from a checkpoint")
+    ap.add_argument("--kd-alpha", type=float, default=0.5, help="CE weight (rest = KL)")
+    ap.add_argument("--kd-tau", type=float, default=2.0, help="distillation temperature")
+    ap.add_argument("--tempo-aug", action="store_true",
+                    help="rhythm-preserving tempo/pitch augmentation in mel domain")
+    ap.add_argument("--sibling", action="store_true",
+                    help="easier-course skeleton hint in the prefix (easy subset-of hard)")
+    ap.add_argument("--style", action="store_true",
+                    help="charting-intent style token (cache/styles.json)")
+    ap.add_argument("--align", action="store_true",
+                    help="pointer alignment loss: notes must attend to their onset frame")
+    ap.add_argument("--beat-head", action="store_true",
+                    help="beat/downbeat auxiliary head (labels from TJA measures)")
+    ap.add_argument("--beat-weight", type=float, default=0.1,
+                    help="beat-head loss weight (raise for a dedicated beat model)")
+    ap.add_argument("--sync-token", action="store_true",
+                    help="LHL syncopation-band condition token")
+    ap.add_argument("--plan", action="store_true",
+                    help="song-level plan-block conditioning (cache/plans.json)")
+    ap.add_argument("--mask-infill", type=float, default=0.0,
+                    help="probability of type-mask (skeleton->color) curriculum per window")
+    ap.add_argument("--complexity", action="store_true",
+                    help="rhythmic-complexity band token (density-independent difficulty)")
+    ap.add_argument("--dpo", default=None,
+                    help="reference checkpoint: switch to DPO finetuning with synthetic negatives")
+    ap.add_argument("--dpo-beta", type=float, default=0.1)
+    ap.add_argument("--dpo-hard", action="store_true",
+                    help="near-tolerance negatives (resemble policy errors)")
+    ap.add_argument("--val-every", type=int, default=2000)
+    ap.add_argument("--workers", type=int, default=12)
+    ap.add_argument("--resume", default=None)
+    ap.add_argument("--seed", type=int, default=1234)
+    args = ap.parse_args()
+    torch.manual_seed(args.seed)
+
+    os.makedirs(args.out, exist_ok=True)
+    device = "cuda"
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    train_ids, val_ids, _ = load_split_ids(args.cache, use_aug=args.use_aug)
+    train_ds = ChartWindowDataset(
+        args.cache, train_ids, train=True, cond_drop=args.cond_drop,
+        spec_augment=not args.no_augment, use_ctx=args.ctx, aux=args.aux,
+        global_ctx=args.global_ctx, importance_sampling=args.importance_sampling,
+        tempo_aug=args.tempo_aug, sibling=args.sibling, style=args.style,
+        beat_head=args.beat_head, sync_token=args.sync_token,
+        plan=args.plan, mask_infill=args.mask_infill, complexity=args.complexity,
+    )
+    val_ds = ChartWindowDataset(args.cache, val_ids, train=False, windows_per_chart=1,
+                                use_ctx=args.ctx, aux=args.aux, global_ctx=args.global_ctx,
+                                sibling=args.sibling, style=args.style,
+                                beat_head=args.beat_head, sync_token=args.sync_token,
+                                plan=args.plan, complexity=args.complexity)
+    print(f"train charts {len(train_ds.items)}, val charts {len(val_ds.items)}")
+
+    train_loader = DataLoader(
+        train_ds, batch_size=args.batch, shuffle=True, collate_fn=collate,
+        num_workers=args.workers, pin_memory=True, drop_last=True, persistent_workers=True,
+    )
+    val_loader = DataLoader(
+        val_ds, batch_size=args.batch, shuffle=False, collate_fn=collate, num_workers=4,
+    )
+
+    model = ChartModel(
+        d_model=args.d_model, enc_layers=args.enc_layers, dec_layers=args.dec_layers,
+        ffn=args.ffn, aux=args.aux, global_ctx=args.global_ctx,
+        func_time=args.func_time,
+    ).to(device)
+    if args.align:
+        model.enable_ptr()
+    if args.beat_head:
+        model.enable_beat()
+    model.to(device)
+    print(f"params: {count_params(model)/1e6:.1f}M, vocab {VOCAB.size}")
+    if args.init_from and os.path.exists(args.init_from):
+        src = torch.load(args.init_from, map_location=device)["model"]
+        own = model.state_dict()
+        hit = {k: v for k, v in src.items()
+               if k in own and own[k].shape == v.shape
+               and k.startswith(("decoder", "tok_emb", "out", "dec_pos"))}
+        own.update(hit)
+        model.load_state_dict(own)
+        print(f"initialized {len(hit)} tensors from {args.init_from}")
+    class_weight = build_class_weights(device, args.balance_cap) if args.balanced else None
+
+    teacher = None
+    if args.distill:
+        from .generate import load_model
+
+        teacher = load_model(args.distill, device=device)
+        for p in teacher.parameters():
+            p.requires_grad_(False)
+        t_vocab = teacher.tok_emb.weight.shape[0]
+        print(f"distilling from {args.distill} (teacher vocab {t_vocab})")
+    dpo_ref = None
+    if args.dpo:
+        from .generate import load_model as _lm
+
+        dpo_ref = _lm(args.dpo, device=device)
+        for p_ in dpo_ref.parameters():
+            p_.requires_grad_(False)
+        # initialize the policy from the reference (vocab may have grown since:
+        # copy overlapping embedding rows, everything else matches the 12M arch)
+        ref_sd = dpo_ref.state_dict()
+        own = model.state_dict()
+        for k, v in ref_sd.items():
+            if k in own:
+                if own[k].shape == v.shape:
+                    own[k] = v
+                elif k == "tok_emb.weight":
+                    own[k][: v.shape[0]] = v
+        model.load_state_dict(own)
+        # disable dropout for DPO: logprob-margin noise from dropout on long
+        # sequences swamps the beta-scaled objective (verified: init loss 3.3
+        # instead of log 2)
+        import torch.nn as _nn
+
+        for m_ in model.modules():
+            if isinstance(m_, _nn.Dropout):
+                m_.p = 0.0
+        print(f"DPO: policy initialized from {args.dpo}, beta={args.dpo_beta}, dropout off")
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01, betas=(0.9, 0.98))
+
+    def lr_at(step):
+        if step < args.warmup:
+            return args.lr * step / args.warmup
+        p = (step - args.warmup) / max(1, args.steps - args.warmup)
+        return args.lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * p)))
+
+    step = 0
+    best_val = float("inf")
+    if args.resume and os.path.exists(args.resume):
+        ck = torch.load(args.resume, map_location=device)
+        model.load_state_dict(ck["model"])
+        opt.load_state_dict(ck["opt"])
+        step = ck["step"]
+        best_val = ck.get("best_val", best_val)
+        print(f"resumed from {args.resume} at step {step}")
+
+    log_path = os.path.join(args.out, "train_log.jsonl")
+    model.train()
+    t0 = time.time()
+    run_loss, run_n = 0.0, 0
+    done = False
+    while not done:
+        for mel, tgt, mask, aux_t, gsum, posb, beat_t, in_tgt in train_loader:
+            step += 1
+            if step > args.steps:
+                done = True
+                break
+            for g in opt.param_groups:
+                g["lr"] = lr_at(step)
+            mel, tgt, mask = mel.to(device, non_blocking=True), tgt.to(device), mask.to(device)
+            gsum = gsum.to(device) if gsum is not None else None
+            posb = posb.to(device)
+            in_tgt = in_tgt.to(device)
+            if dpo_ref is not None:
+                gen_rng = torch.Generator(device=device)
+                gen_rng.manual_seed(args.seed * 100003 + step)
+                rej = corrupt_events(tgt, mask, gen_rng, hard=args.dpo_hard)
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    pc = seq_logprob(model, mel, tgt, mask)
+                    pr = seq_logprob(model, mel, rej, mask)
+                    with torch.no_grad():
+                        rc = seq_logprob(dpo_ref, mel, tgt, mask)
+                        rr = seq_logprob(dpo_ref, mel, rej, mask)
+                    loss = -F.logsigmoid(
+                        args.dpo_beta * ((pc - pr) - (rc - rr))).mean()
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                opt.step()
+                run_loss += loss.item()
+                run_n += 1
+                if step % 100 == 0:
+                    msg = {"step": step, "dpo_loss": round(run_loss / run_n, 4)}
+                    print(json.dumps(msg), flush=True)
+                    with open(log_path, "a") as f:
+                        f.write(json.dumps(msg) + "\n")
+                    run_loss, run_n = 0.0, 0
+                if step % args.val_every == 0 or step == args.steps:
+                    ck = {"model": model.state_dict(), "opt": opt.state_dict(),
+                          "step": step, "best_val": 0.0, "args": vars(args)}
+                    torch.save(ck, os.path.join(args.out, "last.pt"))
+                    torch.save(ck, os.path.join(args.out, "best.pt"))
+                continue
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                if args.align or args.beat_head:
+                    logits, aux_logits, ptr_logits, beat_logits = model(
+                        mel, tgt, return_aux=True, return_extras=True,
+                        gsum=gsum, pos_bucket=posb, in_tgt=in_tgt)
+                elif args.aux and aux_t is not None:
+                    logits, aux_logits = model(mel, tgt, return_aux=True,
+                                               gsum=gsum, pos_bucket=posb, in_tgt=in_tgt)
+                    ptr_logits = beat_logits = None
+                else:
+                    logits = model(mel, tgt, gsum=gsum, pos_bucket=posb, in_tgt=in_tgt)
+                    aux_logits = ptr_logits = beat_logits = None
+                aux_loss = 0.0
+                if args.aux and aux_t is not None and aux_logits is not None:
+                    aux_loss = F.binary_cross_entropy_with_logits(
+                        aux_logits.float(), aux_t.to(device))
+                extra = 0.0
+                if ptr_logits is not None:
+                    # pointer alignment: positions whose GOLD token is TIME must
+                    # point at that frame's encoder position
+                    gold = tgt[:, 1:]
+                    is_t = (gold >= VOCAB.time0) & (gold < VOCAB.time0 + 1728) & mask[:, 1:]
+                    if is_t.any():
+                        target_pos = ((gold - VOCAB.time0) // 4).clamp(0, ptr_logits.shape[-1] - 1)
+                        pl = F.cross_entropy(
+                            ptr_logits.reshape(-1, ptr_logits.shape[-1]).float(),
+                            target_pos.reshape(-1), reduction="none",
+                        ).reshape(gold.shape)
+                        extra = extra + 0.2 * (pl * is_t).sum() / is_t.sum()
+                if beat_logits is not None and beat_t is not None:
+                    extra = extra + args.beat_weight * F.binary_cross_entropy_with_logits(
+                        beat_logits.float(), beat_t.to(device))
+                loss = loss_fn(logits, tgt, mask, class_weight=class_weight) + 0.1 * aux_loss + extra
+                if teacher is not None:
+                    with torch.no_grad():
+                        # teacher may have a smaller vocab (no SIB/STYLE tokens):
+                        # clamp unseen ids to UNK and distill over the shared slice
+                        t_tgt = torch.where(tgt >= t_vocab,
+                                            torch.full_like(tgt, VOCAB.unk_cond), tgt)
+                        t_logits = teacher(mel, t_tgt)
+                    V = min(t_logits.shape[-1], logits.shape[-1])
+                    tau = args.kd_tau
+                    kl = F.kl_div(
+                        F.log_softmax(logits[..., :V].float() / tau, dim=-1),
+                        F.softmax(t_logits[..., :V].float() / tau, dim=-1),
+                        reduction="none",
+                    ).sum(-1)
+                    kl = (kl * mask[:, 1:]).sum() / mask[:, 1:].sum().clamp(min=1)
+                    loss = args.kd_alpha * loss + (1 - args.kd_alpha) * (tau * tau) * kl
+            opt.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+            run_loss += loss.item()
+            run_n += 1
+
+            if step % 100 == 0:
+                msg = {
+                    "step": step,
+                    "loss": round(run_loss / run_n, 4),
+                    "lr": round(lr_at(step), 6),
+                    "sec": round(time.time() - t0, 1),
+                }
+                print(json.dumps(msg), flush=True)
+                with open(log_path, "a") as f:
+                    f.write(json.dumps(msg) + "\n")
+                run_loss, run_n = 0.0, 0
+
+            if step % args.val_every == 0 or step == args.steps:
+                vl = validate(model, val_loader, device)
+                msg = {"step": step, "val_loss": round(vl, 4)}
+                print(json.dumps(msg), flush=True)
+                with open(log_path, "a") as f:
+                    f.write(json.dumps(msg) + "\n")
+                ck = {"model": model.state_dict(), "opt": opt.state_dict(),
+                      "step": step, "best_val": best_val, "args": vars(args)}
+                torch.save(ck, os.path.join(args.out, "last.pt"))
+                if vl < best_val:
+                    best_val = vl
+                    ck["best_val"] = best_val
+                    torch.save(ck, os.path.join(args.out, "best.pt"))
+
+    print(f"done. best_val={best_val:.4f}")
+
+
+if __name__ == "__main__":
+    main()

softchart/vocab.py

@@ -0,0 +1,189 @@
+"""Token vocabulary for chart sequences.
+
+Sequence layout (decoder):
+    [BOS] [COURSE_c] [LEVEL_l] [DENS_d] ([TIME_t] [NOTE_e])* [EOS]
+Conditions may be UNK (condition dropout / unspecified at inference).
+"""
+
+from dataclasses import dataclass
+
+# --- audio / windowing constants (single source of truth) ---
+SR = 22050
+N_FFT = 2048
+HOP = 256
+N_MELS = 128
+FPS = SR / HOP  # 86.1328125 frames/sec
+WINDOW = 1728  # frames per training window (~20.06 s)
+MAX_TGT = 768  # max decoder length (prefix + events + eos)
+
+COURSES = ["easy", "normal", "hard", "oni", "ura"]
+N_LEVELS = 12  # levels clamped to 1..12
+N_DENS = 16
+DENS_BUCKET_NPS = 0.75  # bucket width in notes/sec
+
+# canonical note classes; parsed TJA note_type strings are mapped onto these
+NOTE_CLASSES = ["don", "ka", "don_big", "ka_big", "roll", "roll_big", "balloon", "end"]
+
+# mapping from strings observed in the dataset -> canonical class
+# (observed in taiko-1000-parsed: Don, Ka, DonBig, KaBig, Roll, RollBig,
+#  Balloon, BalloonAlt, EndOf)
+NOTE_TYPE_MAP = {
+    "Don": "don",
+    "Ka": "ka",
+    "DonBig": "don_big",
+    "KaBig": "ka_big",
+    "Roll": "roll",
+    "RollBig": "roll_big",
+    "Balloon": "balloon",
+    "BalloonAlt": "balloon",  # kusudama, treated as balloon-class span
+    "EndOf": "end",
+}
+
+
+@dataclass
+class Vocab:
+    pad: int = 0
+    eos: int = 1
+    bos: int = 2
+    unk_cond: int = 3
+
+    def __post_init__(self):
+        base = 4
+        self.course = {c: base + i for i, c in enumerate(COURSES)}
+        base += len(COURSES)
+        self.level = {l: base + l - 1 for l in range(1, N_LEVELS + 1)}
+        base += N_LEVELS
+        self.dens = {d: base + d for d in range(N_DENS)}
+        base += N_DENS
+        self.note = {n: base + i for i, n in enumerate(NOTE_CLASSES)}
+        base += len(NOTE_CLASSES)
+        self.time0 = base
+        base += WINDOW
+        # v2+ tokens appended at the end so older checkpoints stay loadable
+        self.sep = base  # separates prev-window context from conditions
+        self.sib = base + 1  # marks the sibling-chart (easier course) segment
+        self.style = {s: base + 2 + s for s in range(8)}  # charting-intent codes
+        self.sync = {s: base + 10 + s for s in range(6)}  # LHL syncopation bands
+        # plan-realize tokens: [PLAN] marker + block density (8) + block flags (3)
+        self.plan = base + 16
+        self.pdens = {d: base + 17 + d for d in range(8)}
+        self.pflag = {f: base + 25 + f for f in range(3)}  # 0 none / 1 gap / 2 climax
+        self.mask = base + 28  # type-infill placeholder
+        self.cplx = {c: base + 29 + c for c in range(6)}  # rhythmic-complexity band
+        self.size = base + 35
+
+        self.id2note = {v: k for k, v in self.note.items()}
+        self.id2course = {v: k for k, v in self.course.items()}
+
+    def time(self, frame: int) -> int:
+        assert 0 <= frame < WINDOW
+        return self.time0 + frame
+
+    def is_time(self, tok: int) -> bool:
+        return self.time0 <= tok < self.time0 + WINDOW
+
+    def is_note(self, tok: int) -> bool:
+        return tok in self.id2note
+
+    def dens_bucket(self, nps: float) -> int:
+        return min(N_DENS - 1, max(0, int(nps / DENS_BUCKET_NPS)))
+
+
+VOCAB = Vocab()
+
+
+SIB_EVENTS = 12  # fixed number of sibling-chart events in the prefix (24 tokens)
+
+
+PLAN_SLOTS = 5  # fixed number of plan blocks in the prefix (2 tokens each)
+
+
+def complexity_band(frames):
+    """IOI-class entropy of a window's hit sequence -> band 0-5 (rhythmic
+    complexity independent of raw density)."""
+    import math
+    hits = sorted(f for f, c in frames if NOTE_CLASSES[c] in ("don","ka","don_big","ka_big"))
+    if len(hits) < 4:
+        return 0
+    iois = [b - a for a, b in zip(hits, hits[1:]) if b > a]
+    if not iois:
+        return 0
+    from collections import Counter
+    cls = Counter(int(round(math.log2(max(i,1)) * 2)) for i in iois)
+    tot = sum(cls.values())
+    ent = -sum((n/tot) * math.log2(n/tot) for n in cls.values())
+    return min(5, int(ent / 0.5))
+
+
+def encode_window(vocab, course, level, notes, cond_drop=0.0, rng=None, ctx_types=None,
+                  sib_pairs=None, style=None, sync_band=None, plan_slice=None,
+                  complexity=None):
+    """Build a token sequence for one window.
+
+    notes: list of (frame_idx, note_class_id) sorted by frame, frame in [0, WINDOW).
+    ctx_types: optional list of note-class ids from the previous window's tail
+        (pattern continuity context, v2). Encoded as [BOS] ctx.. [SEP] conds..
+    sib_pairs: optional list of (frame, class_id) events from an easier course of
+        the same song (skeleton hint, easy⊂hard). Fixed SIB_EVENTS slots,
+        missing slots filled with UNK. Encoded after [SIB].
+    Returns (tokens, prefix_len) where loss should be applied after the prefix.
+    """
+    n_hits = sum(1 for _, c in notes if NOTE_CLASSES[c] not in ("end",))
+    nps = n_hits / (WINDOW / FPS)
+    d = vocab.dens_bucket(nps)
+
+    def maybe(tok):
+        if cond_drop > 0 and rng is not None and rng.random() < cond_drop:
+            return vocab.unk_cond
+        return tok
+
+    lvl = max(1, min(N_LEVELS, level if level and level > 0 else 1))
+    seq = [vocab.bos]
+    if ctx_types is not None:  # fixed-length, left-padded with UNK (id < 0 = pad)
+        seq += [vocab.note[NOTE_CLASSES[c]] if c >= 0 else vocab.unk_cond
+                for c in ctx_types]
+        seq.append(vocab.sep)
+    if sib_pairs is not None:
+        seq.append(vocab.sib)
+        pairs = list(sib_pairs)[:SIB_EVENTS]
+        for f, c in pairs:
+            seq += [vocab.time(int(f)), vocab.note[NOTE_CLASSES[c]]]
+        seq += [vocab.unk_cond] * (2 * (SIB_EVENTS - len(pairs)))
+    seq += [
+        maybe(vocab.course[course]),
+        maybe(vocab.level[lvl]),
+        maybe(vocab.dens[d]),
+    ]
+    if style is not None:  # charting-intent code (style >= 0; -1 = unknown)
+        seq.append(maybe(vocab.style[style]) if style >= 0 else vocab.unk_cond)
+    if sync_band is not None:  # LHL syncopation band (groove-intensity control)
+        seq.append(maybe(vocab.sync[sync_band]) if sync_band >= 0 else vocab.unk_cond)
+    if complexity is not None:  # rhythmic-complexity band (difficulty beyond density)
+        seq.append(maybe(vocab.cplx[complexity]) if complexity >= 0 else vocab.unk_cond)
+    if plan_slice is not None:  # song-level plan blocks overlapping this window
+        seq.append(vocab.plan)
+        blocks = list(plan_slice)[:PLAN_SLOTS]
+        for d8, fl in blocks:
+            seq += [maybe(vocab.pdens[min(7, max(0, d8))]),
+                    maybe(vocab.pflag[min(2, max(0, fl))])]
+        seq += [vocab.unk_cond] * (2 * (PLAN_SLOTS - len(blocks)))
+    prefix_len = len(seq)
+    for f, c in notes:
+        seq.append(vocab.time(int(f)))
+        seq.append(vocab.note[NOTE_CLASSES[c]])
+    seq.append(vocab.eos)
+    return seq, prefix_len
+
+
+def decode_tokens(vocab, tokens):
+    """Token ids -> list of (frame_idx, note_class_name). Ignores malformed pairs."""
+    out = []
+    cur_t = None
+    for tok in tokens:
+        if vocab.is_time(tok):
+            cur_t = tok - vocab.time0
+        elif tok in vocab.id2note and cur_t is not None:
+            out.append((cur_t, vocab.id2note[tok]))
+        elif tok == vocab.eos:
+            break
+    return out