legacy_softchart/generate.py

"""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}