Improve speed of sage sheet decoder with yue2. Recover from tracks that fail

extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py

@@ -41,6 +41,22 @@ class _Attention(nn.Module):
         b, _, l, _ = out.shape
         return self.out_proj(out.transpose(1, 2).reshape(b, l, -1)), (k, v)
 
+    # static-cache path: explicit casts to the projection dtype (no-ops under autocast with fp32 weights)
+    def kv(self, src):
+        src = src.to(self.k_proj.weight.dtype)
+        return self._split(self.k_proj(src)), self._split(self.v_proj(src))
+
+    def attend(self, x, k, v, mask):
+        q = self._split(self.q_proj(x.to(self.q_proj.weight.dtype))).to(k.dtype)
+        if mask is not None and q.shape[2] == 1:
+            # single-query + mask sends SDPA to the mem-efficient kernel, which tiles over every cached key (~0.3 ms); by hand it is ~20 us
+            scores = torch.matmul(q.float(), k.float().transpose(-1, -2)) * (self.head_dim**-0.5)
+            out = torch.matmul(scores.masked_fill(~mask, float("-inf")).softmax(-1).to(v.dtype), v)
+        else:
+            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        b, _, l, _ = out.shape
+        return self.out_proj(out.transpose(1, 2).reshape(b, l, -1).to(self.out_proj.weight.dtype))
+
 
 class _Layer(nn.Module):
     def __init__(self, dim: int, heads: int, ffn: int):
@@ -63,6 +79,37 @@ class _Layer(nn.Module):
         x = self.final_layer_norm(x + self.fc2(F.gelu(self.fc1(x))))
         return x, self_kv + cross_kv
 
+    def step_static(self, x, cache, i, write_idx, mask):
+        k, v = self.self_attn.kv(x)
+        cache.k[i].index_copy_(2, write_idx, k.to(cache.k.dtype))
+        cache.v[i].index_copy_(2, write_idx, v.to(cache.v.dtype))
+        x = self.self_attn_layer_norm(x + self.self_attn.attend(x, cache.k[i], cache.v[i], mask))
+        x = self.encoder_attn_layer_norm(x + self.encoder_attn.attend(x, cache.ck[i], cache.cv[i], None))
+        return self.final_layer_norm(x + self.fc2(F.gelu(self.fc1(x.to(self.fc1.weight.dtype)))))
+
+
+class StaticCache:
+    """Preallocated batch-1 KV buffers so a decode step has no growing tensors and can be CUDA-graph replayed.
+    ``pos`` (device tensor) is the number of cached tokens; the captured step reads and bumps it in place."""
+
+    def __init__(self, decoder: "ScoreDecoder", max_len: int, memory_len: int, device, dtype):
+        attn = decoder.layers[0].self_attn
+        n, h, d = len(decoder.layers), attn.heads, attn.head_dim
+        self.max_len = max_len
+        self.k = torch.zeros(n, 1, h, max_len, d, device=device, dtype=dtype)
+        self.v = torch.zeros_like(self.k)
+        self.ck = torch.zeros(n, 1, h, memory_len, d, device=device, dtype=dtype)
+        self.cv = torch.zeros_like(self.ck)
+        self.pos = torch.zeros(1, dtype=torch.long, device=device)
+        self.arange = torch.arange(max_len, device=device)
+
+    def fill_cross(self, decoder: "ScoreDecoder", memory):
+        for i, layer in enumerate(decoder.layers):
+            k, v = layer.encoder_attn.kv(memory)
+            self.ck[i].copy_(k)
+            self.cv[i].copy_(v)
+        self.pos.zero_()
+
 
 class ScoreDecoder(nn.Module):
     """Drop-in for the BartDecoder the upstream model builds: ``ScoreDecoder(bart_config, embed_tokens=...)``."""
@@ -79,6 +126,26 @@ class ScoreDecoder(nn.Module):
     def gradient_checkpointing_disable(self):
         pass
 
