From 410206e41c7a323c41b0eff67fc5d1e02d26c97b Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 16 Sep 2026 12:24:39 -0600 Subject: [PATCH] Improve speed of sage sheet decoder with yue2. Recover from tracks that fail --- .../yue2/src/sheetsage_decoder.py | 67 ++++++ .../yue2/src/sheetsage_generate.py | 198 ++++++++++++++++++ .../audio_models/yue2/src/sheetsage_repair.py | 190 +++++++++++++++++ .../audio_models/yue2/src/tokenizer.py | 19 +- .../audio_models/yue2/yue2_model.py | 7 + toolkit/basic.py | 5 + toolkit/dataloader_mixins.py | 17 +- 7 files changed, 499 insertions(+), 4 deletions(-) create mode 100644 extensions_built_in/audio_models/yue2/src/sheetsage_generate.py create mode 100644 extensions_built_in/audio_models/yue2/src/sheetsage_repair.py diff --git a/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py b/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py index 5ae0200..2d17203 100644 --- a/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py +++ b/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 diff --git a/extensions_built_in/audio_models/yue2/src/sheetsage_generate.py b/extensions_built_in/audio_models/yue2/src/sheetsage_generate.py new file mode 100644 index 0000000..c25956d --- /dev/null +++ b/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) diff --git a/extensions_built_in/audio_models/yue2/src/sheetsage_repair.py b/extensions_built_in/audio_models/yue2/src/sheetsage_repair.py new file mode 100644 index 0000000..49bb0ac --- /dev/null +++ b/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 diff --git a/extensions_built_in/audio_models/yue2/src/tokenizer.py b/extensions_built_in/audio_models/yue2/src/tokenizer.py index 392aab3..c0954ae 100644 --- a/extensions_built_in/audio_models/yue2/src/tokenizer.py +++ b/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 diff --git a/extensions_built_in/audio_models/yue2/yue2_model.py b/extensions_built_in/audio_models/yue2/yue2_model.py index bdad93b..a58ccd1 100644 --- a/extensions_built_in/audio_models/yue2/yue2_model.py +++ b/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"] diff --git a/toolkit/basic.py b/toolkit/basic.py index 3784445..d98ff3c 100644 --- a/toolkit/basic.py +++ b/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 diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 2f5c0ed..e108ec5 100644 --- a/toolkit/dataloader_mixins.py +++ b/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)}")