Restore selectable legacy model generations

README.md

@@ -11,11 +11,19 @@ pinned: false
 license: mit
 models:
   - JacobLinCool/softchart-v18
+  - JacobLinCool/softchart-v17
+  - JacobLinCool/softchart-v17-small
+  - JacobLinCool/softchart-v17-tiny
+  - JacobLinCool/softchart-v16
+  - JacobLinCool/softchart-v16-small
+  - JacobLinCool/softchart-v16-tiny
+  - JacobLinCool/softchart-v15
+  - JacobLinCool/softchart-planner
 datasets:
   - JacobLinCool/taiko-1000-parsed-clean
 ---
 
-# 🥁 SoftChart V1.8
+# 🥁 SoftChart
 
 Generate a **Taiko no Tatsujin-style** chart from an uploaded song with the
 scratch-trained [SoftChart V1.8 hierarchical model](https://huggingface.co/JacobLinCool/softchart-v18).
@@ -23,8 +31,15 @@ The 8,985,091-parameter model combines whole-song hierarchical context, a local
 rhythmic-skeleton auxiliary head, chart generation, and beat/downbeat
 estimation. It uses no external pretrained model and no external song planner.
 
-The Space always uses V1.8's full-song **time generation** path for arbitrary
-uploads. The beat head can estimate BPM and a downbeat grid for TJA
+V1.8 is the default recommendation, not the only option. Users can select
+V1.5, V1.6, or V1.7—including the small and tiny scaling variants—from the
+interface for direct comparison. Those releases run through their frozen
+generation runtime and original
+[legacy planner](https://huggingface.co/JacobLinCool/softchart-planner); this
+does not change V1.8's self-contained, scratch-only architecture.
+
+The Space always uses each selected model's full-song **time generation** path
+for arbitrary uploads. The beat head can estimate BPM and a downbeat grid for TJA
 quantization, but that estimated grid is not treated as the trusted rational
 meter required by exact slot generation. This distinction prevents an
 estimated beat grid from being presented as authored-meter exact timing.

app.py

@@ -1,9 +1,9 @@
-"""SoftChart V1.8 — hierarchical Taiko chart generation for HF Spaces.
+"""SoftChart — selectable Taiko chart model generations for HF Spaces.
 
-The single scratch-trained model owns whole-song hierarchy, local chart
-generation, and beat/downbeat estimation. Arbitrary uploads always use the
-model's time path; an estimated beat grid may quantize the exported TJA but is
-never promoted to the trusted-meter slot path.
+V1.8 is the recommended scratch-trained hierarchical model. V1.5–V1.7 remain
+available through their frozen runtime so users can compare generations.
+Arbitrary uploads always use time generation; an estimated beat grid may
+quantize the exported TJA but is never promoted to a trusted-meter slot path.
 """
 
 import logging
@@ -20,7 +20,17 @@ import torch
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 
-from softchart.generate import generate_song, load_hf
+from legacy_softchart.generate import (
+    generate_song as generate_song_legacy,
+    load_hf as load_hf_legacy,
+)
+from legacy_softchart.grid import (
+    fit_grid_fixed_bpm as fit_grid_fixed_bpm_legacy,
+    fit_grid_piecewise as fit_grid_piecewise_legacy,
+)
+from legacy_softchart.hf import SoftChartPlanner
+from legacy_softchart.tja import append_measure_with_gogo, gogo_measure_mask
+from softchart.generate import generate_song as generate_song_v18, load_hf as load_hf_v18
 from softchart.fonts import cjk_font_path
 from softchart.grid import debias_to_grid, fit_grid_fixed_bpm, fit_grid_piecewise
 from softchart.preview_audio import synthesize_taiko_preview
@@ -31,54 +41,134 @@ from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR
 LOGGER = logging.getLogger("softchart.space")
 STATIC_DIR = Path(__file__).with_name("static")
 
-MODEL_REPO = "JacobLinCool/softchart-v18"
+PLAN_REPO = "JacobLinCool/softchart-planner"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 EXPECTED_PARAMETERS = 8_985_091
 
+MODEL_SPECS = {
+    "v1.8": {
+        "repo": "JacobLinCool/softchart-v18",
+        "display": "V1.8 hierarchical",
+        "runtime": "v18",
+    },
+    "v1.7": {
+        "repo": "JacobLinCool/softchart-v17",
+        "display": "V1.7",
+        "runtime": "legacy",
+    },
+    "v1.7-small": {
+        "repo": "JacobLinCool/softchart-v17-small",
+        "display": "V1.7 small",
+        "runtime": "legacy",
+    },
+    "v1.7-tiny": {
+        "repo": "JacobLinCool/softchart-v17-tiny",
+        "display": "V1.7 tiny",
+        "runtime": "legacy",
+    },
+    "v1.6": {
+        "repo": "JacobLinCool/softchart-v16",
+        "display": "V1.6",
+        "runtime": "legacy",
+    },
+    "v1.6-small": {
+        "repo": "JacobLinCool/softchart-v16-small",
+        "display": "V1.6 small",
+        "runtime": "legacy",
+    },
+    "v1.6-tiny": {
+        "repo": "JacobLinCool/softchart-v16-tiny",
+        "display": "V1.6 tiny",
+        "runtime": "legacy",
+    },
+    "v1.5": {
+        "repo": "JacobLinCool/softchart-v15",
+        "display": "V1.5",
+        "runtime": "legacy",
+    },
+}
+DEFAULT_MODEL = "v1.8"
+
 COURSE_DENS = {"easy": 1, "normal": 2, "hard": 4, "oni": 7}
 CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4",
         "roll": "5", "roll_big": "6", "balloon": "7"}
 SUB = 96
 
-_MODEL = {}
+_MODELS = {}
+_PLANNER = {}
 _MEL_FB = None
 _STFT_WINDOW = None
 
 
-def get_model():
-    if "generator" not in _MODEL:
-        model = load_hf(MODEL_REPO, device=DEVICE)
+def get_model(model_choice):
+    spec = MODEL_SPECS.get(model_choice)
+    if spec is None:
+        raise ValueError(
+            f"Unknown model {model_choice!r}. Choose one of: "
+            f"{', '.join(MODEL_SPECS)}."
+        )
+    if model_choice not in _MODELS:
+        if spec["runtime"] == "v18":
+            model = load_hf_v18(spec["repo"], device=DEVICE)
+        else:
+            model = load_hf_legacy(spec["repo"], device=DEVICE)
         capabilities = getattr(model, "_softchart_capabilities", {})
-        required = {
-            "aux": True,
-            "beat_head": True,
-            "dual": True,
-            "hierarchical_ctx": True,
-            "global_ctx": False,
-            "plan": False,
-        }
-        mismatched = {
-            name: capabilities.get(name)
-            for name, expected in required.items()
-            if capabilities.get(name) is not expected
-        }
-        if (not getattr(model, "_dual", False) or model.beat is None
-                or not getattr(model, "_hierarchical_ctx", False)
-                or getattr(model, "_has_plan", False) or mismatched):
-            raise RuntimeError(
-                f"{MODEL_REPO} is not the expected hierarchical V1.8 artifact"
-            )
-        parameters = sum(parameter.numel() for parameter in model.parameters())
-        if parameters != EXPECTED_PARAMETERS:
-            raise RuntimeError(
-                f"{MODEL_REPO} has {parameters:,} parameters; expected "
-                f"{EXPECTED_PARAMETERS:,}"
-            )
-        _MODEL["generator"] = model
-    return _MODEL["generator"]
+        if spec["runtime"] == "v18":
+            required = {
+                "aux": True,
+                "beat_head": True,
+                "dual": True,
+                "hierarchical_ctx": True,
+                "global_ctx": False,
+                "plan": False,
+            }
+            mismatched = {
+                name: capabilities.get(name)
+                for name, expected in required.items()
+                if capabilities.get(name) is not expected
+            }
+            parameters = sum(parameter.numel() for parameter in model.parameters())
+            if (not getattr(model, "_hierarchical_ctx", False)
+                    or getattr(model, "_has_plan", False) or mismatched
+                    or parameters != EXPECTED_PARAMETERS):
+                raise RuntimeError(
+                    f"{spec['repo']} is not the expected hierarchical V1.8 artifact"
+                )
+            bundle = {
+                "model": model,
+                "generate": generate_song_v18,
+                "fit_grid": fit_grid_piecewise,
+                "fit_grid_fixed": fit_grid_fixed_bpm,
+                "preprocess": load_logmel_v18,
+                "runtime": "v18",
+            }
+        else:
+            if (not getattr(model, "_dual", False) or model.beat is None
+                    or not getattr(model, "_has_plan", False)):
+                raise RuntimeError(
+                    f"{spec['repo']} is not a supported V1.5–V1.7 artifact"
+                )
+            bundle = {
+                "model": model,
+                "generate": generate_song_legacy,
+                "fit_grid": fit_grid_piecewise_legacy,
+                "fit_grid_fixed": fit_grid_fixed_bpm_legacy,
+                "preprocess": load_logmel_legacy,
+                "runtime": "legacy",
+            }
+        _MODELS[model_choice] = bundle
+    return _MODELS[model_choice]
+
+
+def get_legacy_planner():
+    if "planner" not in _PLANNER:
+        _PLANNER["planner"] = (
+            SoftChartPlanner.from_pretrained(PLAN_REPO).to(DEVICE).eval()
+        )
+    return _PLANNER["planner"]
 
 
-def load_logmel(path):
+def load_logmel_v18(path):
     """Apply the exact V1.8 FFmpeg + torch.stft preprocessing contract."""
     import librosa
     global _MEL_FB, _STFT_WINDOW
@@ -119,6 +209,58 @@ def load_logmel(path):
     return mel, wav
 
 
+def load_logmel_legacy(path):
+    """Preserve the preprocessing contract used by V1.5–V1.7."""
+    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 learned_plan(planner, mel, course, bpm, downbeats=None):
+    duration = mel.shape[1] / FPS
+    beat = 60.0 / bpm
+    edges = (
+        list(downbeats[::4]) + [duration]
+        if downbeats is not None and len(downbeats) >= 2
+        else list(np.arange(0, duration, 4 * beat)) + [duration]
+    )
+    features = []
+    spans = []
+    for start, end in zip(edges, edges[1:]):
+        segment = mel[:, int(start * FPS):int(end * FPS)]
+        if segment.shape[1] < 2:
+            continue
+        flux = np.maximum(0, np.diff(segment, axis=1)).sum(0)
+        features.append(np.concatenate([
+            segment.mean(1), segment.std(1),
+            [flux.mean(), flux.std(), flux.max()],
+        ]))
+        spans.append((round(float(start), 3), round(float(end), 3)))
+    if not features:
+        return None
+    course_id = {"easy": 0, "normal": 1, "hard": 2, "oni": 3}[course]
+    inputs = torch.tensor(np.array(features), dtype=torch.float32)[None].to(DEVICE)
+    with torch.no_grad():
+        density_logits, flag_logits = planner(
+            inputs, torch.tensor([course_id], device=DEVICE)
+        )
+    densities = density_logits[0].argmax(-1).cpu().numpy()
+    flags = flag_logits[0].argmax(-1).cpu().numpy()
+    return [
+        [start, end, int(density), int(flag)]
+        for (start, end), density, flag in zip(spans, densities, flags)
+    ]
+
+
 def group_quantize(times, phase, grid, min_run=3):
     n = len(times)
     slots = [0] * n
@@ -156,7 +298,8 @@ def output_tja_path(wave_name, course, directory):
     return os.path.join(directory, f"{stem}_{course}.tja")
 
 
-def write_tja(gen, bpm, title, course, level, wave, downbeats=None, grid_fit=None):
+def write_tja(gen, bpm, title, course, level, wave, downbeats=None, grid_fit=None,
+              plan=None):
     hits = sorted((h["t"], CHAR[h["type"]]) for h in gen["hits"])
     beat = 60.0 / bpm
     grid = beat / (SUB / 4)
@@ -208,10 +351,22 @@ def write_tja(gen, bpm, title, course, level, wave, downbeats=None, grid_fit=Non
             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)
-    ]
+    measure_starts = phase + np.arange(n_meas + 1, dtype=float) * (4 * beat)
+    gogo_mask = gogo_measure_mask(plan, measure_starts, n_meas)
+    lines = []
+    in_gogo = False
+    for measure in range(n_meas):
+        in_gogo = append_measure_with_gogo(
+            lines,
+            "".join(
+                slots.get(measure * SUB + slot, "0") for slot in range(SUB)
+            ) + ",",
+            measure,
+            gogo_mask,
+            in_gogo,
+        )
+    if in_gogo:
+        lines.append("#GOGOEND")
     balloons = [10] * sum(1 for s in gen["spans"] if s["type"] == "balloon")
     return "\n".join([
         f"TITLE:{title} (SoftChart)", f"BPM:{bpm:g}", f"WAVE:{wave}",
@@ -220,7 +375,7 @@ def write_tja(gen, bpm, title, course, level, wave, downbeats=None, grid_fit=Non
         "", "#START", *lines, "#END"]) + "\n"
 
 
-def render_song_structure(mel, title, course, out_path):
+def render_song_structure(mel, title, course, out_path, plan=None):
     import matplotlib
     matplotlib.use("Agg")
     from matplotlib import font_manager
@@ -233,14 +388,41 @@ def render_song_structure(mel, title, course, out_path):
     matplotlib.rcParams["axes.unicode_minus"] = False
 
     dur = mel.shape[1] / FPS
-    fig, ax0 = plt.subplots(figsize=(13, 3.2))
+    fig = plt.figure(figsize=(13, 4.4 if plan else 3.2))
+    grid = fig.add_gridspec(
+        2 if plan else 1, 1,
+        height_ratios=[3.0, 1.0] if plan else [1],
+        hspace=0.14 if plan else 0.0,
+    )
+    ax0 = fig.add_subplot(grid[0])
     ax0.imshow(mel, aspect="auto", origin="lower",
                cmap="magma", extent=[0, dur, 0, N_MELS])
     ax0.set_ylabel("mel")
     ax0.set_title(f"{title} — {course}  |  full-song mel spectrogram")
     ax0.set_xlim(0, dur)
     ax0.grid(axis="x", alpha=0.18)
-    ax0.set_xlabel("time (s)")
+    if plan:
+        ax0.set_xticklabels([])
+        ax1 = fig.add_subplot(grid[1], sharex=ax0)
+        for start, end, density, flag in plan:
+            color = "#d64545" if flag == 2 else (
+                "#4a90d9" if flag == 1 else "#999999"
+            )
+            ax1.bar(
+                (start + end) / 2,
+                max(density, 0.15),
+                width=max((end - start) * 0.92, 0.01),
+                color=color,
+                alpha=0.85,
+            )
+        ax1.set_xlim(0, dur)
+        ax1.set_ylim(0, 8)
+        ax1.set_yticks([0, 4, 8])
+        ax1.set_ylabel("plan", fontsize=8)
+        ax1.set_xlabel("time (s) — legacy plan: density / gap / climax")
+        ax1.grid(axis="x", alpha=0.18)
+    else:
+        ax0.set_xlabel("time (s)")
     fig.savefig(out_path, dpi=130, bbox_inches="tight")
     plt.close(fig)
     return out_path
@@ -310,10 +492,11 @@ async def homepage():
 async def health():
     return {
         "status": "ok",
-        "version": "1.8",
-        "model": MODEL_REPO,
+        "default_model": DEFAULT_MODEL,
+        "available_models": list(MODEL_SPECS),
         "device": DEVICE,
-        "model_loaded": "generator" in _MODEL,
+        "models_loaded": sorted(_MODELS),
+        "planner_loaded": "planner" in _PLANNER,
     }
 
 
@@ -336,6 +519,7 @@ def generate_chart(
     temperature: float,
     top_p: float,
     drum_volume: float,
+    model_choice: str = DEFAULT_MODEL,
 ) -> dict[str, object]:
     """Stream the actual inference stages to the custom frontend."""
     workdir = tempfile.mkdtemp(prefix="softchart-request-")
@@ -347,18 +531,25 @@ def generate_chart(
         temperature = float(temperature)
         top_p = float(top_p)
         drum_volume = float(drum_volume)
+        spec = MODEL_SPECS.get(model_choice)
+        if spec is None:
+            raise ValueError(
+                f"Unknown model {model_choice!r}. Choose one of: "
+                f"{', '.join(MODEL_SPECS)}."
+            )
 
         yield _progress(
             "loading", 0.03, "Preparing",
-            "Loading SoftChart V1.8 hierarchical.",
+            f"Loading SoftChart {spec['display']}.",
         )
-        model = get_model()
+        bundle = get_model(model_choice)
+        model = bundle["model"]
 
         yield _progress(
             "audio", 0.11, "Listening",
             "Mapping rhythm, melody, and timbre.",
         )
-        mel, wav = load_logmel(audio_path)
+        mel, wav = bundle["preprocess"](audio_path)
 
         yield _progress(
             "beat", 0.21, "Finding the beat",
@@ -366,7 +557,7 @@ def generate_chart(
         )
         grid = dbs = None
         if use_beat:
-            grid = fit_grid_piecewise(model, mel, device=DEVICE)
+            grid = bundle["fit_grid"](model, mel, device=DEVICE)
             if grid is not None:
                 dbs = grid["downbeats"] if grid["ok"] else grid["db_peaks"]
                 if not grid["ok"]:
@@ -374,7 +565,9 @@ def generate_chart(
         if bpm_override > 0:
             bpm = bpm_override
             if use_beat and (grid is None or abs(grid["bpm"] - bpm) > 0.5):
-                fixed_grid = fit_grid_fixed_bpm(model, mel, bpm, device=DEVICE)
+                fixed_grid = bundle["fit_grid_fixed"](
+                    model, mel, bpm, device=DEVICE
+                )
                 grid = fixed_grid if fixed_grid is not None and fixed_grid["ok"] else None
                 if grid is not None:
                     dbs = grid["downbeats"]
@@ -397,9 +590,18 @@ def generate_chart(
         if grid is None and abs(bpm - round(bpm)) < 0.06:
             bpm = float(round(bpm))
 
+        plan = None
+        if bundle["runtime"] == "v18":
+            structure_title = "Reading the structure"
+            structure_detail = "Building whole-song hierarchical context."
+        else:
+            structure_title = "Shaping the arc"
+            structure_detail = "Applying the original legacy song planner."
+            plan = learned_plan(
+                get_legacy_planner(), mel, course, bpm, downbeats=dbs
+            )
         yield _progress(
-            "structure", 0.34, "Reading the structure",
-            "Building whole-song hierarchical context.",
+            "structure", 0.34, structure_title, structure_detail,
         )
 
         yield _progress(
@@ -407,15 +609,16 @@ def generate_chart(
             "Writing playable Taiko patterns.",
         )
         title = os.path.splitext(wave_name)[0]
-        generated = generate_song(
+        generated = bundle["generate"](
             model, mel, course, level=level,
             density_bucket=COURSE_DENS[course], greedy=not sampling,
             temperature=temperature, top_p=top_p, seed=0, device=DEVICE,
+            plan=plan,
         )
         generated = snap_chart(generated, bpm)
         tja = write_tja(
             generated, bpm, title, course, level, wave_name, dbs,
-            grid_fit=grid,
+            grid_fit=grid, plan=plan,
         )
 
         yield _progress(
@@ -428,7 +631,9 @@ def generate_chart(
         chart_image_path = os.path.join(workdir, "chart.png")
         structure_image_path = os.path.join(workdir, "song-structure.png")
         render_tja_image(tja, out_path=chart_image_path)
-        render_song_structure(mel, title, course, out_path=structure_image_path)
+        render_song_structure(
+            mel, title, course, out_path=structure_image_path, plan=plan,
+        )
 
         yield _progress(
             "mix", 0.91, "Mixing",
@@ -446,10 +651,15 @@ def generate_chart(
 
         grid_rms = round(float(grid["rms_ms"]), 1) if grid is not None else None
         metrics = {
+            "model": model_choice,
             "bpm": round(float(bpm), 1),
             "notes": len(generated["hits"]),
             "spans": len(generated["spans"]),
-            "timing": "hierarchical time + grid quantization",
+            "timing": (
+                "hierarchical time + grid quantization"
+                if bundle["runtime"] == "v18"
+                else "legacy time + grid quantization"
+            ),
             "grid_rms_ms": grid_rms,
             "preview_hits": int(mix_stats["rendered_hit_count"]),
         }

legacy_softchart/generate.py

@@ -0,0 +1,572 @@
+"""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 {}
+    front_ch = sd["frontend.0.weight"].shape[1]  # 128 (clean) or 130 (legacy slot)
+    emb_w = sd["tok_emb.weight"].shape[1]
+    d_model = a.get("d_model", 512)
+    # clean_phase (v1.6): encoder frontend is 128ch and a phase_proj is present.
+    clean_phase = bool(a.get("clean_phase", False)) or ("phase_proj.0.weight" in sd)
+    # legacy slot models fed phase to the encoder (in_ch=130); clean models keep
+    # in_ch=130 at the data level but only 128 reach the frontend.
+    in_ch = 130 if (clean_phase or front_ch == 130) else front_ch
+    # v1.7 arch flags: prefer saved args; fall back to state-dict keys so
+    # args-less checkpoints of the new arch still load
+    adapter_rank_ffn = int(a.get("adapter_rank_ffn", 0) or 0)
+    if not adapter_rank_ffn and "decoder.ffn_adapters.0.down.weight" in sd:
+        adapter_rank_ffn = sd["decoder.ffn_adapters.0.down.weight"].shape[0]
+    depth_emb = bool(a.get("depth_emb", False)) or ("decoder.depth_emb" in sd)
+    model = ChartModel(
+        d_model=d_model, nhead=a.get("nhead", 8),
+        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), in_ch=in_ch,
+        emb_factor=emb_w if emb_w != d_model else None,
+        clean_phase=clean_phase,
+        enc_share=a.get("enc_share") or None,
+        dec_share=a.get("dec_share") or None,
+        adapter_rank=int(a.get("adapter_rank", 0) or 0),
+        unique_layernorm=bool(a.get("unique_layernorm", False)),
+        adapter_rank_ffn=adapter_rank_ffn, depth_emb=depth_emb,
+        unshare_last_dec=bool(a.get("unshare_last_dec", 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(hires=any(k.startswith("beat.up") for k in sd))
+    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))
+        model._dual = bool(a.get("dual", False))
+        model._slot = bool(a.get("slot", False)) or model._dual
+    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._slot = False
+        model._dual = 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._dual = bool(caps.get("dual", False))
+    m._slot = bool(caps.get("slot", False)) or m._dual
+    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, mode=None):
+    v = VOCAB
+    seq = [v.bos]
+    if mode is not None:  # dual-mode output-semantics switch
+        seq.append(v.mode[mode])
+    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, min_gap=MIN_GAP_FRAMES):
+    """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)
+    mels_dev = mels.to(device)
+    with _autocast(device):
+        memory = model.encode(
+            mels_dev,
+            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,
+        )
+        # clean_phase (v1.6): inject the grid-phase channels into the memory the
+        # DECODER reads. No-op for legacy models (in_ch already carried phase).
+        memory = model.phase_mem(memory, mels_dev)
+
+    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, 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).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) & (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, mode=None,
+                 on_chunk=None):
+    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, mode=mode,
+                         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))
+        if on_chunk:
+            on_chunk(min(i + batch_windows, len(wins)), len(wins))
+    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, on_progress=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))
+    dual = getattr(model, "_dual", False)
+    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)))
+        if dual:  # gridless time mode: phase channels = -1
+            w = torch.cat([w, torch.full((2, WINDOW), -1.0)], dim=0)
+        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,
+        mode=("time" if dual else None), on_chunk=on_progress)
+
+    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,
+            mode=("time" if dual else None))
+
+    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,
+                mode=("time" if dual else None))
+            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}
+
+
+@torch.no_grad()
+def generate_song_slot(model, mel, grid, course, level=None, density_bucket=None,
+                       device="cuda", greedy=False, temperature=1.0, top_p=0.95,
+                       seed=0, batch_windows=8, plan=None, on_progress=None):
+    """Slot-mode generation: windows are anchored at the fitted barlines and the
+    decoder emits exact TJA lattice indices (measure*96 + slot). No
+    quantization step exists — 'hits_slots' ARE the chart.
+
+    mel: (n_mels, T); grid: fit_grid() result (must be trustworthy).
+    Returns {hits, spans (seconds, for rendering/metrics),
+             hits_slots [(measure, slot, cls)], spans_slots, n_measures}.
+    """
+    from .vocab import MEAS_MAX, SLOTS
+
+    if isinstance(mel, np.ndarray):
+        mel = torch.from_numpy(mel.astype(np.float32))
+    T = mel.shape[1]
+    dur = T / FPS
+    db = np.asarray(grid["downbeats"], np.float64)
+    bar = float(grid["bar"])
+    edges = np.append(db, db[-1] + bar)  # measure m spans [edges[m], edges[m+1])
+
+    wins, metas = [], []  # meta = (first_measure_idx, K)
+    j = 0
+    while j < len(db):
+        K = 0
+        while (j + K < len(db) and K < MEAS_MAX
+               and (edges[j + K + 1] - edges[j]) * FPS <= WINDOW):
+            K += 1
+        if K == 0:
+            break
+        t0 = edges[j]
+        f0 = int(round(t0 * FPS))
+        src = max(1, int(round((edges[j + K] - t0) * FPS)))
+        x = mel[:, f0 : f0 + src]
+        if x.shape[1] < src:
+            x = torch.nn.functional.pad(x, (0, src - x.shape[1]), value=float(np.log(1e-5)))
+        if x.shape[1] < WINDOW:
+            x = torch.nn.functional.pad(x, (0, WINDOW - x.shape[1]), value=float(np.log(1e-5)))
+        ph = torch.full((2, WINDOW), -1.0)
+        ef = (edges[j : j + K + 1] - t0) * FPS
+        for m in range(K):
+            a, b = ef[m], ef[m + 1]
+            i0, i1 = int(np.ceil(a - 1e-6)), min(int(np.ceil(b - 1e-6)), WINDOW)
+            if i1 <= i0:
+                continue
+            frac = (torch.arange(i0, i1, dtype=torch.float32) - a) / max(b - a, 1e-6)
+            ph[0, i0:i1] = frac
+            ph[1, i0:i1] = (frac * 4) % 1.0
+        wins.append(torch.cat([x, ph], dim=0))
+        metas.append((j, K))
+        j += K
+
+    sib_default = getattr(model, "_has_sib", False)
+    style_val = -1 if getattr(model, "_has_style", False) else None
+    plan_blocks = plan if getattr(model, "_has_plan", False) else None
+    plan_default = getattr(model, "_has_plan", False) and plan is None
+
+    events = []  # (measure_global, slot, cls) in decode order
+    for i0 in range(0, len(wins), batch_windows):
+        chunk = torch.stack(wins[i0 : i0 + batch_windows])
+        prefixes = []
+        for b_ in range(chunk.shape[0]):
+            jj, KK = metas[i0 + b_]
+            t0, t1 = edges[jj], edges[jj + KK]
+            prefixes.append(build_prefix(
+                course, level, density_bucket,
+                mode=("slot" if getattr(model, "_dual", False) else None),
+                sib_pairs=([] if sib_default else None), style=style_val,
+                plan_slice=([(b[2], b[3]) for b in plan_blocks
+                             if b[1] > t0 and b[0] < t1] if plan_blocks is not None
+                            else ([] if plan_default else None))))
+        evs = decode_windows(model, chunk, prefixes, device=device, greedy=greedy,
+                             temperature=temperature, top_p=top_p, seed=seed + i0,
+                             n_cond=4 if style_val is not None else 3, min_gap=2)
+        for b_, w_ev in enumerate(evs):
+            jj, KK = metas[i0 + b_]
+            for g, cls in w_ev:
+                if g < KK * SLOTS:
+                    events.append((jj + g // SLOTS, g % SLOTS, cls))
+        if on_progress:
+            on_progress(min(i0 + batch_windows, len(wins)), len(wins))
+
+    events.sort(key=lambda e: (e[0], e[1]))
+    def _t(me, sl):
+        return float(edges[me] + sl / SLOTS * (edges[me + 1] - edges[me])) \
+            if me < len(edges) - 1 else float(edges[-1])
+
+    hits, spans, hits_slots, spans_slots = [], [], [], []
+    open_span = None
+    last_key = None
+    for me, sl, cls in events:
+        t = _t(me, sl)
+        if cls in HIT_CLASSES:
+            if open_span is not None:
+                if t - open_span[2] > SPAN_MAX.get(open_span[3], 6.5):
+                    t1s = open_span[2] + SPAN_MAX.get(open_span[3], 6.5)
+                    spans.append({"t0": round(open_span[2], 4), "t1": round(t1s, 4),
+                                  "type": open_span[3]})
+                    spans_slots.append((open_span[0], open_span[1], me, sl, open_span[3]))
+                    open_span = None
+                else:
+                    continue
+            if (me, sl) == last_key:
+                continue
+            hits.append({"t": round(t, 4), "type": cls})
+            hits_slots.append((me, sl, cls))
+            last_key = (me, sl)
+        elif cls in SPAN_CLASSES:
+            if open_span is None:
+                open_span = (me, sl, t, cls)
+        elif cls == "end":
+            if open_span is not None and t - open_span[2] > 0.05:
+                t1s = min(t, open_span[2] + SPAN_MAX.get(open_span[3], 6.5))
+                spans.append({"t0": round(open_span[2], 4), "t1": round(t1s, 4),
+                              "type": open_span[3]})
+                spans_slots.append((open_span[0], open_span[1], me, sl, open_span[3]))
+                open_span = None
+    n_meas = (metas[-1][0] + metas[-1][1]) if metas else 0
+    return {"hits": hits, "spans": spans, "hits_slots": hits_slots,
+            "spans_slots": spans_slots, "n_measures": n_meas,
+            "course": course, "level": level, "density_bucket": density_bucket}

legacy_softchart/grid.py

@@ -0,0 +1,609 @@
+"""Precise global beat-grid estimation from the beat head's activations.
+
+Why this exists: the beat head predicts beat/downbeat activations on the
+encoder's 4-frame (~46.4 ms) grid. Raw peak picking + median(diff) is far too
+imprecise for barline anchoring — each peak carries +-23 ms bin-quantization
+noise, a single-peak anchor offsets the whole chart by 1-3 TJA slots, and a
+0.1% period error drifts ~200 ms over a three-minute song. That is exactly the
+"notes are not on the barline" failure mode.
+
+Fix: a TJA needs one rigid grid anyway (constant BPM + one OFFSET), so we
+estimate (period, phase) GLOBALLY from every peak in the song:
+
+  1. activation curves over OVERLAPPING windows (hop = WINDOW/2), using BOTH
+     channels (ch0 = beat, ch1 = downbeat) — ~5x more evidence than
+     downbeats-only;
+  2. sub-bin peak times via parabolic interpolation (46 ms bins -> ~5-10 ms);
+  3. robust periodic regression t_k ~= phase + k*period with integer-index
+     assignment and outlier rejection, initialized at bar level and refined at
+     beat level (precision grows ~1/(N*span): 100+ bars pins BPM to <0.01);
+  4. integer / half-integer BPM snap, accepted only when the residual does not
+     degrade (community charts are almost always integer BPM);
+  5. bar phase chosen by downbeat-activation voting over the meter offsets.
+
+Residual statistics are returned so callers can fall back (no barline
+anchoring) when the song does not have a constant tempo.
+"""
+
+import numpy as np
+import torch
+
+from .vocab import FPS, WINDOW
+
+BIN = 4.0 / FPS  # activation bin length in seconds (~46.4 ms)
+
+
+@torch.no_grad()
+def beat_activations(model, mel, device="cuda", hop=WINDOW // 2):
+    """Averaged beat/downbeat activation curves over overlapping windows.
+
+    mel: (n_mels, T). Returns (acc, bin_s): acc is (2, ceil(T/factor)) with
+    row 0 = beat, row 1 = downbeat; bin_s is the bin length in seconds.
+    The head resolution is auto-detected: the classic linear head runs at
+    encoder rate (4-frame bins, ~46 ms), the hi-res head at frame rate
+    (~11.6 ms bins).
+    """
+    from .generate import _autocast
+
+    if isinstance(mel, torch.Tensor):
+        mel = mel.numpy()
+    T = mel.shape[1]
+    L = WINDOW // 4
+    acc = cnt = None
+    factor = 4
+    for st in range(0, max(T - hop, 1), hop):
+        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)))
+        if w.shape[0] < model.in_ch:  # dual/slot models: blank grid channels
+            w = torch.cat([w, torch.full((model.in_ch - w.shape[0], WINDOW), -1.0)])
+        with _autocast(device):
+            mem = model.encode(w[None].to(device))
+            pr = torch.sigmoid(model.beat(mem[:, -L:]).float())[0].cpu().numpy().T
+        if acc is None:
+            factor = WINDOW // pr.shape[1]
+            nbin = (T + factor - 1) // factor
+            acc = np.zeros((2, nbin))
+            cnt = np.zeros(nbin)
+        b0 = st // factor
+        n = min(pr.shape[1], nbin - b0)  # skip bins past the real song end
+        if n > 0:
+            acc[:, b0 : b0 + n] += pr[:, :n]
+            cnt[b0 : b0 + n] += 1
+    if acc is None:
+        return np.zeros((2, 1)), factor / FPS
+    return acc / np.maximum(cnt, 1), factor / FPS
+
+
+def _refined_peaks(act, height, min_dist_s, bin_s=BIN):
+    """Peak times (seconds) with sub-bin parabolic interpolation."""
+    from scipy.signal import find_peaks
+
+    idx, _ = find_peaks(act, height=height, distance=max(2, int(min_dist_s / bin_s)))
+    out = []
+    for i in idx:
+        dx = 0.0
+        if 0 < i < len(act) - 1:
+            d = act[i - 1] - 2 * act[i] + act[i + 1]
+            if d < -1e-9:
+                dx = float(np.clip(0.5 * (act[i - 1] - act[i + 1]) / d, -0.5, 0.5))
+        # +0.5: training targets bin FLOOR(t*FPS/div), so a peak at bin i
+        # means the event sits at the bin CENTER — without this everything
+        # reads one half-bin early
+        out.append((i + dx + 0.5) * bin_s)
+    return np.asarray(out)
+
+
+def _fit_periodic(t, period, phase, reject=0.16, iters=6):
+    """Robust regression t_i ~= phase + k_i*period with unknown integer k_i.
+    Returns (period, phase, inlier_mask, rms_seconds)."""
+    t = np.asarray(t, float)
+    keep = np.ones(len(t), bool)
+    for _ in range(iters):
+        k = np.round((t - phase) / period)
+        r = t - (phase + k * period)
+        keep = np.abs(r) < reject * period
+        if keep.sum() < 4:
+            return period, phase, keep, float("inf")
+        kk, tt = k[keep], t[keep]
+        km = kk.mean()
+        denom = float(((kk - km) ** 2).sum())
+        if denom < 1e-9:
+            break
+        period = float(((kk - km) * (tt - tt.mean())).sum() / denom)
+        phase = float(tt.mean() - period * km)
+    r = t - (phase + np.round((t - phase) / period) * period)
+    rms = float(np.sqrt(np.mean(r[keep] ** 2)))
+    return period, phase, keep, rms
+
+
+def _refit_phase(t, period, phase, reject=0.22, iters=3):
+    """Phase-only refit for a FIXED period (used by the BPM-snap step)."""
+    t = np.asarray(t, float)
+    keep = np.ones(len(t), bool)
+    for _ in range(iters):
+        k = np.round((t - phase) / period)
+        r = t - (phase + k * period)
+        keep = np.abs(r) < reject * period
+        if keep.sum() < 4:
+            return phase, keep, float("inf")
+        phase += float(np.median(r[keep]))
+    r = t - (phase + np.round((t - phase) / period) * period)
+    rms = float(np.sqrt(np.mean(r[keep] ** 2)))
+    return phase, keep, rms
+
+
+@torch.no_grad()
+def fit_grid(model, mel, device="cuda", meter=None):
+    """Estimate a rigid (BPM, barline phase) grid for the whole song.
+
+    Returns a dict with bpm / beat / bar / phase (a barline time), synthesized
+    downbeats+beats covering the song, and fit diagnostics (inlier_frac,
+    rms_ms, ok). Returns None when there are not enough peaks to fit.
+    """
+    if isinstance(mel, torch.Tensor):
+        mel_np = mel.numpy()
+    else:
+        mel_np = mel
+    act, bin_s = beat_activations(model, mel_np, device)
+    dur = mel_np.shape[1] / FPS
+    t_d = _refined_peaks(act[1], height=0.35, min_dist_s=0.9, bin_s=bin_s)
+    t_b = _refined_peaks(act[0], height=0.35, min_dist_s=0.22, bin_s=bin_s)
+    if len(t_d) < 4:
+        return None
+
+    # 1) bar-level robust fit (several phase inits; first peak may be spurious)
+    p0 = float(np.median(np.diff(t_d)))
+    fits = [_fit_periodic(t_d, p0, float(ph)) for ph in t_d[:3]]
+    p_bar0, phi_d, keep_d, rms_d = min(fits, key=lambda f: f[3])
+
+    # 2) meter from the bar/beat period ratio; taiko is almost always 4/4
+    if meter is None:
+        meter = 4
+        if len(t_b) >= 8:
+            m = p_bar0 / max(float(np.median(np.diff(t_b))), 1e-6)
+            if abs(m - 3) < 0.25:
+                meter = 3
+
+    # 3) beat-level refinement with ALL peaks (downbeats are beats too)
+    t_all = np.concatenate([t_b, t_d]) if len(t_b) else t_d
+    p_beat, phi_beat, keep_a, rms_a = _fit_periodic(
+        t_all, p_bar0 / meter, phi_d, reject=0.22)
+    if not np.isfinite(rms_a):
+        p_beat, phi_beat, rms_a = p_bar0 / meter, phi_d, rms_d
+
+    # 4) integer / half-integer BPM snap + octave normalization: prefer the
+    #    common-range notation (a 120 song fit at 60 plays identically but
+    #    notates ugly), accept only if the residual holds up
+    bpm = 60.0 / p_beat
+    cands = [round(bpm), round(bpm * 2) / 2]
+    if bpm < 80:   # octave normalization: prefer common-range notation
+        cands = [round(bpm * 2), round(bpm * 2) * 1.0] + cands
+    elif bpm > 210:
+        cands = [round(bpm / 2)] + cands
+    # rank: in [80, 210] first, integers before halves, keep insertion order
+    ranked = sorted(dict.fromkeys(cands),
+                    key=lambda c: (not (80 <= c <= 210), abs(c - round(c)) > 0,
+                                   cands.index(c)))
+    best = (rms_a, bpm, p_beat, phi_beat)
+    for cand in ranked:
+        if not 40 <= cand <= 320 or not (0.4 <= cand / max(bpm, 1e-6) <= 2.5):
+            continue
+        pb = 60.0 / cand
+        ph, _, rms = _refit_phase(t_all, pb, phi_beat)
+        if rms <= best[0] + 0.0015:  # snap unless it costs >1.5 ms rms
+            best = (rms, float(cand), pb, ph)
+            break
+    rms_a, bpm, p_beat, phi_beat = best
+    p_bar = meter * p_beat
+
+    # 5) bar phase: among the `meter` beat offsets, pick the one whose barline
+    #    comb collects the most downbeat activation
+    nb = act.shape[1]
+
+    def _comb_score(ph):
+        ts = np.arange(ph if ph >= 0 else ph + p_bar, dur, p_bar)
+        ix = np.clip(np.round(ts / bin_s).astype(int), 0, nb - 1)
+        return float(act[1][ix].sum())
+
+    base = phi_beat + np.round((phi_d - phi_beat) / p_beat) * p_beat
+    offs = [(base + j * p_beat - p_bar * np.floor((base + j * p_beat) / p_bar))
+            for j in range(meter)]
+    phase = max(offs, key=_comb_score)
+
+    k_max = int(np.floor((dur - phase) / p_bar)) + 1
+    downbeats = phase + np.arange(0, max(k_max, 1)) * p_bar
+    downbeats = downbeats[(downbeats >= 0) & (downbeats < dur)]
+    beats = phase + np.arange(0, int((dur - phase) / p_beat) + 1) * p_beat
+    beats = beats[(beats >= 0) & (beats < dur)]
+
+    inlier = float(keep_a.mean()) if len(keep_a) else 0.0
+    rms_ms = rms_a * 1000.0
+    return {
+        "bpm": float(bpm), "beat": float(p_beat), "bar": float(p_bar),
+        "meter": int(meter), "phase": float(phase),
+        "downbeats": downbeats, "beats": beats,
+        "n_db_peaks": int(len(t_d)), "n_beat_peaks": int(len(t_b)),
+        "inlier_frac": inlier, "rms_ms": float(rms_ms),
+        "db_peaks": t_d,  # raw refined peaks, for diagnostics/eval
+        "ok": bool(inlier >= 0.6 and rms_ms <= 30.0 and len(t_d) >= 8),
+    }
+
+
+@torch.no_grad()
+def fit_grid_fixed_bpm(model, mel, bpm, device="cuda", meter=4):
+    """Grid fit with a USER-SUPPLIED BPM: the period is trusted, only the
+    phase is estimated (a far easier problem than the free fit — this often
+    unlocks the slot-exact path on songs whose free fit fails on octave or
+    period confusion). Returns a grid dict or None."""
+    if isinstance(mel, torch.Tensor):
+        mel = mel.numpy()
+    act, bin_s = beat_activations(model, mel, device)
+    dur = mel.shape[1] / FPS
+    p_beat = 60.0 / float(bpm)
+    p_bar = meter * p_beat
+    t_d = _refined_peaks(act[1], height=0.30, min_dist_s=0.5, bin_s=bin_s)
+    t_b = _refined_peaks(act[0], height=0.30, min_dist_s=0.22, bin_s=bin_s)
+    t_all = np.concatenate([t_b, t_d]) if len(t_b) else t_d
+    if len(t_all) < 8:
+        return None
+    # beat-level phase (dense evidence), then bar phase by activation voting
+    ph, keep, rms = _refit_phase(t_all, p_beat, float(np.median(t_all[:5])), iters=5)
+    if not np.isfinite(rms):
+        return None
+    nb = act.shape[1]
+
+    def _score(o):
+        ts = np.arange(o - np.floor(o / p_bar) * p_bar, dur, p_bar)
+        ix = np.clip(np.round(ts / bin_s).astype(int), 0, nb - 1)
+        return float(act[1][ix].sum())
+
+    phase = max((ph + j * p_beat for j in range(meter)), key=_score)
+    phase -= np.floor(phase / p_bar) * p_bar
+    downbeats = phase + np.arange(0, int((dur - phase) / p_bar) + 1) * p_bar
+    downbeats = downbeats[(downbeats >= 0) & (downbeats < dur)]
+    beats = phase + np.arange(0, int((dur - phase) / p_beat) + 1) * p_beat
+    beats = beats[(beats >= 0) & (beats < dur)]
+    inlier = float(keep.mean()) if len(keep) else 0.0
+    return {
+        "bpm": float(bpm), "beat": p_beat, "bar": p_bar, "meter": meter,
+        "phase": float(phase), "downbeats": downbeats, "beats": beats,
+        "n_db_peaks": int(len(t_d)), "n_beat_peaks": int(len(t_b)),
+        "inlier_frac": inlier, "rms_ms": float(rms * 1000),
+        "db_peaks": t_d, "fixed_bpm": True,
+        # period is user-trusted, so the gate only needs the PHASE to be sane
+        "ok": bool(inlier >= 0.5 and rms * 1000 <= 40.0),
+    }
+
+
+# --- adaptive piecewise-constant tempo fit ------------------------------------
+# The BPM census over all 1155 songs: 34% are multi-BPM; among them the main
+# tempo covers a median 91% of measures, contiguous-segment count is <=4 for
+# 82% of songs but p99 = 29 (chained ritardando). So segmentation must be
+# ADAPTIVE (change-point DP with a per-segment penalty), not a fixed K.
+
+
+def _seg_fit(t, reject=0.16):
+    """Fit t_k ~= phase + k*period over ONE segment of consecutive downbeat
+    peaks (missing peaks allowed: k advances by round(gap/period)).
+    Returns (period, phase, sse_seconds^2, n_inlier)."""
+    t = np.asarray(t, float)
+    if len(t) < 2:
+        return None
+    d = np.diff(t)
+    p0 = float(np.median(d))
+    if p0 <= 0:
+        return None
+    k = np.concatenate([[0], np.cumsum(np.maximum(1, np.round(d / p0)))])
+    for _ in range(3):
+        km, tm = k.mean(), t.mean()
+        denom = float(((k - km) ** 2).sum())
+        if denom < 1e-9:
+            return None
+        p = float(((k - km) * (t - tm)).sum() / denom)
+        if p <= 0:
+            return None
+        phi = float(tm - p * km)
+        r = t - (phi + k * p)
+        k = k + np.round(-r / p)  # re-assign indices after refinement
+    r = t - (phi + k * p)
+    keep = np.abs(r) < reject * p
+    sse = float((r[keep] ** 2).sum()) + float((~keep).sum()) * (reject * p) ** 2
+    return p, phi, sse, int(keep.sum())
+
+
+def fit_grid_piecewise(model, mel, device="cuda", dev_thresh=0.18,
+                       min_run=3, seg_rms_ms=25.0):
+    """Main-grid + spliced exceptions (v2 after the DP version over-segmented).
+
+    The BPM census says multi-BPM songs are ~91% main-tempo, so: fit the rigid
+    grid first (full-song precision, outliers rejected), then find CONTIGUOUS
+    runs of downbeat peaks that consistently deviate from it, refit each run
+    locally, and splice. Segment count stays small by construction and the
+    main grid keeps its precision. Returns None if the main fit fails.
+    """
+    if isinstance(mel, torch.Tensor):
+        mel = mel.numpy()
+    g = fit_grid(model, mel, device=device)
+    if g is None:
+        return None
+    t_d = g["db_peaks"]
+    if len(t_d) < 8:
+        return None
+    if not g["ok"]:
+        # the single rigid grid failed: true multi-tempo song. Try every
+        # coarse split point, fit both halves rigidly, keep the best pair
+        # (recursion depth 1 -> up to 2 segments; census: covers most cases).
+        best = None
+        for cut in range(6, len(t_d) - 6, 3):
+            fa = _seg_fit(t_d[:cut])
+            fb = _seg_fit(t_d[cut:])
+            if fa is None or fb is None:
+                continue
+            rms = np.sqrt((fa[2] + fb[2]) / len(t_d))
+            if best is None or rms < best[0]:
+                best = (rms, cut, fa, fb)
+        if best is not None and best[0] * 1000 <= 30.0:
+            _, cut, (pa, fia, _, na), (pb, fib, _, nb2) = best
+            dur = mel.shape[1] / FPS
+            t_mid = 0.5 * (t_d[cut - 1] + t_d[cut])
+            da = fia + np.arange(0, int(np.ceil((t_mid - fia) / pa)) + 1) * pa
+            da = da[(da >= 0) & (da < t_mid)]
+            db2 = fib + np.arange(0, int(np.ceil((dur - fib) / pb)) + 1) * pb
+            db2 = db2[(db2 >= t_mid) & (db2 < dur)]
+            downbeats = np.sort(np.concatenate([da, db2]))
+            inl = (na + nb2) / len(t_d)
+            return {**g, "downbeats": downbeats,
+                    "bpm": 240.0 / pa, "beat": 60.0 / (240.0 / pa) if pa else g["beat"],
+                    "bar": pa, "phase": float(downbeats[0]) if len(downbeats) else g["phase"],
+                    "segments": [(float(fia), 240.0 / pa), (float(t_mid), 240.0 / pb)],
+                    "piecewise": True, "n_segments": 2,
+                    "inlier_frac": float(inl), "rms_ms": float(best[0] * 1000),
+                    "ok": bool(inl >= 0.7)}
+        return {**g, "piecewise": False, "n_segments": 1,
+                "segments": [(float(g["phase"]), float(g["bpm"]))]}
+    bar = g["bar"]
+    phi = g["phase"]
+    dur = mel.shape[1] / FPS
+    r = t_d - (phi + np.round((t_d - phi) / bar) * bar)
+    bad = np.abs(r) > dev_thresh * bar
+    # contiguous deviant runs of >= min_run peaks
+    runs = []
+    i = 0
+    while i < len(t_d):
+        if bad[i]:
+            j = i
+            while j + 1 < len(t_d) and bad[j + 1]:
+                j += 1
+            if j - i + 1 >= min_run:
+                runs.append((i, j))
+            i = j + 1
+        else:
+            i += 1
+    segments = [(float(phi), float(g["bpm"]))]
+    if not runs:  # effectively constant: return the rigid fit unchanged
+        return {**g, "piecewise": False, "n_segments": 1,
+                "segments": segments}
+
+    spliced = []
+    n_bad_fixed = 0
+    for i, j in runs:
+        f = _seg_fit(t_d[i : j + 1])
+        if f is None:
+            continue
+        p_loc, phi_loc, sse, ninl = f
+        rms = np.sqrt(sse / max(j - i + 1, 1)) * 1000
+        if rms > seg_rms_ms or not (0.8 <= p_loc <= 8.0):
+            continue  # local section too messy: leave it to the main grid
+        bpm_loc = 240.0 / p_loc
+        if abs(bpm_loc - round(bpm_loc)) < 0.35:
+            p2 = 240.0 / round(bpm_loc)
+            rr = t_d[i:j+1] - (phi_loc + np.round((t_d[i:j+1] - phi_loc) / p2) * p2)
+            if float(np.sqrt(np.mean(rr ** 2))) * 1000 <= rms * 1.3 + 4:
+                p_loc = p2
+                phi_loc = phi_loc + float(np.median(rr))
+                bpm_loc = float(round(bpm_loc))
+        t_lo = t_d[i] - 0.25 * p_loc
+        t_hi = (t_d[j] + p_loc) if j + 1 >= len(t_d) else t_d[j + 1] - 0.25 * bar
+        spliced.append((t_lo, t_hi, p_loc, phi_loc, bpm_loc))
+        n_bad_fixed += j - i + 1
+
+    if not spliced:  # nothing splice-worthy: keep the rigid fit verbatim
+        return {**g, "piecewise": False, "n_segments": 1, "segments": segments}
+
+    # main-grid barlines outside spliced spans + local barlines inside
+    main_db = phi + np.arange(0, int(np.ceil((dur - phi) / bar)) + 1) * bar
+    main_db = main_db[(main_db >= 0) & (main_db < dur)]
+    keep = np.ones(len(main_db), bool)
+    downbeats = []
+    for t_lo, t_hi, p_loc, phi_loc, bpm_loc in spliced:
+        keep &= ~((main_db >= t_lo) & (main_db < t_hi))
+        loc = phi_loc + np.arange(0, int(np.ceil((t_hi - phi_loc) / p_loc)) + 1) * p_loc
+        loc = loc[(loc >= max(t_lo, 0)) & (loc < min(t_hi, dur))]
+        downbeats.append(loc)
+        segments.append((float(max(t_lo, 0.0)), float(bpm_loc)))
+    downbeats.append(main_db[keep])
+    downbeats = np.concatenate(downbeats)
+    downbeats = np.sort(downbeats)
+    dk = np.concatenate([[True], np.diff(downbeats) > 0.3])
+    downbeats = downbeats[dk]
+
+    # residual after splicing (all peaks vs nearest final barline)
+    err = np.array([np.min(np.abs(downbeats - t)) for t in t_d])
+    inlier = float(np.mean(err < dev_thresh * bar))
+    rms_ms = float(np.sqrt(np.mean(np.minimum(err, dev_thresh * bar) ** 2)) * 1000)
+    return {
+        **g,
+        "downbeats": downbeats, "segments": sorted(segments),
+        "piecewise": bool(spliced), "n_segments": 1 + len(spliced),
+        "inlier_frac": inlier, "rms_ms": rms_ms,
+        # splices only ADD local corrections to an already-vetted main grid —
+        # never downgrade the rigid fit's verdict
+        "ok": bool(g["ok"] or (inlier >= 0.75 and rms_ms <= 35.0)),
+    }
+
+
+# --- onset evidence: metrically weighted slot-comb refinement -----------------
+# Audio onsets (spectral flux) are ~4x sharper in time than the beat head's
+# 46 ms bins and 10-20x denser than downbeats. They cannot determine the
+# metrical LEVEL (which line is beat 1) — but given the beat head's fit as an
+# anchor, they lock (period, phase) far more precisely, and they can rescue
+# songs whose beat-head activations are confused but whose percussion is clean.
+# The comb is the TJA slot lattice (beat/24) with metrical-hierarchy weights:
+# occupancy concentrates on strong positions, which is what makes a 16 ms
+# lattice identifiable at all (a flat comb that fine would fit noise).
+
+_W24 = np.full(24, 0.05)
+_W24[0] = 1.0                    # beat
+_W24[12] = 0.55                  # 8th
+_W24[[6, 18]] = 0.30             # 16ths
+_W24[[8, 16]] = 0.22             # 8th triplets
+_W24[[3, 9, 15, 21]] = 0.10      # 32nds
+_W24[[4, 20]] = 0.10             # 16th triplets
+
+
+def onset_peaks(mel, max_n=3000):
+    """Spectral-flux onset times (seconds) with sub-frame refinement."""
+    from scipy.signal import find_peaks
+
+    flux = np.maximum(0, np.diff(mel.astype(np.float32), axis=1)).sum(0)
+    flux = np.concatenate([[0.0], flux])
+    med = np.median(flux)
+    idx, props = find_peaks(flux, height=med * 1.5, distance=max(2, int(0.035 * FPS)))
+    if len(idx) > max_n:  # keep the strongest
+        keep = np.argsort(props["peak_heights"])[-max_n:]
+        idx = np.sort(idx[keep])
+    out = []
+    for i in idx:
+        dx = 0.0
+        if 0 < i < len(flux) - 1:
+            d = flux[i - 1] - 2 * flux[i] + flux[i + 1]
+            if d < -1e-9:
+                dx = float(np.clip(0.5 * (flux[i - 1] - flux[i + 1]) / d, -0.5, 0.5))
+        out.append((i + dx) / FPS)
+    return np.asarray(out)
+
+
+def comb_score(onsets, period_beat, phase, sig=0.012):
+    """Sum over onsets of hierarchy-weight x Gaussian(dist to nearest slot)."""
+    step = period_beat / 24.0
+    x = (onsets - phase) / step
+    k = np.round(x)
+    dist = (x - k) * step
+    w = _W24[(k.astype(int)) % 24]
+    return float(np.sum(w * np.exp(-0.5 * (dist / sig) ** 2)))
+
+
+def onset_polish(grid, mel, span_bpm=0.6, span_ms=30, n_p=25, n_f=31):
+    """Local (period, phase) refinement of a trusted fit on onset evidence.
+    Returns an updated grid dict (or the original if no improvement)."""
+    ons = onset_peaks(mel)
+    if len(ons) < 30 or grid is None:
+        return grid
+    p0, f0 = grid["beat"], grid["phase"]
+    base = comb_score(ons, p0, f0)
+    best = (base, p0, f0)
+    for p in np.linspace(60.0 / (grid["bpm"] + span_bpm), 60.0 / (grid["bpm"] - span_bpm), n_p):
+        for f in f0 + np.linspace(-span_ms / 1000, span_ms / 1000, n_f):
+            s = comb_score(ons, p, f)
+            if s > best[0]:
+                best = (s, float(p), float(f))
+    s1, p1, f1 = best
+    if s1 <= base * 1.02:  # <2% gain: keep the beat-head fit
+        return grid
+    out = dict(grid)
+    meter = grid.get("meter", 4)
+    out["beat"], out["bpm"], out["bar"] = p1, 60.0 / p1, meter * p1
+    # keep the bar phase on the SAME barline, re-expressed on the new beat grid
+    out["phase"] = f1 + round((grid["phase"] - f1) / p1) * p1
+    dur = mel.shape[1] / FPS
+    ph, bar = out["phase"], out["bar"]
+    ph -= np.floor(ph / bar) * bar
+    out["downbeats"] = np.arange(ph, dur, bar)
+    out["beats"] = np.arange(ph - np.floor(ph / p1) * p1, dur, p1)
+    out["onset_gain"] = s1 / max(comb_score(ons, p0, f0), 1e-9)
+    return out
+
+
+def onset_rescue(mel, act=None, bpm_range=(65, 210)):
+    """Grid estimation from onsets alone (+ optional downbeat activation for
+    the bar phase) — for songs where the beat-head fit failed. Returns a grid
+    dict with ok=True only when the comb evidence is decisive."""
+    ons = onset_peaks(mel)
+    if len(ons) < 40:
+        return None
+    dur = mel.shape[1] / FPS
+    # tempo candidates from the IOI histogram (mode + octaves)
+    iois = np.diff(ons)
+    iois = iois[(iois > 0.08) & (iois < 2.0)]
+    if len(iois) < 20:
+        return None
+    hist, edges = np.histogram(iois, bins=192, range=(0.08, 2.0))
+    cand_p = []
+    for i in np.argsort(hist)[-6:]:
+        c = 0.5 * (edges[i] + edges[i + 1])
+        for mult in (0.5, 1.0, 2.0, 4.0):
+            p = c * mult
+            if 60.0 / bpm_range[1] <= p <= 60.0 / bpm_range[0]:
+                cand_p.append(p)
+    best = (0.0, None, None)
+    for p in sorted(set(np.round(cand_p, 4))):
+        for f in np.arange(0.0, p, p / 48):
+            s = comb_score(ons, p, f, sig=0.015)
+            if s > best[0]:
+                best = (s, float(p), float(f))
+    if best[1] is None:
+        return None
+    # fine polish around the winner
+    g0 = {"beat": best[1], "bpm": 60.0 / best[1], "phase": best[2], "meter": 4,
+          "bar": 4 * best[1]}
+    # integer-BPM snap when it costs <2% of the comb score
+    bpm = g0["bpm"]
+    for candb in (round(bpm), round(bpm * 2) / 2):
+        if abs(candb - bpm) < 0.6 and candb > 0:
+            pb = 60.0 / candb
+            fb = max(np.arange(0.0, pb, pb / 96),
+                     key=lambda f: comb_score(ons, pb, f))
+            if comb_score(ons, pb, fb) >= 0.98 * best[0]:
+                g0.update(beat=pb, bpm=float(candb), phase=float(fb), bar=4 * pb)
+                break
+    # bar phase: downbeat-activation voting among the 4 beat offsets
+    ph = g0["phase"]
+    if act is not None:
+        if isinstance(act, tuple):
+            act, _bin = act
+        else:
+            _bin = BIN
+        nb = act.shape[1]
+
+        def _sc(o):
+            ts = np.arange(o - np.floor(o / g0["bar"]) * g0["bar"], dur, g0["bar"])
+            ix = np.clip(np.round(ts / _bin).astype(int), 0, nb - 1)
+            return float(act[1][ix].sum())
+
+        ph = max((g0["phase"] + jj * g0["beat"] for jj in range(4)), key=_sc)
+    ph -= np.floor(ph / g0["bar"]) * g0["bar"]
+    # decisiveness gate: winner must clearly beat the runner-up octave
+    alt = comb_score(ons, g0["beat"] * 2, ph) + comb_score(ons, g0["beat"] / 2, ph)
+    score = comb_score(ons, g0["beat"], ph)
+    ok = bool(score > 0.25 * len(ons) and score > 0.75 * alt)
+    return {"bpm": g0["bpm"], "beat": g0["beat"], "bar": g0["bar"], "meter": 4,
+            "phase": float(ph),
+            "downbeats": np.arange(ph, dur, g0["bar"]),
+            "beats": np.arange(ph - np.floor(ph / g0["beat"]) * g0["beat"], dur, g0["beat"]),
+            "n_db_peaks": 0, "n_beat_peaks": int(len(ons)),
+            "inlier_frac": float(score / max(len(ons), 1)), "rms_ms": -1.0,
+            "db_peaks": np.asarray([]), "ok": ok, "rescued": True}
+
+
+def debias_to_grid(times, phase, step):
+    """Remove the generator's systematic latency: median signed offset of
+    `times` to the nearest grid subdivision. Returns (shifted_times, offset).
+    Relative timing (groove, tuplets) is untouched — this is a global shift."""
+    t = np.asarray(times, float)
+    if len(t) == 0:
+        return t, 0.0
+    r = (t - phase + step / 2) % step - step / 2
+    off = float(np.median(r))
+    return t - off, off

legacy_softchart/hf.py

@@ -0,0 +1,104 @@
+"""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, in_ch=None, emb_factor=None,
+                 capabilities=None, nhead=8, clean_phase=False,
+                 enc_share=None, dec_share=None, adapter_rank=0,
+                 unique_layernorm=False, adapter_rank_ffn=0, depth_emb=False,
+                 unshare_last_dec=False):
+        super().__init__()
+        # nhead and the v1.6/v1.7 sharing flags (enc_share/dec_share/adapter_rank/
+        # unique_layernorm/adapter_rank_ffn/depth_emb/unshare_last_dec) and
+        # clean_phase must be persisted in config.json: from_pretrained rebuilds
+        # the architecture from these kwargs alone, before any weights are loaded.
+        # Defaults reproduce the pre-v1.6 unshared architecture for older repos.
+        self.net = ChartModel(
+            d_model=d_model, nhead=nhead, 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, in_ch=in_ch, emb_factor=emb_factor,
+            clean_phase=clean_phase, enc_share=enc_share, dec_share=dec_share,
+            adapter_rank=adapter_rank, unique_layernorm=unique_layernorm,
+            adapter_rank_ffn=adapter_rank_ffn, depth_emb=depth_emb,
+            unshare_last_dec=unshare_last_dec,
+        )
+        # 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)