+    def prefill_static(self, input_ids, cache: StaticCache):
+        """[1, P] into an empty cache; returns hidden [1, P, dim]."""
+        p = input_ids.shape[1]
+        idx = torch.arange(p, device=input_ids.device)
+        x = self.layernorm_embedding(self.embed_tokens(input_ids) + self.embed_positions(idx + POSITION_OFFSET)[None])
+        mask = (cache.arange[None, :] <= idx[:, None])[None, None]
+        for i, layer in enumerate(self.layers):
+            x = layer.step_static(x, cache, i, idx, mask)
+        cache.pos.fill_(p)
+        return x
+
+    def step_static(self, input_ids, cache: StaticCache):
+        """[1, 1] at position ``cache.pos``; graph-capturable (no host reads)."""
+        x = self.layernorm_embedding(self.embed_tokens(input_ids) + self.embed_positions(cache.pos + POSITION_OFFSET)[None])
+        mask = (cache.arange <= cache.pos)[None, None, None, :]
+        for i, layer in enumerate(self.layers):
+            x = layer.step_static(x, cache, i, cache.pos, mask)
+        cache.pos += 1
+        return x
+
     def forward(self, input_ids, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None,
                 past_key_values=None, use_cache=False, output_hidden_states=False, return_dict=True, **_):
         b, l = input_ids.shape

extensions_built_in/audio_models/yue2/src/sheetsage_generate.py

@@ -0,0 +1,198 @@
+"""Launch-light replacement for SheetSage2's ``constrained_prompt_generate``.
+
+Upstream decodes one token per Python iteration: a growing KV cache re-concatenated every step,
+the grammar mask rebuilt from ~10 slice kernels, and a host sync per token, which leaves the
+6-layer / 512-wide decoder at ~5 ms per token on an idle CPU and far worse under dataloader
+load. Here the grammar (``PromptGrammarState``) is tabulated once into ``allowed[state, token]``
+and ``next_state[state, token_type]`` lookups, the decoder runs on a ``StaticCache``, and the
+whole per-token step (mask, argmax, state update, decoder forward for the chosen token) is one
+CUDA graph replayed with a host sync every ``SYNC_EVERY`` tokens. Batch 1 only (what upstream's
+``constrained_prompt_generate`` accepts); anything else falls back to the upstream function.
+"""
+
+import copy
+import math
+from typing import Optional
+
+import torch
+
+from .sheetsage_decoder import StaticCache
+
+SYNC_EVERY = 32
+TYPE_NAMES = ("subbeat_shift", "time", "meter", "eighth_position", "structure", "key", "chord_full", "pitch", "duration")
+EOS_TYPE = len(TYPE_NAMES)
+OTHER_TYPE = EOS_TYPE + 1
+INCOMPLETE = (None, "rhythm_after_meter", "melody_after_pitch")
+N_STATES = 2 * 2 * 5 * 7 * 3
+NO_STOP = 1 << 40
+
+
+def state_index(state) -> int:
+    """PromptGrammarState -> row; ``allowed``/``update`` only read these fields (shift_run only as < 4)."""
+    key = (int(state.payload_count > 0), int(state.in_shift), min(int(state.shift_run), 4), int(state.last_field_index) + 1, INCOMPLETE.index(state.incomplete))
+    return (((key[0] * 2 + key[1]) * 5 + key[2]) * 7 + key[3]) * 3 + key[4]
+
+
+def build_tables(tokenizer, state_cls):
+    """Enumerate the whole grammar key space: allowed [S, V] bool, next_state [S, n_types] long, token type [V] long."""
+    attrs = [{"chord_full": "full_chord"}.get(n, n) for n in TYPE_NAMES]
+    ranges = [(getattr(tokenizer, f"{a}_token_start"), getattr(tokenizer, f"{a}_token_end")) for a in attrs]
+    type_of = torch.full((tokenizer.n_tokens,), OTHER_TYPE, dtype=torch.long)
+    for t, (s, e) in enumerate(ranges):
+        type_of[s:e] = t
+    type_of[tokenizer.eos_token] = EOS_TYPE
+    allowed = torch.zeros(N_STATES, tokenizer.n_tokens, dtype=torch.bool)
+    next_state = torch.arange(N_STATES, dtype=torch.long)[:, None].repeat(1, OTHER_TYPE + 1)
+    for payload in (0, 1):
+        for in_shift in (False, True):
+            for run in range(5):
+                for lfi in range(-1, 6):
+                    for inc in INCOMPLETE:
+                        st = state_cls(tokenizer)
+                        st.payload_count, st.in_shift, st.shift_run, st.last_field_index, st.incomplete = payload, in_shift, run, lfi, inc
+                        i = state_index(st)
+                        allowed[i] = st.allowed(torch.device("cpu"))
+                        for t, (s, e) in enumerate(ranges):
+                            if s < e:
+                                nxt = copy.copy(st)
+                                nxt.update(s)
+                                next_state[i, t] = state_index(nxt)
+    return allowed, next_state, type_of
+
+
+def stop_time_id(stop_seconds: Optional[float], time_hz: int) -> int:
+    """Smallest time id with ``id / time_hz >= stop`` under float division, as upstream compares."""
+    if stop_seconds is None:
+        return NO_STOP
+    t = max(0, math.ceil(float(stop_seconds) * time_hz))
+    while t > 0 and (t - 1) / time_hz >= stop_seconds:
+        t -= 1
+    while t / time_hz < stop_seconds:
+        t += 1
+    return t
+
+
+class FastConstrainedGenerate:
+    """Callable with upstream's ``constrained_prompt_generate`` signature; ``upstream`` handles unsupported calls."""
+
+    def __init__(self, upstream):
+        self.upstream = upstream
+        self.model = None
+        self.key = None
+        self.graph = None
+
+    def __call__(self, model, audio, prompts, max_sequence_length, prefix_tokens=None, autocast_dtype=torch.bfloat16,
+                 stop_time_seconds=None, progress_callback=None, memory=None, step_callback=None):
+        if step_callback is not None or audio.device.type != "cuda" or audio.shape[0] != 1:
+            return self.upstream(model, audio, prompts, max_sequence_length, prefix_tokens=prefix_tokens, autocast_dtype=autocast_dtype,
+                                 stop_time_seconds=stop_time_seconds, progress_callback=progress_callback, memory=memory, step_callback=step_callback)
+        tokenizer = model.tokenizer
+        prefix = list(prefix_tokens) if prefix_tokens is not None else tokenizer.prompt_prefix(prompts)
+        if not prefix or prefix[0] != tokenizer.sos_token:
+            raise ValueError("generation prefix must begin with <|sos|>")
+        if prefix[-1] == tokenizer.eos_token:
+            prefix = prefix[:-1]
+        max_len = int(max_sequence_length)
+        if len(prefix) >= max_len:
+            return torch.tensor(prefix + [tokenizer.eos_token], dtype=torch.long)
+        autocast = torch.autocast(device_type="cuda", dtype=autocast_dtype, cache_enabled=False) if autocast_dtype is not None else torch.autocast(device_type="cuda", enabled=False)
+        if memory is None:
+            with autocast:
+                memory = model.encode(audio)
+        self._ensure_built(model, audio.device, max_len, memory.shape[1], autocast_dtype)
+
+        state = self._state_cls(tokenizer)
+        for token in prefix[prefix.index(tokenizer.out_token) + 1 :]:
+            state.update(token)
+        self.cache.fill_cross(self.dec, memory)
+        hidden = self.dec.prefill_static(torch.tensor([prefix], dtype=torch.long, device=audio.device), self.cache)
+        self.logits.copy_(self._project(hidden))
+        self.state.fill_(state_index(state))
+        self.stop_id.fill_(stop_time_id(stop_time_seconds, tokenizer.time_hz))
+        self.gen_count.zero_()
+        self.finish_pos.fill_(-1)
+        self.done.zero_()
+
+        budget = max_len - len(prefix)
+        steps = 0
+        while steps < budget:
+            for _ in range(min(SYNC_EVERY, budget - steps)):
+                self.graph.replay()
+                steps += 1
+            if progress_callback is not None:
+                progress_callback(len(prefix) + steps)
+            if self.done.item():
+                break
+        n = int(self.finish_pos.item()) + 1 if self.done.item() else steps
+        generated = self.out[:n].tolist()
+        if any(int(self.type_of[t]) == OTHER_TYPE for t in generated):
+            raise RuntimeError("Unexpected prompt token type in constrained generation")
+        tokens = prefix + generated
+        if tokens[-1] != tokenizer.eos_token:
+            tokens.append(tokenizer.eos_token)
+        return torch.tensor(tokens, dtype=torch.long)
+
+    def _project(self, hidden):
+        """last position of [1, L, dim] -> fp32 logits [1, V]"""
+        return torch.nn.functional.linear(hidden[:, -1].to(self.proj_w.dtype), self.proj_w).float()
+
+    # ------------------------------------------------------------------
+    def _ensure_built(self, model, device, max_len, memory_len, autocast_dtype):
+        key = (id(model), str(device), max_len, memory_len, autocast_dtype)
+        if self.key == key and self.graph is not None:
+            return
+        import importlib
+
+        gen = importlib.import_module(type(model).__module__.rsplit(".", 1)[0] + ".generation_sheetsage2")
+        self._state_cls = gen.PromptGrammarState
+        self.model, self.key = model, key
+        tok = model.tokenizer
+        allowed, next_state, type_of = build_tables(tok, self._state_cls)
+        self.allowed, self.next_state = allowed.to(device), next_state.to(device)
+        self.type_of = type_of  # host copy for the post-check
+        self.type_of_dev = type_of.to(device)
+        cache_len = max_len + SYNC_EVERY  # replays may overshoot the budget by < SYNC_EVERY after finishing
+        self.cache = StaticCache(model.decoder, cache_len, memory_len, device, autocast_dtype or torch.float32)
+        self.logits = torch.zeros(1, tok.n_tokens, device=device)
+        self.state = torch.zeros(1, dtype=torch.long, device=device)
+        self.stop_id = torch.zeros(1, dtype=torch.long, device=device)
+        self.gen_count = torch.zeros(1, dtype=torch.long, device=device)
+        self.finish_pos = torch.full((1,), -1, dtype=torch.long, device=device)
+        self.done = torch.zeros(1, dtype=torch.bool, device=device)
+        self.out = torch.zeros(cache_len, dtype=torch.long, device=device)
+        self.cur = torch.zeros(1, 1, dtype=torch.long, device=device)
+        self.eos = tok.eos_token
+        self.time_start = tok.time_token_start
+        # bf16 copies of the decoder linears and the tied projection table: what autocast would cast on every step
+        dtype = autocast_dtype or torch.float32
+        self.dec = copy.deepcopy(model.decoder, memo={id(model.decoder.embed_tokens): model.decoder.embed_tokens})
+        for m in self.dec.modules():
+            if isinstance(m, torch.nn.Linear):
+                m.to(dtype)
+        self.proj_w = model.output_projection.weight.detach().to(dtype)
+
+        def step():
+            masked = self.logits.masked_fill(~self.allowed[self.state], float("-inf"))
+            tok_id = masked.argmax(-1)  # [1]
+            ttype = self.type_of_dev[tok_id]
+            stop_hit = (ttype == 1) & ((tok_id - self.time_start) >= self.stop_id)
+            finished = (tok_id == self.eos) | stop_hit
+            self.out.index_copy_(0, self.gen_count, tok_id)
+            self.finish_pos.copy_(torch.where(finished & (self.finish_pos < 0), self.gen_count, self.finish_pos))
+            self.done.logical_or_(finished)
+            self.gen_count += 1
+            self.state.copy_(self.next_state[self.state, ttype])
+            self.cur.copy_(tok_id[:, None])
+            self.logits.copy_(self._project(self.dec.step_static(self.cur, self.cache)))
+
+        # warm up on a side stream (cuBLAS workspaces, autocast paths), then capture
+        s = torch.cuda.Stream(device)
+        s.wait_stream(torch.cuda.current_stream(device))
+        with torch.cuda.stream(s):
+            for _ in range(3):
+                step()
+        torch.cuda.current_stream(device).wait_stream(s)
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph, stream=s):
+            step()
+        torch.cuda.synchronize(device)