legacy_softchart/model.py

@@ -0,0 +1,440 @@
+"""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 BeatHiRes(nn.Module):
+    """Frame-resolution beat/downbeat head: encoder states (T/4) are
+    upsampled 4x so supervision and peaks live on the 11.6 ms frame grid
+    instead of 46 ms encoder bins (4x finer timing per peak)."""
+
+    def __init__(self, d):
+        super().__init__()
+        self.up = nn.ConvTranspose1d(d, d // 2, 4, stride=4)
+        self.out = nn.Conv1d(d // 2, 2, 3, padding=1)
+
+    def forward(self, mem):  # (B, L, d) -> (B, 4L, 2)
+        h = nn.functional.gelu(self.up(mem.transpose(1, 2)))
+        return self.out(h).transpose(1, 2)
+
+
+def _parse_pattern(spec, n):
+    """Layer-sharing schedule. spec is either an int (number of UNIQUE
+    physical layers, cycled to length n) or a comma string like "E0,E1,E0,E1"
+    (labels; each distinct label is one physical layer). Returns a list of n
+    physical-layer indices, e.g. [0,1,0,1] for 2-unique/4-physical-depth."""
+    if spec is None:
+        return list(range(n))  # no sharing: every layer unique
+    if isinstance(spec, int):
+        return [k % spec for k in range(n)]
+    labels = [s.strip() for s in str(spec).split(",") if s.strip()]
+    assert len(labels) == n, f"pattern {spec} has {len(labels)} labels, need {n}"
+    order = {}
+    out = []
+    for lb in labels:
+        if lb not in order:
+            order[lb] = len(order)
+        out.append(order[lb])
+    return out
+
+
+class Adapter(nn.Module):
+    """Rank-r residual adapter applied after a (possibly reused) transformer
+    block: x + gate * up(gelu(down(x))). gate init 0 so the whole stack starts
+    identical to hard weight sharing; the adapter only differentiates a reused
+    physical layer as training moves the gate off zero."""
+
+    def __init__(self, d, rank):
+        super().__init__()
+        self.down = nn.Linear(d, rank, bias=False)
+        self.up = nn.Linear(rank, d, bias=False)
+        nn.init.normal_(self.down.weight, std=0.02)
+        nn.init.zeros_(self.up.weight)
+        self.gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        return x + self.gate * self.up(nn.functional.gelu(self.down(x)))
+
+
+class SharedEncoder(nn.Module):
+    """Encoder whose physical layers are reused according to a sharing pattern.
+    Each POSITION in the depth keeps its own (a) rank-r adapter and (b) final
+    LayerNorm — so ΔW starts at 0 (hard sharing) but every position can drift.
+    With pattern=None this is exactly n unique layers + a norm, matching the
+    stock nn.TransformerEncoder param count (adapters/pos-norms are opt-in).
+
+    v1.7 depth-diversity options (all default-off, init == v1.6 behaviour):
+      depth_emb: per-depth learned d-vector ADDED to the block input at each
+        reuse (params = depth x d, zero-init so init == hard sharing).
+      adapter_rank_ffn > 0: SPLIT adapters — instead of one post-block adapter,
+        each depth owns a rank-`adapter_rank` adapter on the attention sublayer
+        output and a rank-`adapter_rank_ffn` adapter on the FFN sublayer output
+        (patterning lives in the FFN). Gate init 0 keeps init == hard sharing.
+    """
+
+    def __init__(self, d, nhead, ffn, dropout, pattern, adapter_rank=0,
+                 unique_layernorm=False, adapter_rank_ffn=0, depth_emb=False):
+        super().__init__()
+        # `pattern` is the resolved list of physical-layer indices per depth
+        n_phys = max(pattern) + 1
+        self.layers = nn.ModuleList([
+            nn.TransformerEncoderLayer(d, nhead, ffn, dropout, activation="gelu",
+                                       batch_first=True, norm_first=True)
+            for _ in range(n_phys)])
+        self.plan = pattern
+        self.split = adapter_rank_ffn > 0
+        if self.split:  # v1.7: per-sublayer adapters (attention rank-a, FFN rank-f)
+            self.attn_adapters = nn.ModuleList([
+                Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
+                for _ in pattern])
+            self.ffn_adapters = nn.ModuleList([
+                Adapter(d, adapter_rank_ffn) for _ in pattern])
+        else:  # v1.6: one residual adapter after the whole block
+            self.adapters = nn.ModuleList([
+                Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
+                for _ in pattern])
+        self.depth_emb = (nn.Parameter(torch.zeros(len(pattern), d))
+                          if depth_emb else None)
+        self.pos_norms = nn.ModuleList([
+            nn.LayerNorm(d) if unique_layernorm else nn.Identity()
+            for _ in pattern])
+        self.norm = nn.LayerNorm(d)
+
+    def forward(self, x, src_key_padding_mask=None):
+        for k, phys in enumerate(self.plan):
+            if self.depth_emb is not None:
+                x = x + self.depth_emb[k]
+            layer = self.layers[phys]
+            if self.split:
+                # norm_first decomposition of nn.TransformerEncoderLayer with
+                # per-depth adapters applied to each sublayer OUTPUT (Houlsby
+                # placement); at gate=0 this is bit-equal to the stock layer
+                x = x + self.attn_adapters[k](
+                    layer._sa_block(layer.norm1(x), None, src_key_padding_mask))
+                x = x + self.ffn_adapters[k](layer._ff_block(layer.norm2(x)))
+            else:
+                x = layer(x, src_key_padding_mask=src_key_padding_mask)
+                x = self.adapters[k](x)
+            x = self.pos_norms[k](x)
+        return self.norm(x)
+
+
+class SharedDecoder(nn.Module):
+    """Decoder counterpart of SharedEncoder (see that docstring). Split mode
+    puts rank-`adapter_rank` adapters on BOTH attention sublayers (self and
+    cross) and rank-`adapter_rank_ffn` on the FFN sublayer."""
+
+    def __init__(self, d, nhead, ffn, dropout, pattern, adapter_rank=0,
+                 unique_layernorm=False, adapter_rank_ffn=0, depth_emb=False):
+        super().__init__()
+        n_phys = max(pattern) + 1
+        self.layers = nn.ModuleList([
+            nn.TransformerDecoderLayer(d, nhead, ffn, dropout, activation="gelu",
+                                       batch_first=True, norm_first=True)
+            for _ in range(n_phys)])
+        self.plan = pattern
+        self.split = adapter_rank_ffn > 0
+        if self.split:
+            self.sa_adapters = nn.ModuleList([
+                Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
+                for _ in pattern])
+            self.ca_adapters = nn.ModuleList([
+                Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
+                for _ in pattern])
+            self.ffn_adapters = nn.ModuleList([
+                Adapter(d, adapter_rank_ffn) for _ in pattern])
+        else:
+            self.adapters = nn.ModuleList([
+                Adapter(d, adapter_rank) if adapter_rank else nn.Identity()
+                for _ in pattern])
+        self.depth_emb = (nn.Parameter(torch.zeros(len(pattern), d))
+                          if depth_emb else None)
+        self.pos_norms = nn.ModuleList([
+            nn.LayerNorm(d) if unique_layernorm else nn.Identity()
+            for _ in pattern])
+        self.norm = nn.LayerNorm(d)
+
+    def forward(self, tgt, memory, tgt_mask=None, tgt_key_padding_mask=None,
+                tgt_is_causal=False):
+        x = tgt
+        for k, phys in enumerate(self.plan):
+            if self.depth_emb is not None:
+                x = x + self.depth_emb[k]
+            layer = self.layers[phys]
+            if self.split:
+                x = x + self.sa_adapters[k](layer._sa_block(
+                    layer.norm1(x), tgt_mask, tgt_key_padding_mask, tgt_is_causal))
+                x = x + self.ca_adapters[k](layer._mha_block(
+                    layer.norm2(x), memory, None, None))
+                x = x + self.ffn_adapters[k](layer._ff_block(layer.norm3(x)))
+            else:
+                x = layer(x, memory, tgt_mask=tgt_mask,
+                          tgt_key_padding_mask=tgt_key_padding_mask,
+                          tgt_is_causal=tgt_is_causal)
+                x = self.adapters[k](x)
+            x = self.pos_norms[k](x)
+        return self.norm(x)
+
+
+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,
+        in_ch=None,
+        emb_factor=None,
+        clean_phase=False,
+        enc_share=None,
+        dec_share=None,
+        adapter_rank=0,
+        unique_layernorm=False,
+        adapter_rank_ffn=0,
+        depth_emb=False,
+        unshare_last_dec=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
+        # clean_phase (v1.6): the encoder never sees the 2 grid-phase channels
+        # (in_ch stays N_MELS), and phase is injected into the DECODER memory
+        # only, via a small projection. This removes the condition leak that
+        # forced beat supervision to be masked on slot windows in v1.5.
+        self.clean_phase = clean_phase
+        self.in_ch = in_ch or N_MELS  # legacy slot mode appends 2 grid-phase channels
+        front_ch = N_MELS if clean_phase else self.in_ch
+
+        # conv frontend: (B, front_ch, T) -> (B, T/4, d)
+        self.frontend = nn.Sequential(
+            nn.Conv1d(front_ch, 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)
+
+        # phase injection (clean_phase only): 2 grid-phase channels at frame
+        # rate -> encoder rate (T/4) -> additive conditioning on the memory the
+        # DECODER reads. The beat head reads the pre-injection (clean) memory.
+        # phase_proj input = [measure_phase, beat_phase, valid_flag]. valid=0 on
+        # time-mode windows (their raw phase is -1), giving the decoder a clean
+        # zero-phase conditioning there; slot windows pass valid=1 + real phase.
+        if clean_phase:
+            self.phase_proj = nn.Sequential(
+                nn.Conv1d(3, d_model, 3, stride=2, padding=1),
+                nn.GELU(),
+                nn.Conv1d(d_model, d_model, 3, stride=2, padding=1),
+            )
+        else:
+            self.phase_proj = None
+
+        # sharing schedule (v1.6-small/tiny): reuse physical layers per pattern.
+        # enc_share/dec_share: int (#unique layers) or "E0,E1,E0,E1" label string.
+        self.enc_share = enc_share
+        self.dec_share = dec_share
+        if (enc_share is None and dec_share is None and not adapter_rank
+                and not adapter_rank_ffn and not depth_emb and not unshare_last_dec):
+            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))
+        else:
+            enc_plan = _parse_pattern(enc_share, enc_layers)
+            dec_plan = _parse_pattern(dec_share, dec_layers)
+            # v1.7: un-share the LAST decoder layer (pattern-shaping layer in
+            # front of the generation head gets its own physical parameters);
+            # no-op if that depth is already unique.
+            if unshare_last_dec and dec_plan.count(dec_plan[-1]) > 1:
+                dec_plan = dec_plan[:-1] + [max(dec_plan) + 1]
+            self.encoder = SharedEncoder(d_model, nhead, ffn, dropout, enc_plan,
+                                         adapter_rank, unique_layernorm,
+                                         adapter_rank_ffn, depth_emb)
+            self.decoder = SharedDecoder(d_model, nhead, ffn, dropout, dec_plan,
+                                         adapter_rank, unique_layernorm,
+                                         adapter_rank_ffn, depth_emb)
+
+        # ALBERT-style factorized embeddings (v1.5): vocab->E->d. Position
+        # tokens dominate the table (1728/1810 rows), so E=64 cuts the
+        # embedding block ~3.5x; logits stay tied via the projected matrix.
+        self.emb_factor = emb_factor
+        if emb_factor:
+            self.tok_emb = nn.Embedding(vocab_size, emb_factor, padding_idx=VOCAB.pad)
+            nn.init.normal_(self.tok_emb.weight, std=0.02)
+            self.emb_proj = nn.Linear(emb_factor, d_model, bias=False)
+            nn.init.normal_(self.emb_proj.weight, std=1.0 / (emb_factor ** 0.5) * 0.16)
+            self.out = None
+        else:
+            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.emb_proj = None
+            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.
+        # beat="hires" builds the frame-resolution head (11.6 ms bins)
+        self.beat = (BeatHiRes(d_model) if beat == "hires"
+                     else (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, C, T). clean_phase: C may be N_MELS (audio only) or N_MELS+2
+        # (audio + 2 grid-phase channels); only the audio channels reach the
+        # encoder, so the returned memory is CLEAN (phase-free) and safe for the
+        # beat head. gsum: (B, G, n_mels) whole-song summary chunks;
+        # pos_bucket: (B,) window-position bucket in [0, 16).
+        audio = mel[:, :N_MELS] if self.clean_phase else mel
+        h = self.frontend(audio).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 phase_mem(self, memory, mel):
+        """Phase-conditioned copy of `memory` for the DECODER to cross-attend.
+        Adds a projection of the grid-phase channels to the audio-frame slice
+        of memory; the beat head keeps reading the clean `memory`. No-op unless
+        clean_phase is on and phase channels are present. Time-mode windows
+        (raw phase = -1) get valid=0 and zeroed phase -> ~zero conditioning."""
+        nch = N_MELS
+        if not self.clean_phase or mel.shape[1] < nch + 2:
+            return memory
+        ph = mel[:, nch:nch + 2]  # (B, 2, T) in [0,1); -1 marks no-grid frames
+        valid = (ph[:, :1] >= 0).float()  # 1 where a grid exists, else 0
+        ph_in = torch.cat([ph.clamp(min=0.0), valid], dim=1)  # (B, 3, T)
+        p = self.phase_proj(ph_in).transpose(1, 2)  # (B, T/4, d)
+        L = p.shape[1]
+        out = memory.clone()
+        out[:, -L:] = out[:, -L:] + p  # audio frames are the last L memory slots
+        return out
+
+    def emb_matrix(self):
+        """Full embedding matrix; factorized and/or functional rows resolved."""
+        E = self.tok_emb.weight
+        if self.emb_proj is not None:
+            E = self.emb_proj(E)
+        if self.time_proj is not None:
+            t = self.time_proj(self.time_basis)  # (WINDOW, d)
+            E = torch.cat([E[: VOCAB.time0], t, E[VOCAB.time0 + WINDOW :]], dim=0)
+        return E
+
+    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, hires=False):
+        self.beat = BeatHiRes(self.d_model) if hires else 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)  # CLEAN
+        dec_mem = self.phase_mem(memory, mel)  # phase-conditioned for the decoder
+        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), dec_mem, 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/ptr/beat heads read the CLEAN memory (no phase leak)
+        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)

legacy_softchart/tja.py

@@ -0,0 +1,119 @@
+"""TJA emission for slot-mode charts: slot indices ARE the chart.
+
+No quantization happens here — generate_song_slot() already produced exact
+lattice positions (measure, slot in 0..95), so this writer just prints them.
+Preview rendering and the .tja are identical by construction. Piecewise
+grids additionally emit #BPMCHANGE lines between measures.
+"""
+
+import numpy as np
+
+from .vocab import SLOTS
+
+CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4",
+        "roll": "5", "roll_big": "6", "balloon": "7"}
+
+
+def grid_measure_starts(grid, n_measures):
+    """Return measure boundary times for a slot grid, extending if needed."""
+    db = np.asarray(grid.get("downbeats", []), dtype=float)
+    if len(db) >= n_measures + 1:
+        return db[: n_measures + 1]
+    if len(db) >= 2:
+        step = float(np.median(np.diff(db)))
+        start = float(db[0])
+    else:
+        step = 240.0 / float(grid["bpm"])
+        start = float(db[0]) if len(db) else 0.0
+    return start + np.arange(n_measures + 1, dtype=float) * step
+
+
+def gogo_measure_mask(plan, measure_starts, n_measures):
+    """Map plan climax blocks (flag == 2) to TJA measures."""
+    mask = [False] * n_measures
+    if not plan:
+        return mask
+    starts = np.asarray(measure_starts, dtype=float)
+    if len(starts) < n_measures + 1:
+        return mask
+    for block in plan:
+        if len(block) < 4 or int(block[3]) != 2:
+            continue
+        a, b = float(block[0]), float(block[1])
+        if b <= a:
+            continue
+        for m in range(n_measures):
+            if starts[m] < b and starts[m + 1] > a:
+                mask[m] = True
+    return mask
+
+
+def append_measure_with_gogo(lines, measure_line, measure_idx, gogo_mask, in_gogo):
+    """Append a measure line, opening/closing #GOGO commands at boundaries."""
+    want_gogo = bool(gogo_mask[measure_idx]) if measure_idx < len(gogo_mask) else False
+    if in_gogo and not want_gogo:
+        lines.append("#GOGOEND")
+        in_gogo = False
+    if want_gogo and not in_gogo:
+        lines.append("#GOGOSTART")
+        in_gogo = True
+    lines.append(measure_line)
+    return in_gogo
+
+
+def write_tja_slots(gen, grid, title, course, level, wave, out_path=None,
+                    balloon_count=10, plan=None):
+    """gen: generate_song_slot() result; grid: fit_grid() result (or a dict
+    with downbeats+bpm). Returns the TJA text (and writes it if out_path)."""
+    slots = {}
+    for me, sl, cls in gen["hits_slots"]:
+        slots.setdefault((me, sl), CHAR[cls])
+    for m0, s0, m1, s1, typ in gen.get("spans_slots", []):
+        a = (m0, s0)
+        while a in slots:  # span start yields to hits: next free slot
+            a = (a[0] + (a[1] + 1) // SLOTS, (a[1] + 1) % SLOTS)
+        b = (m1, s1)
+        while b in slots or b <= a:
+            b = (b[0] + (b[1] + 1) // SLOTS, (b[1] + 1) % SLOTS)
+        slots[a] = CHAR[typ]
+        slots[b] = "8"
+    n_meas = max(gen.get("n_measures", 0),
+                 (max(m for m, _ in slots) + 1) if slots else 1)
+    # piecewise-tempo grids: emit #BPMCHANGE whenever the per-measure BPM
+    # (from consecutive fitted barlines) moves; TJA measure lines are unchanged
+    db = np.asarray(grid.get("downbeats", []), float) if grid.get("piecewise") else None
+    measure_starts = grid_measure_starts(grid, n_meas)
+    gogo_mask = gogo_measure_mask(plan, measure_starts, n_meas)
+    lines = []
+    cur_bpm = float(grid["bpm"])
+    in_gogo = False
+    for m in range(n_meas):
+        if in_gogo and not gogo_mask[m]:
+            lines.append("#GOGOEND")
+            in_gogo = False
+        if db is not None and m + 1 < len(db):
+            bpm_m = 240.0 / (db[m + 1] - db[m])
+            if abs(bpm_m - round(bpm_m)) < 0.05:
+                bpm_m = float(round(bpm_m))
+            if abs(bpm_m - cur_bpm) > 0.05:
+                lines.append(f"#BPMCHANGE {bpm_m:g}")
+                cur_bpm = bpm_m
+        in_gogo = append_measure_with_gogo(
+            lines, "".join(slots.get((m, k), "0") for k in range(SLOTS)) + ",",
+            m, gogo_mask, in_gogo)
+    if in_gogo:
+        lines.append("#GOGOEND")
+    balloons = [balloon_count] * sum(1 for s in gen.get("spans_slots", [])
+                                     if s[4] == "balloon")
+    offset = float(grid["downbeats"][0]) if len(grid.get("downbeats", [])) else 0.0
+    tja = "\n".join([
+        f"TITLE:{title} (SoftChart)", f"BPM:{grid['bpm']:g}", f"WAVE:{wave}",
+        f"OFFSET:{-offset:.3f}",
+        f"COURSE:{'Oni' if course == 'oni' else str(course).capitalize()}",
+        f"LEVEL:{level}",
+        f"BALLOON:{','.join(map(str, balloons))}" if balloons else "BALLOON:",
+        "", "#START", *lines, "#END"]) + "\n"
+    if out_path:
+        with open(out_path, "w") as f:
+            f.write(tja)
+    return tja

legacy_softchart/vocab.py

@@ -0,0 +1,199 @@
+"""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 = 1536  # max decoder length (prefix + events + eos)
+
+SLOTS = 96  # TJA lattice: slots per 4/4 measure (LCM of 16ths/triplets/32nds)
+MEAS_MAX = WINDOW // SLOTS  # max whole measures per window in slot mode (18)
+SLOTS_PER_BEAT = 24  # meter-general lattice (v3): 3/4 -> 72, 5/4 -> 120 slots
+BEATS_MAX = WINDOW // SLOTS_PER_BEAT  # 72 beats of token space per window
+
+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
+        # Extended condition tokens append after the original token ids.
+        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
+        # dual-mode models: explicit output-semantics switch (v3). mode "slot"
+        # = positions are lattice indices on a supplied grid; "time" = frames.
+        self.mode = {"slot": base + 35, "time": base + 36}
+        self.size = base + 37
+
+        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, mode=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 mode is not None:  # dual-mode: explicit output-semantics token (never dropped)
+        seq.append(vocab.mode[mode])
+    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

static/app.css

@@ -594,6 +594,29 @@ audio {
   font-weight: 750;
 }
 
+.model-select {
+  width: 100%;
+  margin-top: 8px;
+  padding: 10px 12px;
+  color: var(--ink);
+  background: var(--surface-solid);
+  border: 1px solid var(--line);
+  border-radius: 10px;
+  font-size: 13px;
+  font-weight: 600;
+  cursor: pointer;
+  transition: border-color .18s, box-shadow .18s;
+}
+
+.model-select:hover {
+  border-color: var(--line-strong);
+}
+
+.model-select:focus-visible {
+  outline: 3px solid rgba(39, 111, 145, .35);
+  outline-offset: 1px;
+}
+
 .range {
   --range-progress: 88.89%;
   width: 100%;

static/app.js

@@ -15,6 +15,7 @@ const ui = {
   fileMeta: $("#file-meta"),
   removeFile: $("#remove-file"),
   fileError: $("#file-error"),
+  model: $("#model"),
   level: $("#level"),
   levelOutput: $("#level-output"),
   bpm: $("#bpm"),
@@ -107,6 +108,7 @@ const stageDefaults = {
 };
 
 const metricLabels = {
+  model: "Model",
   bpm: "BPM",
   notes: "Notes",
   spans: "Drumroll sections",
@@ -525,7 +527,8 @@ function requestValues() {
     ui.sampling.checked,
     Number(ui.temperature.value),
     Number(ui.topP.value),
-    Number(ui.drumVolume.value)
+    Number(ui.drumVolume.value),
+    ui.model.value
   ];
 }
 

static/index.html

@@ -83,9 +83,26 @@
         <div class="form-section split-rule">
           <div class="field-group">
             <div class="label-row">
-              <span>Model</span>
-              <strong class="field-hint">V1.8 · Hierarchical · 9.0M</strong>
+              <label for="model">Model</label>
             </div>
+            <select id="model" class="model-select">
+              <optgroup label="V1.8 — hierarchical scratch (recommended)">
+                <option value="v1.8" selected>V1.8 · 9.0M — best research model</option>
+              </optgroup>
+              <optgroup label="V1.7 — original planner path">
+                <option value="v1.7">V1.7 · 7.9M — best legacy quality</option>
+                <option value="v1.7-small">V1.7 small · 3.6M</option>
+                <option value="v1.7-tiny">V1.7 tiny · 1.4M — fastest</option>
+              </optgroup>
+              <optgroup label="V1.6 — original planner path">
+                <option value="v1.6">V1.6 · 7.9M</option>
+                <option value="v1.6-small">V1.6 small · 3.6M</option>
+                <option value="v1.6-tiny">V1.6 tiny · 1.4M</option>
+              </optgroup>
+              <optgroup label="V1.5 — original planner path">
+                <option value="v1.5">V1.5 · 7.9M</option>
+              </optgroup>
+            </select>
           </div>
 
           <fieldset class="field-group">