extensions_built_in/audio_models/yue2/src/sheetsage_repair.py

@@ -0,0 +1,190 @@
+"""Lenient ABC rebuild for SheetSage2 transcriptions whose strict notation export failed.
+
+Upstream's ``export_result`` rejects the whole sheet on any inconsistency in the decoded events
+(non-consecutive beat IDs in a bar, a note shorter than a subbeat, overlapping quantized notes,
+a chord change shorter than the grid, an unparsable chord symbol). The decoded events are still
+mostly right, so this rebuilds the same inputs with the offending rows repaired or dropped and
+runs upstream's own ``generate_abc_from_data`` again, so the result is still validated ABC.
+A quality gate refuses sheets that needed more than ``MAX_REPAIR_FRACTION`` of beats or notes touched.
+"""
+
+import importlib
+from typing import List, Tuple
+
+import numpy as np
+
+MAX_REPAIR_FRACTION = 0.2
+
+
+class AbcRepairError(ValueError):
+    pass
+
+
+def _modules(model):
+    pkg = type(model).__module__.rsplit(".", 1)[0]
+    return importlib.import_module(pkg + ".notation_sheetsage2"), importlib.import_module(pkg + ".exports_sheetsage2"), importlib.import_module(pkg + ".midi_sheetsage2")
+
+
+def _beat_rows(events) -> Tuple[List[list], int]:
+    """exports.rhythm_rows without raising on off-grid eighth positions; returns rows + skipped count."""
+    from fractions import Fraction
+
+    rows, meter, skipped = [], None, 0
+    for event in events:
+        rhythm = event["values"].get("rhythm", {})
+        meter = rhythm.get("meter", meter)
+        eighth = rhythm.get("eighth_position")
+        if eighth is None or meter is None:
+            continue
+        position = Fraction(int(eighth) * int(meter[1]), 8)
+        if position.denominator != 1 or not 0 <= position < meter[0]:
+            skipped += 1
+            continue
+        rows.append([float(event["time"]), int(position) + 1, int(meter[0]), int(meter[1])])
+    return rows, skipped
+
+
+def _repair_beats(rows: List[list]) -> Tuple[List[list], int]:
+    """Strictly increasing times; consecutive beat IDs inside each bar. A repeated ID within half a beat
+    period is a duplicated row (dropped); any other non-increase starts a new bar; skipped IDs are filled."""
+    if len(rows) < 2:
+        raise AbcRepairError("fewer than two decoded beats")
+    period = float(np.median(np.diff([r[0] for r in rows]))) if len(rows) > 2 else 0.5
+    out, changed = [], 0
+    for row in rows:
+        if out and row[0] <= out[-1][0]:
+            changed += 1
+            continue
+        if out and row[1] == out[-1][1] and row[0] - out[-1][0] < 0.5 * period:
+            changed += 1
+            continue
+        if out:
+            expected = 1 if row[1] <= out[-1][1] else out[-1][1] + 1
+            if expected > row[2]:  # ran past the declared numerator: the model missed a downbeat
+                expected = 1
+            if expected != row[1]:
+                changed += 1
+                row = [row[0], expected, row[2], row[3]]
+        out.append(row)
+    if all(r[1] != 1 for r in out):
+        raise AbcRepairError("no downbeat in the decoded beats")
+    return out, changed
+
+
+def _extend_beats(rows, duration, notes):
+    """Same tail extension as export_result: continue the final tempo to the audio/last note end."""
+    period = float(np.median(np.diff([b[0] for b in rows[-9:]])))
+    if period <= 0:
+        raise AbcRepairError("decoded beats must increase in time")
+    beats = [list(b) for b in rows]
+    end = max(duration, max((n[1] for n in notes), default=0))
+    while beats[-1][0] < end - 1e-6:
+        prev = beats[-1]
+        beats.append([prev[0] + period, prev[1] % prev[2] + 1, prev[2], prev[3]])
+    return beats
+
+
+def _merge_short_intervals(rows, subbeat_times, quantize) -> Tuple[List[list], int]:
+    """Rows that quantize to zero subbeats are absorbed into the previous row (or the next, for the first)."""
+    out, merged = [], 0
+    for row in rows:
+        start_t, end_t = quantize(row[0], subbeat_times), quantize(row[1], subbeat_times)
+        if end_t <= start_t:
+            merged += 1
+            if out:
+                out[-1][1] = max(out[-1][1], row[1])
+            continue
+        if out and out[-1][1] > row[0]:  # keep contiguity after a merge
+            row = [out[-1][1], row[1], row[2]]
+            if row[1] <= row[0]:
+                merged += 1
+                continue
+        out.append(list(row))
+    return out, merged
+
+
+def _filter_notes(notes, subbeat_times, notation) -> Tuple[List[list], int]:
+    """Drop notes that quantize to nothing or land on subbeats another note of the same track already holds."""
+    boundaries = notation._subbeat_boundaries(subbeat_times)
+    kept, dropped = [], 0
+    for track in (0, 1):
+        taken = np.zeros(len(subbeat_times), dtype=bool)
+        for note in sorted((n for n in notes if n[3] == track), key=lambda n: (n[0], n[1], n[2])):
+            start_t = max(0, min(int(np.searchsorted(boundaries, note[0])), len(taken) - 1))
+            end_t = max(0, min(int(np.searchsorted(boundaries, note[1])), len(taken) - 1))
+            if end_t <= start_t or taken[start_t:end_t].any():
+                dropped += 1
+                continue
+            taken[start_t:end_t] = True
+            kept.append(list(note))
+    return sorted(kept), dropped
+
+
+def repair_abc(model, result: dict, melody_only: bool = False) -> Tuple[str, str]:
+    """Rebuild ABC for a transcribe() result whose ``abc`` is None. Returns (abc, summary) or raises AbcRepairError."""
+    notation, exports, midi_mod = _modules(model)
+    events = result["events"]
+    duration = float(result["duration_seconds"])
+    notes = []
+    for event in events:
+        start = float(event["time"])
+        for note in event["values"].get("melody", ()):
+            end = min(duration, float(note["end_time"]))
+            if end > start:
+                notes.append([start, end, int(note["pitch"]), int(note["track"])])
+    notes.sort()
+    clean, _ = exports.notation_notes(notes)
+
+    rows, skipped = _beat_rows(events)
+    rows, changed = _repair_beats(rows)
+    beats = _extend_beats(rows, duration, clean)
+    n_beats = len(rows) + skipped
+    if (changed + skipped) > MAX_REPAIR_FRACTION * n_beats:
+        raise AbcRepairError(f"{changed + skipped}/{n_beats} beat rows needed repair")
+
+    parsed = notation._parse_beats(notation._row_entries([[str(x) for x in b] for b in beats], "beats"), "beats")
+    measures, _ = notation.infer_measures(parsed, meter_conflict="infer")
+    subbeat_times, _, _ = notation._build_grid(parsed, measures)
+
+    intervals = {}
+    merged = 0
+    for field in ("chord", "key", "structure"):
+        rows = exports.interval_rows(events, field, duration)
+        rows = [[max(beats[0][0], a), min(beats[-1][0], b), v] for a, b, v in rows if b > beats[0][0] and a < beats[-1][0]]
+        if field != "structure":
+            rows, m = _merge_short_intervals(rows, subbeat_times, notation._quantize_time)
+            merged += m
+        intervals[field] = rows
+    if not intervals["key"]:
+        raise AbcRepairError("no key was decoded")
+    bad_chords = 0
+    for row in intervals["chord"]:
+        try:
+            notation.chord_symbol_to_abc(row[2])
+        except notation.AbcRebuildError:
+            row[2] = "N"
+            bad_chords += 1
+    keys = []
+    for row in intervals["key"]:
+        try:
+            notation.key_symbol_to_abc(row[2])
+            keys.append(row)
+        except notation.AbcRebuildError:
+            if keys:
+                keys[-1][1] = row[1]
+    if not keys:
+        raise AbcRepairError("no usable key")
+    intervals["key"] = keys
+
+    kept, dropped = _filter_notes(clean, subbeat_times, notation)
+    if dropped > MAX_REPAIR_FRACTION * max(1, len(clean)):
+        raise AbcRepairError(f"{dropped}/{len(clean)} notes could not be placed on the beat grid")
+
+    try:
+        text, _ = notation.generate_abc_from_data(
+            midi_mod.midi_bytes(exports._midi(kept)), beats, intervals["chord"], intervals["key"], intervals["structure"], melody_only=melody_only
+        )
+    except ValueError as exc:
+        raise AbcRepairError(f"rebuild still failed: {exc}") from exc
+    summary = f"beats repaired {changed + skipped}/{n_beats}, notes dropped {dropped}/{len(clean)}, intervals merged {merged}, chords blanked {bad_chords}"
+    return text, summary

extensions_built_in/audio_models/yue2/src/tokenizer.py

@@ -15,6 +15,8 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 
+from toolkit.basic import UnusableFileError
+
 from .model import ABC_END, ABC_START, CODEC_SIZE, EOD, INSTRUCTIONS, MUSIC_START
 
 MERT_REPO = "m-a-p/MERT-v2-FullSong"
@@ -107,6 +109,7 @@ class SheetSage2Transcriber:
         from .shims import install_shims
 
         install_shims()  # mir_eval.chord / pretty_midi stand-ins unless the real packages exist
+        import importlib
         import sys
 
         from safetensors.torch import load_file
@@ -114,11 +117,16 @@ class SheetSage2Transcriber:
         from transformers.dynamic_module_utils import get_class_from_dynamic_module
 
         from .sheetsage_decoder import ScoreDecoder
+        from .sheetsage_generate import FastConstrainedGenerate
 
         cfg = AutoConfig.from_pretrained(repo, trust_remote_code=True)
         cfg.weights_format = "merged"  # the Comfy file carries the LoRA already folded into the encoder
         cls = get_class_from_dynamic_module(cfg.auto_map["AutoModel"], repo)
         sys.modules[cls.__module__].BartDecoder = ScoreDecoder  # upstream targets transformers 4.45's BartDecoder
+        pipeline = importlib.import_module(cls.__module__.rsplit(".", 1)[0] + ".pipeline_sheetsage2")
+        if not isinstance(pipeline.constrained_prompt_generate, FastConstrainedGenerate):
+            # per-window decode loop: static cache + CUDA graph instead of upstream's per-token Python loop
+            pipeline.constrained_prompt_generate = FastConstrainedGenerate(pipeline.constrained_prompt_generate)
         cls.post_init = lambda self: None  # transformers-5 tied-weight bookkeeping rejects upstream's list; weights load explicitly
         model = cls(cfg)
         sd = {k: v.to(torch.float32) for k, v in load_file(weights_path).items()}
@@ -130,6 +138,7 @@ class SheetSage2Transcriber:
         self.model = model.eval()
         self.model.requires_grad_(False)
         _rebuild_rotary(self.model)
+        self.warnings: List[str] = []  # repair notes for the caller to print with the track path
 
     def to(self, device):
         self.model.to(device)
@@ -142,11 +151,19 @@ class SheetSage2Transcriber:
     @torch.no_grad()
     def transcribe(self, waveform: torch.Tensor, sample_rate: int, cot: str = "full") -> str:
         """waveform [C, samples] -> ABC text (``cot`` "full" or "melody")."""
+        from .sheetsage_repair import AbcRepairError, repair_abc
+
         mono = waveform.float().mean(0) if waveform.dim() == 2 else waveform.float()
         result = self.model.transcribe(mono.cpu(), sampling_rate=sample_rate, melody_only=cot == "melody")
         abc = result.get("abc") or ""
         if not abc.strip():
-            raise RuntimeError(f"SheetSage2 produced no ABC: {result.get('abc_error') or 'unknown reason'}")
+            # strict export rejected the sheet: rebuild it with the bad rows repaired, else the cacher drops the track
+            error = result.get("abc_error") or "unknown reason"
+            try:
+                abc, summary = repair_abc(self.model, result, melody_only=cot == "melody")
+            except AbcRepairError as exc:
+                raise UnusableFileError(f"SheetSage2 produced no ABC: {error}; repair failed: {exc}") from exc
+            self.warnings.append(f"SheetSage2 ABC repaired ({error}): {summary}")
         return abc
 
 

extensions_built_in/audio_models/yue2/yue2_model.py

@@ -265,6 +265,13 @@ class YuE2AudioModel(BaseAudioModel):
             self.transcriber = SheetSage2Transcriber(resolve_named_file(self.sheetsage_path, component="sheetsage2"))
         return self.transcriber
 
+    def pop_encode_warnings(self) -> List[str]:
+        """Repair notes from the last encode_audio call; the latent cacher prints them with the file path."""
+        if self.transcriber is None:
+            return []
+        warnings, self.transcriber.warnings = self.transcriber.warnings, []
+        return warnings
+
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["model.layers"]
 

toolkit/basic.py

@@ -4,6 +4,11 @@ import os
 import torch
 
 
+class UnusableFileError(Exception):
+    """A model raises this from encode when a dataset file cannot be used (e.g. no sheet could be transcribed);
+    the latent cacher drops the file instead of failing the job."""
+
+
 def value_map(inputs, min_in, max_in, min_out, max_out):
     return (inputs - min_in) * (max_out - min_out) / (max_in - min_in) + min_out
 

toolkit/dataloader_mixins.py

@@ -19,7 +19,7 @@ from tqdm import tqdm
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection, SiglipImageProcessor
 
 from toolkit.audio.preserve_pitch import time_stretch_preserve_pitch
-from toolkit.basic import flush, value_map
+from toolkit.basic import UnusableFileError, flush, value_map
 from toolkit.buckets import get_bucket_for_image_size, get_resolution
 from toolkit.config_modules import ControlTypes
 from toolkit.control_generator import ControlGenerator
@@ -2005,7 +2005,16 @@ class LatentCachingMixin:
                     if needs_encode and not did_move:
                         self.sd.set_device_state_preset('cache_latents')
                         did_move = True
-                    self._cache_one_latent(file_item, latent_path, cached_state_dict, needs_encode, to_disk, to_memory)
+                    try:
+                        self._cache_one_latent(file_item, latent_path, cached_state_dict, needs_encode, to_disk, to_memory)
+                    except UnusableFileError as e:
+                        pbar.write(f"Skipping unusable file {file_item.path}: {e}")
+                        failed_items.append(file_item)
+                        pbar.update(1)
+                        continue
+                    # models may report per-file repairs (e.g. a patched sheet); print them on their own line with the path
+                    for warning in getattr(self.sd, 'pop_encode_warnings', lambda: [])():
+                        pbar.write(f"{file_item.path}: {warning}")
                     file_item.is_latent_cached = True
                     i += 1
                     pbar.update(1)
@@ -2014,7 +2023,7 @@ class LatentCachingMixin:
                 pbar.close()
 
             if failed_items:
-                print_acc(f"Removed {len(failed_items)} files from the dataset that failed to load")
+                print_acc(f"Removed {len(failed_items)} files from the dataset that failed to load or were unusable")
                 self._remove_file_items(failed_items)
 
             # restore device state
@@ -2116,6 +2125,8 @@ class LatentCachingMixin:
                         for k, v in latent.extras.items():
                             if torch.is_tensor(v):
                                 state_dict[f'{DISK_PREFIX}{k}'] = v.clone().detach().cpu()
+            except UnusableFileError:
+                raise
             except Exception as e:
                 print_acc(f"Error processing image: {file_item.path}")
                 print_acc(f"Error: {str(e)}")