diff --git a/extensions_built_in/audio_models/yue2/src/model.py b/extensions_built_in/audio_models/yue2/src/model.py index 4bbcec0..fb80541 100644 --- a/extensions_built_in/audio_models/yue2/src/model.py +++ b/extensions_built_in/audio_models/yue2/src/model.py @@ -192,7 +192,8 @@ class YuE2AR(YuE2Expert): self.model.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) def embed(self, ids: torch.Tensor): - return self.model.embed_tokens(ids.to(self.model.embed_tokens.weight.device)) + # never touch embed_tokens.weight: the int8 checkpoint's table dequantizes the whole vocab on that property + return self.model.embed_tokens(ids.to(self.device)) def prefill(self, inputs_embeds: torch.Tensor, return_hidden: bool = False): """Causal pass over [B, L, H]. Returns per-layer post-rope (k, v) as @@ -355,17 +356,49 @@ class YuE2Model(nn.Module, OstrisModelMixin): @classmethod def load_from_state_dict(cls, sd: dict, dtype=torch.bfloat16): - """``sd`` is the Comfy-Org all-in-one checkpoint (raw keys).""" + """``sd`` is a Comfy-Org all-in-one checkpoint (raw keys): bf16, or the int8 convrot repack whose + ``comfy_quant`` markers attach the shipped quantization directly (no requantization).""" + from toolkit.util.comfy_quant_import import Int8Embedding, import_comfy_quantized_layers + from toolkit.util.ostris_quant import OstrisLinear + model = cls() nar_sd = {k[len("model.diffusion_model.") :]: v for k, v in sd.items() if k.startswith("model.diffusion_model.")} ar_sd = {k[len("text_encoders.") :]: v for k, v in sd.items() if k.startswith("text_encoders.") and k != "text_encoders.yue2_tokenizer_json"} + prequantized = any(k.endswith(".comfy_quant") for k in sd) for name, module, part in (("NAR", model.nar, nar_sd), ("AR", model.ar, ar_sd)): - missing, unexpected = module.load_state_dict({k: v.to(dtype) for k, v in part.items()}, strict=False) + allowed_missing = set() + if prequantized: + quantized_paths = {k[: -len(".comfy_quant")] for k in part if k.endswith(".comfy_quant")} + part, n = import_comfy_quantized_layers(module, part, orig_dtype=dtype) + for mod_name, m in module.named_modules(): + if isinstance(m, OstrisLinear): + allowed_missing.add(f"{mod_name}.weight") + if m.bias is not None: + allowed_missing.add(f"{mod_name}.bias") + m.bias.data = m.bias.data.to(dtype) + elif isinstance(m, (Int8Embedding, nn.Embedding)) and mod_name in quantized_paths: + allowed_missing.add(f"{mod_name}.weight") + part = {k: (v.to(dtype) if v.is_floating_point() else v) for k, v in part.items()} + missing, unexpected = module.load_state_dict(part, strict=False) + missing = [k for k in missing if k not in allowed_missing] if missing: raise ValueError(f"YuE2 {name} missing keys: {missing[:5]} (+{max(0, len(missing) - 5)})") if unexpected: print(f" YuE2 {name} unexpected: {len(unexpected)} (first 3: {unexpected[:3]})") - return model.to(dtype) + if not prequantized: + return model.to(dtype) + # cast the remaining float tensors (norms, embeddings, buffers) without touching the quantized layers + for m in model.modules(): + if isinstance(m, (OstrisLinear, Int8Embedding)): + continue + for pname, prm in list(m.named_parameters(recurse=False)): + if prm.is_floating_point(): + prm.data = prm.data.to(dtype) + for bname, buf in list(m.named_buffers(recurse=False)): + if buf is not None and buf.is_floating_point(): + setattr(m, bname, buf.to(dtype)) + model.aitk_is_quantized = True # aitk_post_load keeps the shipped convrot8 layers when convrot8 is requested + return model @torch.no_grad() @@ -376,6 +409,8 @@ def merge_nar_lora(model: YuE2Model, ckpt_path: str, scale: float = 1.0): ck = torch.load(ckpt_path, map_location="cpu", weights_only=True) except Exception: ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) + if getattr(model, "aitk_is_quantized", False): + raise ValueError("merge_nar_lora needs the bf16 checkpoint; the int8 convrot repack cannot take merged weights") tensors = iter(ck["lora"]) cfg = model.cfg inner = cfg.num_attention_heads * cfg.head_dim diff --git a/extensions_built_in/audio_models/yue2/src/pipeline.py b/extensions_built_in/audio_models/yue2/src/pipeline.py index af516d0..e17d884 100644 --- a/extensions_built_in/audio_models/yue2/src/pipeline.py +++ b/extensions_built_in/audio_models/yue2/src/pipeline.py @@ -11,11 +11,13 @@ from typing import Callable, List, Optional, Tuple import torch from .model import ( + ABC_END, CODEC_OFFSET, CODEC_SIZE, CONTEXT, - MUSIC_END, + EOD, KVCache, + MUSIC_END, YuE2Model, rope_cos_sin, ) @@ -40,14 +42,22 @@ def sample_logits( penalty_window: int, min_tokens: int, generator: torch.Generator, + phase: str = "semantic", + legacy_off: bool = False, ) -> int: - scores = logits.float().clone() + """ComfyUI's ``distribution`` for one row: ``abc`` allows the text vocabulary + ABC_END, + ``semantic`` the codec ids + MUSIC_END. ``legacy_off`` = cot "off" (bf16 scores, keep top 3).""" + scores = logits.clone() if legacy_off else logits.float().clone() + end = ABC_END if phase == "abc" else MUSIC_END allowed = torch.full_like(scores, -torch.inf) - allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 - allowed[MUSIC_END] = 0 + if phase == "abc": + allowed[:EOD] = 0 + else: + allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 + allowed[end] = 0 scores += allowed if step < min_tokens: - scores[MUSIC_END] = -torch.inf + scores[end] = -torch.inf if repetition_penalty != 1.0 and history: recent = torch.tensor(history[-penalty_window:], dtype=torch.long, device=scores.device) counts = torch.zeros_like(scores).scatter_add_(0, recent, torch.ones_like(recent, dtype=scores.dtype)) @@ -62,8 +72,7 @@ def sample_logits( values, indices = scores.sort(descending=True) probs = values.softmax(-1) removed = probs.cumsum(-1) - probs > top_p - # released "off" mode keeps the top 3 candidates regardless of mass - removed[:3] = False + removed[: 3 if legacy_off else 1] = False values.masked_fill_(removed, -torch.inf) scores = values.scatter(-1, indices, values) probs = scores.softmax(-1) @@ -89,6 +98,10 @@ class ARDecodeGraph: self.token = torch.zeros((1, 1), device=device, dtype=torch.long) self.pos = torch.zeros((1, 1), device=device, dtype=torch.long) self.index = torch.arange(capacity, device=device) + # the graph covers the layer stack only; embedding lookup and lm_head run outside it + # (the int8 repack's embedding/lm_head kernels invalidate stream capture) + self.x_in = torch.zeros((1, 1, cfg.hidden_size), device=device, dtype=dtype) + self.h_out = torch.zeros((1, cfg.hidden_size), device=device, dtype=dtype) self.logits = torch.empty((1, cfg.vocab_size), device=device, dtype=dtype) self.graph = None @@ -98,9 +111,9 @@ class ARDecodeGraph: self.v[i][:, :, :length] = v self.pos.fill_(length) - def _step(self): + def _layers(self): ar = self.ar - x = ar.model.embed_tokens(self.token) + x = self.x_in cos, sin = rope_cos_sin(self.pos, self.head_dim, ar.cfg.rope_theta) mask = (self.index <= self.pos)[None, None] # [1, 1, 1, capacity] scale = self.head_dim ** -0.5 @@ -115,25 +128,31 @@ class ARDecodeGraph: out = torch.matmul(scores, self.v[i]).reshape(1, 1, -1) x = x + layer.self_attn.o_proj(out) x = x + layer.mlp(layer.post_attention_layernorm(x)) - self.logits.copy_(ar.model.lm_head(ar.model.norm(x[:, -1]))) + self.h_out.copy_(ar.model.norm(x[:, -1])) + + def _step(self): + ar = self.ar + self.x_in.copy_(ar.model.embed_tokens(self.token).to(self.x_in.dtype)) + if self.graph is not None: + self.graph.replay() + else: + self._layers() + self.logits.copy_(ar.model.lm_head(self.h_out)) def capture(self): stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for _ in range(2): - self._step() + self._layers() torch.cuda.current_stream().wait_stream(stream) self.graph = torch.cuda.CUDAGraph() with torch.cuda.graph(self.graph): - self._step() + self._layers() def step(self, token: int): self.token.fill_(token) - if self.graph is not None: - self.graph.replay() - else: - self._step() + self._step() self.pos.add_(1) return self.logits[0] @@ -146,19 +165,22 @@ class YuE2Pipeline: self.use_cuda_graphs = True @torch.no_grad() - def generate_codec_tokens( + def generate_tokens( self, prefix_embeds: torch.Tensor, # [1, L, H] + phase: str, max_tokens: int, seed: int, - temperature: float = 1.0, - top_p: float = 0.95, - top_k: int = 100, - repetition_penalty: float = 1.2, - penalty_window: int = 50, - min_tokens: int = 200, + temperature: float, + top_p: float, + top_k: int, + repetition_penalty: float, + penalty_window: int, + min_tokens: int, progress: Optional[Callable[[int, int], None]] = None, + legacy_off: bool = False, ) -> List[int]: + """Sample AR tokens after ``prefix_embeds`` until the phase's end token (raw vocab ids).""" ar = self.model.ar device = prefix_embeds.device prefix_len = prefix_embeds.shape[1] @@ -185,11 +207,13 @@ class YuE2Pipeline: logits = ar.prefill_into_cache(prefix_embeds.to(ar.dtype), cache)[0] history: List[int] = [] min_tokens = min(min_tokens, max_tokens) + end = ABC_END if phase == "abc" else MUSIC_END for step in range(max_tokens): token = sample_logits( - logits, history, step, temperature, top_p, top_k, repetition_penalty, penalty_window, min_tokens, generator + logits, history, step, temperature, top_p, top_k, repetition_penalty, penalty_window, min_tokens, generator, + phase=phase, legacy_off=legacy_off, ) - if token == MUSIC_END: + if token == end: break history.append(token) if progress is not None: @@ -203,7 +227,46 @@ class YuE2Pipeline: if graph is None: del cache del graph - return [t - CODEC_OFFSET for t in history] + return history + + def generate_codec_tokens( + self, + prefix_embeds: torch.Tensor, + max_tokens: int, + seed: int, + temperature: float = 1.0, + top_p: float = 0.95, + top_k: int = 100, + repetition_penalty: float = 1.2, + penalty_window: int = 50, + min_tokens: int = 200, + progress: Optional[Callable[[int, int], None]] = None, + legacy_off: bool = False, + ) -> List[int]: + """Music phase (ComfyUI defaults). Returns codec ids without the vocab offset.""" + ids = self.generate_tokens( + prefix_embeds, "semantic", max_tokens, seed, temperature, top_p, top_k, repetition_penalty, penalty_window, + min_tokens, progress, legacy_off=legacy_off, + ) + return [t - CODEC_OFFSET for t in ids] + + def generate_abc_tokens( + self, + prefix_embeds: torch.Tensor, + seed: int, + max_tokens: int = 8192, + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 30, + repetition_penalty: float = 1.005, + penalty_window: int = 100, + progress: Optional[Callable[[int, int], None]] = None, + ) -> List[int]: + """ABC phase (ComfyUI "Generate ABC" defaults): text ids of the sheet, without ABC_END.""" + return self.generate_tokens( + prefix_embeds, "abc", max_tokens, seed, temperature, top_p, top_k, repetition_penalty, penalty_window, + min(32, max_tokens), progress, + ) @torch.no_grad() def synthesize( diff --git a/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py b/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py new file mode 100644 index 0000000..5ae0200 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py @@ -0,0 +1,110 @@ +"""SheetSage2's score decoder (BART-style, post-LN, learned positions with offset 2) in plain +torch, so the m-a-p/SheetSage2 remote code runs on current transformers without BartDecoder. +Parameter names follow the released checkpoint (decoder.embed_positions, decoder.layers.N.self_attn.*, +encoder_attn.*, fc1/fc2, *_layer_norm, layernorm_embedding). KV cache = per-layer tuples with the +batch dimension first, which the upstream generation loop can index.""" + +from types import SimpleNamespace +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +POSITION_OFFSET = 2 # BART learned positional embedding convention + + +class _Attention(nn.Module): + def __init__(self, dim: int, heads: int): + super().__init__() + self.heads = heads + self.head_dim = dim // heads + self.q_proj = nn.Linear(dim, dim) + self.k_proj = nn.Linear(dim, dim) + self.v_proj = nn.Linear(dim, dim) + self.out_proj = nn.Linear(dim, dim) + + def _split(self, x): + b, l, _ = x.shape + return x.view(b, l, self.heads, self.head_dim).transpose(1, 2) + + def forward(self, x, kv_source=None, mask=None, past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, cache_static=False): + q = self._split(self.q_proj(x)) + if cache_static and past is not None: + k, v = past # cross attention: memory keys never change + else: + src = x if kv_source is None else kv_source + k, v = self._split(self.k_proj(src)), self._split(self.v_proj(src)) + if past is not None and not cache_static: + k, v = torch.cat([past[0], k], 2), torch.cat([past[1], v], 2) + 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)), (k, v) + + +class _Layer(nn.Module): + def __init__(self, dim: int, heads: int, ffn: int): + super().__init__() + self.self_attn = _Attention(dim, heads) + self.self_attn_layer_norm = nn.LayerNorm(dim) + self.encoder_attn = _Attention(dim, heads) + self.encoder_attn_layer_norm = nn.LayerNorm(dim) + self.fc1 = nn.Linear(dim, ffn) + self.fc2 = nn.Linear(ffn, dim) + self.final_layer_norm = nn.LayerNorm(dim) + + def forward(self, x, memory, self_mask, past): + self_past = None if past is None else past[:2] + cross_past = None if past is None else past[2:] + h, self_kv = self.self_attn(x, mask=self_mask, past=self_past) + x = self.self_attn_layer_norm(x + h) + h, cross_kv = self.encoder_attn(x, kv_source=memory, past=cross_past, cache_static=True) + x = self.encoder_attn_layer_norm(x + h) + x = self.final_layer_norm(x + self.fc2(F.gelu(self.fc1(x)))) + return x, self_kv + cross_kv + + +class ScoreDecoder(nn.Module): + """Drop-in for the BartDecoder the upstream model builds: ``ScoreDecoder(bart_config, embed_tokens=...)``.""" + + def __init__(self, config, embed_tokens: nn.Embedding): + super().__init__() + dim, heads = config.d_model, config.decoder_attention_heads + self.embed_tokens = embed_tokens + self.embed_positions = nn.Embedding(config.max_position_embeddings + POSITION_OFFSET, dim) + self.layernorm_embedding = nn.LayerNorm(dim) + self.layers = nn.ModuleList([_Layer(dim, heads, config.decoder_ffn_dim) for _ in range(config.decoder_layers)]) + self.pad_token_id = config.pad_token_id + + def gradient_checkpointing_disable(self): + pass + + 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 + past_len = 0 if past_key_values is None else past_key_values[0][0].shape[2] + positions = torch.arange(past_len, past_len + l, device=input_ids.device) + POSITION_OFFSET + x = self.embed_tokens(input_ids) + self.embed_positions(positions)[None] + x = self.layernorm_embedding(x) + mask = None + if l > 1: + causal = torch.ones(l, past_len + l, dtype=torch.bool, device=x.device).tril(past_len) + mask = causal[None, None] + if attention_mask is not None: + mask = mask & attention_mask.bool()[:, None, None, :] + elif attention_mask is not None and past_len == 0: + mask = attention_mask.bool()[:, None, None, :] + if mask is not None: + mask = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, float("-inf")) + hidden_states = [x] if output_hidden_states else None + new_cache = [] + for i, layer in enumerate(self.layers): + x, kv = layer(x, encoder_hidden_states, mask, None if past_key_values is None else past_key_values[i]) + new_cache.append(kv) + if hidden_states is not None: + hidden_states.append(x) + return SimpleNamespace( + last_hidden_state=x, + past_key_values=tuple(new_cache) if use_cache else None, + hidden_states=tuple(hidden_states) if hidden_states is not None else None, + ) diff --git a/extensions_built_in/audio_models/yue2/src/shims/__init__.py b/extensions_built_in/audio_models/yue2/src/shims/__init__.py new file mode 100644 index 0000000..7f29dc2 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/shims/__init__.py @@ -0,0 +1,27 @@ +"""Stand-ins for mir_eval.chord and pretty_midi so m-a-p/SheetSage2's remote code imports without +those packages. Installed into sys.modules only when the real packages are absent.""" + +import importlib +import sys +import types + + +def install_shims(): + for name in ("mir_eval", "pretty_midi"): + try: + importlib.import_module(name) + continue # real package available; leave it alone + except ImportError: + pass + if name == "mir_eval": + from . import mir_eval_chord + + pkg = types.ModuleType("mir_eval") + pkg.chord = mir_eval_chord + pkg.__path__ = [] + sys.modules["mir_eval"] = pkg + sys.modules["mir_eval.chord"] = mir_eval_chord + else: + from . import pretty_midi + + sys.modules["pretty_midi"] = pretty_midi diff --git a/extensions_built_in/audio_models/yue2/src/shims/mir_eval_chord.py b/extensions_built_in/audio_models/yue2/src/shims/mir_eval_chord.py new file mode 100644 index 0000000..e37eeb9 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/shims/mir_eval_chord.py @@ -0,0 +1,518 @@ +"""Chord-label helpers vendored from mir_eval (mir_eval/chord.py), so the SheetSage2 remote code +runs without the mir_eval package. Unmodified except for the removed module preamble. + +The MIT License (MIT) + +Copyright (c) 2014 Colin Raffel + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +import re +import warnings + +import numpy as np + + +BITMAP_LENGTH = 12 +NO_CHORD = "N" +NO_CHORD_ENCODED = -1, np.array([0] * BITMAP_LENGTH), -1 +X_CHORD = "X" +X_CHORD_ENCODED = -1, np.array([-1] * BITMAP_LENGTH), -1 + + +class InvalidChordException(Exception): + r"""Exception class for suspect / invalid chord labels""" + + def __init__(self, message="", chord_label=None): + self.message = message + self.chord_label = chord_label + self.name = self.__class__.__name__ + super().__init__(message) + + +# --- Chord Primitives --- +def _pitch_classes(): + r"""Map from pitch class (str) to semitone (int).""" + pitch_classes = ["C", "D", "E", "F", "G", "A", "B"] + semitones = [0, 2, 4, 5, 7, 9, 11] + return {c: s for c, s in zip(pitch_classes, semitones)} + + +def _scale_degrees(): + r"""Map scale degrees (str) to semitones (int).""" + degrees = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13"] + semitones = [0, 2, 4, 5, 7, 9, 11, 12, 14, 16, 17, 19, 21] + return {d: s for d, s in zip(degrees, semitones)} + + +# Maps pitch classes (strings) to semitone indexes (ints). +PITCH_CLASSES = _pitch_classes() + + +def pitch_class_to_semitone(pitch_class): + r"""Convert a pitch class to semitone. + + Parameters + ---------- + pitch_class : str + Spelling of a given pitch class, e.g. 'C#', 'Gbb' + + Returns + ------- + semitone : int + Semitone value of the pitch class. + + """ + semitone = 0 + for idx, char in enumerate(pitch_class): + if char == "#" and idx > 0: + semitone += 1 + elif char == "b" and idx > 0: + semitone -= 1 + elif idx == 0: + semitone = PITCH_CLASSES.get(char) + else: + raise InvalidChordException( + "Pitch class improperly formed: %s" % pitch_class + ) + return semitone % 12 + + +# Maps scale degrees (strings) to semitone indexes (ints). +SCALE_DEGREES = _scale_degrees() + + +def scale_degree_to_semitone(scale_degree): + r"""Convert a scale degree to semitone. + + Parameters + ---------- + scale_degree : str + Spelling of a relative scale degree, e.g. 'b3', '7', '#5' + + Returns + ------- + semitone : int + Relative semitone of the scale degree, wrapped to a single octave + + Raises + ------ + InvalidChordException if `scale_degree` is invalid. + """ + semitone = 0 + offset = 0 + if scale_degree.startswith("#"): + offset = scale_degree.count("#") + scale_degree = scale_degree.strip("#") + elif scale_degree.startswith("b"): + offset = -1 * scale_degree.count("b") + scale_degree = scale_degree.strip("b") + + semitone = SCALE_DEGREES.get(scale_degree, None) + if semitone is None: + raise InvalidChordException( + "Scale degree improperly formed: {}, expected one of {}.".format( + scale_degree, list(SCALE_DEGREES.keys()) + ) + ) + return semitone + offset + + +def scale_degree_to_bitmap(scale_degree, modulo=False, length=BITMAP_LENGTH): + """Create a bitmap representation of a scale degree. + + Note that values in the bitmap may be negative, indicating that the + semitone is to be removed. + + Parameters + ---------- + scale_degree : str + Spelling of a relative scale degree, e.g. 'b3', '7', '#5' + modulo : bool, default=True + If a scale degree exceeds the length of the bit-vector, modulo the + scale degree back into the bit-vector; otherwise it is discarded. + length : int, default=12 + Length of the bit-vector to produce + + Returns + ------- + bitmap : np.ndarray, in [-1, 0, 1], len=`length` + Bitmap representation of this scale degree. + """ + sign = 1 + if scale_degree.startswith("*"): + sign = -1 + scale_degree = scale_degree.strip("*") + edit_map = [0] * length + sd_idx = scale_degree_to_semitone(scale_degree) + if sd_idx < length or modulo: + edit_map[sd_idx % length] = sign + return np.array(edit_map) + + +# Maps quality strings to bitmaps, corresponding to relative pitch class +# semitones, i.e. vector[0] is the tonic. +QUALITIES = { + # 1 2 3 4 5 6 7 + "maj": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0], + "min": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0], + "aug": [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0], + "dim": [1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0], + "sus4": [1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0], + "sus2": [1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], + "7": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0], + "maj7": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1], + "min7": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0], + "minmaj7": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1], + "maj6": [1, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0], + "min6": [1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0], + "dim7": [1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0], + "hdim7": [1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0], + "aug7": [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0], + "maj9": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1], + "min9": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0], + "9": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0], + "min11": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0], + "11": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0], + "maj13": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1], + "min13": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0], + "13": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0], + "1": [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "5": [1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0], + "": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], +} + + +def quality_to_bitmap(quality): + """Return the bitmap for a given quality. + + Parameters + ---------- + quality : str + Chord quality name. + + Returns + ------- + bitmap : np.ndarray + Bitmap representation of this quality (12-dim). + + """ + if quality not in QUALITIES: + raise InvalidChordException( + "Unsupported chord quality shorthand: '%s' " + "Did you mean to reduce extended chords?" % quality + ) + return np.array(QUALITIES[quality]) + + +# Maps extended chord qualities to the subset above, translating additional +# voicings to extensions as a set of scale degrees (strings). +# TODO(ejhumphrey): Revisit how minmaj7's are mapped. This is how TMC did it, +# but MMV handles it like a separate quality (rather than an add7). +EXTENDED_QUALITY_REDUX = { + "minmaj7": ("min", {"7"}), + "maj9": ("maj7", {"9"}), + "min9": ("min7", {"9"}), + "9": ("7", {"9"}), + "b9": ("7", {"b9"}), + "#9": ("7", {"#9"}), + "11": ("7", {"9", "11"}), + "#11": ("7", {"9", "#11"}), + "13": ("7", {"9", "11", "13"}), + "b13": ("7", {"9", "11", "b13"}), + "min11": ("min7", {"9", "11"}), + "maj13": ("maj7", {"9", "11", "13"}), + "min13": ("min7", {"9", "11", "13"}), +} + + +def reduce_extended_quality(quality): + """Map an extended chord quality to a simpler one, moving upper voices to + a set of scale degree extensions. + + Parameters + ---------- + quality : str + Extended chord quality to reduce. + + Returns + ------- + base_quality : str + New chord quality. + extensions : set + Scale degrees extensions for the quality. + + """ + return EXTENDED_QUALITY_REDUX.get(quality, (quality, set())) + + +# --- Chord Label Parsing --- +# This monster regexp is pulled from the JAMS chord namespace, +# which is in turn derived from the context-free grammar of +# Harte et al., 2005. +CHORD_RE = re.compile( + r"""^((N|X)|(([A-G](b*|#*))((:(maj|min|dim|aug|1|5|sus2|sus4|maj6|min6|7|maj7|min7|dim7|hdim7|minmaj7|aug7|9|maj9|min9|11|maj11|min11|13|maj13|min13)(\((\*?((b*|#*)([1-9]|1[0-3]?))(,\*?((b*|#*)([1-9]|1[0-3]?)))*)\))?)|(:\((\*?((b*|#*)([1-9]|1[0-3]?))(,\*?((b*|#*)([1-9]|1[0-3]?)))*)\)))?((/((b*|#*)([1-9]|1[0-3]?)))?)?))$""" +) # nopep8 + + +def validate_chord_label(chord_label): + """Test for well-formedness of a chord label. + + Parameters + ---------- + chord_label : str + Chord label to validate. + """ + if not CHORD_RE.match(chord_label): + raise InvalidChordException("Invalid chord label: " "{}".format(chord_label)) + pass + + +def split(chord_label, reduce_extended_chords=False): + """Parse a chord label into its four constituent parts: + - root + - quality shorthand + - scale degrees + - bass + + Note: Chords lacking quality AND interval information are major. + - If a quality is specified, it is returned. + - If an interval is specified WITHOUT a quality, the quality field is + empty. + + Some examples:: + + 'C' -> ['C', 'maj', {}, '1'] + 'G#:min(*b3,*5)/5' -> ['G#', 'min', {'*b3', '*5'}, '5'] + 'A:(3)/6' -> ['A', '', {'3'}, '6'] + + Parameters + ---------- + chord_label : str + A chord label. + reduce_extended_chords : bool + Whether to map the upper voicings of extended chords (9's, 11's, 13's) + to semitone extensions. (Default value = False) + + Returns + ------- + chord_parts : list + Split version of the chord label. + + """ + chord_label = str(chord_label) + validate_chord_label(chord_label) + if chord_label == NO_CHORD: + return [chord_label, "", set(), ""] + + bass = "1" + if "/" in chord_label: + chord_label, bass = chord_label.split("/") + + scale_degrees = set() + omission = False + if "(" in chord_label: + chord_label, scale_degrees = chord_label.split("(") + omission = "*" in scale_degrees + scale_degrees = scale_degrees.strip(")") + scale_degrees = {i.strip() for i in scale_degrees.split(",")} + + # Note: Chords lacking quality AND added interval information are major. + # If a quality shorthand is specified, it is returned. + # If an interval is specified WITHOUT a quality, the quality field is + # empty. + # Intervals specifying omissions MUST have a quality. + if omission and ":" not in chord_label: + raise InvalidChordException( + "Intervals specifying omissions MUST have a quality." + ) + quality = "" if scale_degrees else "maj" + if ":" in chord_label: + chord_root, quality_name = chord_label.split(":") + # Extended chords (with ":"s) may not explicitly have Major qualities, + # so only overwrite the default if the string is not empty. + if quality_name: + quality = quality_name.lower() + else: + chord_root = chord_label + + if reduce_extended_chords: + quality, addl_scale_degrees = reduce_extended_quality(quality) + scale_degrees.update(addl_scale_degrees) + + return [chord_root, quality, scale_degrees, bass] + + +def join(chord_root, quality="", extensions=None, bass=""): + r"""Join the parts of a chord into a complete chord label. + + Parameters + ---------- + chord_root : str + Root pitch class of the chord, e.g. 'C', 'Eb' + quality : str + Quality of the chord, e.g. 'maj', 'hdim7' + (Default value = '') + extensions : list + Any added or absent scaled degrees for this chord, e.g. ['4', '\*3'] + (Default value = None) + bass : str + Scale degree of the bass note, e.g. '5'. + (Default value = '') + + Returns + ------- + chord_label : str + A complete chord label. + + """ + chord_label = chord_root + if quality or extensions: + chord_label += ":%s" % quality + if extensions: + chord_label += "(%s)" % ",".join(extensions) + if bass and bass != "1": + chord_label += "/%s" % bass + validate_chord_label(chord_label) + return chord_label + + +# --- Chords to Numerical Representations --- +def encode(chord_label, reduce_extended_chords=False, strict_bass_intervals=False): + """Translate a chord label to numerical representations for evaluation. + + Parameters + ---------- + chord_label : str + Chord label to encode. + reduce_extended_chords : bool + Whether to map the upper voicings of extended chords (9's, 11's, 13's) + to semitone extensions. + (Default value = False) + strict_bass_intervals : bool + Whether to require that the bass scale degree is present in the chord. + (Default value = False) + + Returns + ------- + root_number : int + Absolute semitone of the chord's root. + semitone_bitmap : np.ndarray, dtype=int + 12-dim vector of relative semitones in the chord spelling. + bass_number : int + Relative semitone of the chord's bass note, e.g. 0=root, 7=fifth, etc. + """ + if chord_label == NO_CHORD: + return NO_CHORD_ENCODED + if chord_label == X_CHORD: + return X_CHORD_ENCODED + chord_root, quality, scale_degrees, bass = split( + chord_label, reduce_extended_chords=reduce_extended_chords + ) + + root_number = pitch_class_to_semitone(chord_root) + bass_number = scale_degree_to_semitone(bass) % 12 + + semitone_bitmap = quality_to_bitmap(quality) + semitone_bitmap[0] = 1 + + for scale_degree in scale_degrees: + semitone_bitmap += scale_degree_to_bitmap(scale_degree, reduce_extended_chords) + + semitone_bitmap = (semitone_bitmap > 0).astype(np.int64) + if not semitone_bitmap[bass_number] and strict_bass_intervals: + raise InvalidChordException( + "Given bass scale degree is absent from this chord: " "%s" % chord_label, + chord_label, + ) + else: + semitone_bitmap[bass_number] = 1 + return root_number, semitone_bitmap, bass_number + + +def encode_many(chord_labels, reduce_extended_chords=False): + """Translate a set of chord labels to numerical representations for sane + evaluation. + + Parameters + ---------- + chord_labels : list + Set of chord labels to encode. + reduce_extended_chords : bool + Whether to map the upper voicings of extended chords (9's, 11's, 13's) + to semitone extensions. + (Default value = False) + + Returns + ------- + root_number : np.ndarray, dtype=int + Absolute semitone of the chord's root. + interval_bitmap : np.ndarray, dtype=int + 12-dim vector of relative semitones in the given chord quality. + bass_number : np.ndarray, dtype=int + Relative semitones of the chord's bass notes. + + """ + num_items = len(chord_labels) + roots, basses = np.zeros([2, num_items], dtype=np.int64) + semitones = np.zeros([num_items, 12], dtype=np.int64) + local_cache = dict() + for i, label in enumerate(chord_labels): + result = local_cache.get(label, None) + if result is None: + result = encode(label, reduce_extended_chords) + local_cache[label] = result + roots[i], semitones[i], basses[i] = result + return roots, semitones, basses + + +def rotate_bitmap_to_root(bitmap, chord_root): + """Circularly shift a relative bitmap to its absolute pitch classes. + + For clarity, the best explanation is an example. Given 'G:Maj', the root + and quality map are as follows:: + + root=5 + quality=[1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0] # Relative chord shape + + After rotating to the root, the resulting bitmap becomes:: + + abs_quality = [0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1] # G, B, and D + + Parameters + ---------- + bitmap : np.ndarray, shape=(12,) + Bitmap of active notes, relative to the given root. + chord_root : int + Absolute pitch class number. + + Returns + ------- + bitmap : np.ndarray, shape=(12,) + Absolute bitmap of active pitch classes. + + """ + bitmap = np.asarray(bitmap) + assert bitmap.ndim == 1, "Currently only 1D bitmaps are supported." + idxs = list(np.nonzero(bitmap)) + idxs[-1] = (idxs[-1] + chord_root) % 12 + abs_bitmap = np.zeros_like(bitmap) + abs_bitmap[tuple(idxs)] = 1 + return abs_bitmap diff --git a/extensions_built_in/audio_models/yue2/src/shims/pretty_midi.py b/extensions_built_in/audio_models/yue2/src/shims/pretty_midi.py new file mode 100644 index 0000000..7ee2eb4 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/shims/pretty_midi.py @@ -0,0 +1,89 @@ +"""Minimal in-memory stand-in for the parts of pretty_midi that SheetSage2's ABC path uses: +PrettyMIDI(resolution=...) / PrettyMIDI(bytes-like), .instruments, .resolution, .write(stream); +Instrument(program, name=...) with .notes/.name/.program/.is_drum; Note(velocity, pitch, start, end). + +"Bytes" are our own serialization (the only consumer is this shim), but note times are rounded to +MIDI ticks on write exactly as a real file round trip would (120 bpm default tempo: ticks per second += 2 * resolution), because the ABC builder is documented to follow the serialized MIDI timing.""" + +import json +from io import BytesIO + +_MAGIC = b"AITK-MIDI-JSON\n" + + +class Note: + def __init__(self, velocity, pitch, start, end): + self.velocity = int(velocity) + self.pitch = int(pitch) + self.start = float(start) + self.end = float(end) + + def __repr__(self): + return f"Note(start={self.start:.6f}, end={self.end:.6f}, pitch={self.pitch}, velocity={self.velocity})" + + +class Instrument: + def __init__(self, program, is_drum=False, name=""): + self.program = int(program) + self.is_drum = bool(is_drum) + self.name = name + self.notes = [] + + +class PrettyMIDI: + def __init__(self, midi_file=None, resolution=220, initial_tempo=120.0): + self.resolution = int(resolution) + self.initial_tempo = float(initial_tempo) + self.instruments = [] + if midi_file is not None: + if isinstance(midi_file, (bytes, bytearray)): + data = bytes(midi_file) + elif hasattr(midi_file, "read"): + data = midi_file.read() + else: + with open(str(midi_file), "rb") as f: + data = f.read() + if not data.startswith(_MAGIC): + raise ValueError("pretty_midi shim can only read MIDI written by this shim") + payload = json.loads(data[len(_MAGIC) :].decode("utf-8")) + self.resolution = int(payload["resolution"]) + self.initial_tempo = float(payload["initial_tempo"]) + for inst in payload["instruments"]: + instrument = Instrument(inst["program"], inst["is_drum"], inst["name"]) + instrument.notes = [Note(n["velocity"], n["pitch"], n["start"], n["end"]) for n in inst["notes"]] + self.instruments.append(instrument) + + def _ticks_per_second(self): + return self.resolution * self.initial_tempo / 60.0 + + def _quantize(self, seconds): + tps = self._ticks_per_second() + return int(round(seconds * tps)) / tps + + def write(self, filename): + payload = { + "resolution": self.resolution, + "initial_tempo": self.initial_tempo, + "instruments": [ + { + "program": i.program, + "is_drum": i.is_drum, + "name": i.name, + "notes": [ + {"velocity": n.velocity, "pitch": n.pitch, "start": self._quantize(n.start), "end": self._quantize(n.end)} + for n in i.notes + ], + } + for i in self.instruments + ], + } + data = _MAGIC + json.dumps(payload).encode("utf-8") + if hasattr(filename, "write"): + filename.write(data) + else: + with open(str(filename), "wb") as f: + f.write(data) + + def get_end_time(self): + return max((n.end for i in self.instruments for n in i.notes), default=0.0) diff --git a/extensions_built_in/audio_models/yue2/src/tokenizer.py b/extensions_built_in/audio_models/yue2/src/tokenizer.py index 064d44f..392aab3 100644 --- a/extensions_built_in/audio_models/yue2/src/tokenizer.py +++ b/extensions_built_in/audio_models/yue2/src/tokenizer.py @@ -38,11 +38,26 @@ class YuE2TextTokenizer: def decode(self, ids: List[int]) -> str: return self.tokenizer.decode(ids, skip_special_tokens=True) + def prefix_head_ids(self, style: str, lyrics: str, cot: str = "off") -> List[int]: + """Prompt up to and including ABC_START (ComfyUI's ``prefix``); the ABC block and + ``[ABC_END, MUSIC_START]`` follow (see ``abc_tail_ids``).""" + prompt = f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n" + return [EOD] + self.encode(prompt) + [ABC_START] + + @staticmethod + def abc_tail_ids(abc_ids: List[int]) -> List[int]: + return list(abc_ids) + [ABC_END, MUSIC_START] + + def encode_abc(self, abc: str) -> List[int]: + ids = self.encode(abc) + if any(not 0 <= t < EOD for t in ids): + raise ValueError("ABC ids must stay inside the ordinary text vocabulary") + return ids + def prefix_ids(self, style: str, lyrics: str, cot: str = "off", abc: str = "") -> List[int]: """Everything the AR expert sees before the first codec token.""" - prompt = f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n" - abc_ids = [] if cot == "off" else self.encode(abc) - return [EOD] + self.encode(prompt) + [ABC_START] + abc_ids + [ABC_END, MUSIC_START] + abc_ids = [] if cot == "off" else self.encode_abc(abc) + return self.prefix_head_ids(style, lyrics, cot) + self.abc_tail_ids(abc_ids) class TokenizerHead(nn.Module): @@ -75,6 +90,66 @@ def _rebuild_rotary(model: nn.Module) -> int: return n +SHEETSAGE_REPO = "m-a-p/SheetSage2" +SHEETSAGE_FILE = "Comfy-Org/YuE2/audio_encoders/sheetsage2_bf16.safetensors" +# tied to token_embedding.weight (loaded); everything else in the Comfy file maps 1:1 onto upstream +NONPERSISTENT_OK = {"decoder.embed_tokens.weight"} +SHEETSAGE_SAMPLE_RATE = 24000 + + +class SheetSage2Transcriber: + """Audio -> ABC lead sheet with the released SheetSage2 (transformers remote code), the same + transcriber ComfyUI ships for cover mode. ``full`` keeps chords, ``melody`` drops them.""" + + def __init__(self, weights_path: str, repo: str = SHEETSAGE_REPO): + """``weights_path``: the Comfy repack ``audio_encoders/sheetsage2_bf16.safetensors`` (merged + weights); ``repo`` supplies config + code (transformers remote code) only.""" + from .shims import install_shims + + install_shims() # mir_eval.chord / pretty_midi stand-ins unless the real packages exist + import sys + + from safetensors.torch import load_file + from transformers import AutoConfig + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + from .sheetsage_decoder import ScoreDecoder + + 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 + 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()} + missing, unexpected = model.load_state_dict(sd, strict=False) + missing = [k for k in missing if k not in NONPERSISTENT_OK] + if missing or unexpected: + raise ValueError(f"SheetSage2 weights do not match: missing {missing[:5]} unexpected {unexpected[:5]}") + model.output_projection.weight = model.token_embedding.weight # keep the tie after loading + self.model = model.eval() + self.model.requires_grad_(False) + _rebuild_rotary(self.model) + + def to(self, device): + self.model.to(device) + return self + + @property + def device(self): + return next(self.model.parameters()).device + + @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").""" + 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'}") + return abc + + class SemanticTokenizer(nn.Module): """Waveform -> per-frame codec ids (0..32767) at 25 Hz.""" diff --git a/extensions_built_in/audio_models/yue2/yue2_model.py b/extensions_built_in/audio_models/yue2/yue2_model.py index d94a7d4..14b3f36 100644 --- a/extensions_built_in/audio_models/yue2/yue2_model.py +++ b/extensions_built_in/audio_models/yue2/yue2_model.py @@ -1,6 +1,7 @@ """YuE2 (m-a-p/YuE2-3B) for ai-toolkit. -Weights: the Comfy-Org all-in-one checkpoint (``checkpoints/yue2_3b_bf16.safetensors``) +Weights: the Comfy-Org all-in-one checkpoint (default ``checkpoints/yue2_3b_int8_convrot.safetensors``, +shipped int8 convrot layers attached as-is; ``yue2_3b_bf16.safetensors`` also loads) resolved under MODELS_PATH in ComfyUI's folder layout and downloaded there when missing. The official audio -> semantic-token tokenizer is unreleased; training conditioning comes from the community head (MERT-v2-FullSong + classifier) in @@ -15,8 +16,10 @@ flow loss on the real latents. With ``model_kwargs.ar_loss_weight`` > 0 the same also trains the AR expert with next-token cross entropy over the codec tokens, so one LoRA network on both experts learns composition and rendering together. -Captions use the ace-step tag format: (style), , (max seconds -for samples). +Prompt modes follow ComfyUI (``model_kwargs.cot``): "full" (default, Generate ABC path: the AR +writes a chord-annotated ABC sheet, then the codec tokens), "melody" (sheet without chords) or +"off" (no sheet; what Comfy runs when the ABC input is empty). For full/melody the training sheet +comes from SheetSage2 at cache time and rides the latent cache as ``abc_ids``. """ import os @@ -25,6 +28,7 @@ import re from typing import List, Optional import torch +import torch.utils.checkpoint from safetensors.torch import load_file from tqdm import tqdm @@ -38,12 +42,15 @@ from toolkit.print import print_acc from toolkit.prompt_utils import PromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler -from .src.model import CODEC_OFFSET, FRAMES_PER_SECOND, MUSIC_END, YuE2Model, merge_nar_lora +from .src.model import ABC_END, CODEC_OFFSET, CONTEXT, FRAMES_PER_SECOND, INSTRUCTIONS, MUSIC_END, MUSIC_START, YuE2Model, merge_nar_lora from .src.pipeline import YuE2Pipeline -from .src.tokenizer import HEAD_FILE, HEAD_REPO, NAR_LORA_FILE, SemanticTokenizer, YuE2TextTokenizer +from .src.tokenizer import HEAD_FILE, HEAD_REPO, NAR_LORA_FILE, SHEETSAGE_FILE, SemanticTokenizer, SheetSage2Transcriber, YuE2TextTokenizer from .src.vae import SAMPLE_RATE, YuE2VAE -DEFAULT_CHECKPOINT = "Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors" +COT_CODES = {"off": 0, "melody": 1, "full": 2} +COT_NAMES = {v: k for k, v in COT_CODES.items()} + +DEFAULT_CHECKPOINT = "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors" # community tokenizer files have no ComfyUI folder; they live under /ai_toolkit/ AI_TOOLKIT_DIR = "ai_toolkit" @@ -130,10 +137,10 @@ def parse_caption(text: str) -> dict: class YuE2TextEncoder(torch.nn.Module): - """Prompt prefix -> AR token embeddings. Holds no weights of its own; the - embedding table lives in the AR expert (looked up on whatever device it is on).""" + """Prompt -> AR token embeddings of the prefix head (instruction, tags, lyrics, ABC_START). + Holds no weights of its own; the embedding table lives in the AR expert.""" - def __init__(self, model: YuE2Model, tokenizer: YuE2TextTokenizer, cot: str = "off"): + def __init__(self, model: YuE2Model, tokenizer: YuE2TextTokenizer, cot: str = "full"): super().__init__() self._model = [model] self.tokenizer = tokenizer @@ -146,14 +153,11 @@ class YuE2TextEncoder(torch.nn.Module): @property def dtype(self): - return self._model[0].ar.model.embed_tokens.weight.dtype - - def prefix_ids(self, style: str, lyrics: str) -> List[int]: - return self.tokenizer.prefix_ids(style, lyrics, cot=self.cot) + return self._model[0].ar.dtype @torch.no_grad() def forward(self, style: str, lyrics: str) -> torch.Tensor: - ids = torch.tensor([self.prefix_ids(style, lyrics)], dtype=torch.long) + ids = torch.tensor([self.tokenizer.prefix_head_ids(style, lyrics, cot=self.cot)], dtype=torch.long) return self._model[0].ar.embed(ids) @@ -182,6 +186,15 @@ class YuE2AudioModel(BaseAudioModel): self.merge_nar_lora = bool(kw.get("merge_nar_lora", False)) self.nar_lora_path = kw.get("nar_lora_path", f"{HEAD_REPO}/{NAR_LORA_FILE}") self.semantic_head_path = kw.get("semantic_head_path", f"{HEAD_REPO}/{HEAD_FILE}") + # ComfyUI's Generate ABC path: "full" (chords) or "melody"; "off" skips the sheet (empty ABC in Comfy) + self.cot = str(kw.get("cot", "full")) + if self.cot not in INSTRUCTIONS: + raise ValueError(f"model_kwargs.cot must be one of {list(INSTRUCTIONS)}") + self.abc_max_tokens = int(kw.get("abc_max_tokens", 8192)) + # fraction of training items fed without the sheet (off-mode prefix), so one LoRA serves both Comfy paths + self.abc_dropout = float(kw.get("abc_dropout", 0.5)) + self.sheetsage_path = kw.get("sheetsage_path", SHEETSAGE_FILE) + self.transcriber: Optional[SheetSage2Transcriber] = None self.sample_max_seconds = float(kw.get("sample_max_seconds", 120.0)) self.sample_ar_temperature = float(kw.get("sample_ar_temperature", 1.0)) # reference sampler value; 1.0 lets a memorized song replay (the penalty knocks it off path) @@ -230,7 +243,7 @@ class YuE2AudioModel(BaseAudioModel): flush() self.model.to(device) self.vae.to(self.vae_device_torch) - self.text_encoder = YuE2TextEncoder(self.model, self.tokenizer) + self.text_encoder = YuE2TextEncoder(self.model, self.tokenizer, cot=self.cot) self.pipeline = YuE2Pipeline(self.model, self.vae) self.pipeline.do_tiled_decoding = True self.pipeline.use_cuda_graphs = bool(self.model_config.model_kwargs.get("ar_cuda_graphs", True)) @@ -243,6 +256,12 @@ class YuE2AudioModel(BaseAudioModel): self.semantic_tokenizer = SemanticTokenizer(head_path, debug=self.debug) return self.semantic_tokenizer + def _get_transcriber(self) -> SheetSage2Transcriber: + if self.transcriber is None: + print_acc("Loading SheetSage2 (audio -> ABC)") + self.transcriber = SheetSage2Transcriber(resolve_named_file(self.sheetsage_path, component="sheetsage2")) + return self.transcriber + def get_transformer_block_names(self) -> Optional[List[str]]: return ["model.layers"] @@ -269,7 +288,8 @@ class YuE2AudioModel(BaseAudioModel): return PromptEmbeds(torch.cat(embeds, 0), attention_mask=torch.cat(masks, 0)) def encode_audio(self, audio_tensor: torch.Tensor, device=None, dtype=None): - """[B, 2, samples] at 48 kHz -> DTO [B, T, 64] with codec ``tokens`` [B, T].""" + """[B, 2, samples] at 48 kHz -> DTO [B, T, 64] with codec ``tokens`` [B, T] and, unless cot is + "off", the SheetSage2 sheet as ``abc_ids`` [B, N] (text ids, -1 padded).""" if device is None: device = self.vae_device_torch if self.vae.device == torch.device("cpu"): @@ -282,7 +302,20 @@ class YuE2AudioModel(BaseAudioModel): tokens = torch.stack([tokenizer.tokenize(wav, SAMPLE_RATE) for wav in audio_tensor]) # [B, T'] n = min(latents.shape[1], tokens.shape[1]) latents = latents[:, :n].to(dtype=self.torch_dtype if dtype is None else dtype) - return DTO(latents, tokens=tokens[:, :n].to(latents.device, torch.int32)) + extras = {"tokens": tokens[:, :n].to(latents.device, torch.int32)} + if self.cot != "off": + transcriber = self._get_transcriber() + if transcriber.device != self.device_torch: + transcriber.to(self.device_torch) + sheets = [self.tokenizer.encode_abc(transcriber.transcribe(wav, SAMPLE_RATE, cot=self.cot)) for wav in audio_tensor] + width = max(len(x) for x in sheets) + abc = torch.full((len(sheets), width), -1, dtype=torch.int32) + for i, ids in enumerate(sheets): + abc[i, : len(ids)] = torch.tensor(ids, dtype=torch.int32) + extras["abc_ids"] = abc.to(latents.device) + # which sheet flavor the cache holds; training refuses a mismatch instead of silently using it + extras["abc_mode"] = torch.full((len(sheets),), COT_CODES[self.cot], dtype=torch.int32, device=latents.device) + return DTO(latents, **extras) # ------------------------------------------------------------------ # training @@ -330,42 +363,67 @@ class YuE2AudioModel(BaseAudioModel): length = int((prefix_mask > 0.5).sum().item()) return prefix_embeds[:length] - def _ar_info(self, logits: torch.Tensor, ids: torch.Tensor, prefix: torch.Tensor, total: int, whole_song: bool) -> str: - """Readout diagnostics: what the AR loss was computed over.""" - if not self.debug or self._loss_log_every <= 0 or (self._loss_log_step + 1) % self._loss_log_every != 0: - return "" - with torch.no_grad(): - p = logits.softmax(-1) - p_target = p.gather(1, ids[:, None]).squeeze(1) - return (f"[ar targets {ids.shape[0]} unique {ids.unique().numel()} (song {total}, whole {whole_song}) " - f"prefix {prefix.shape[0]} p(END|prefix) {p[0, MUSIC_END].item():.3f} p(target) median {p_target.median().item():.3f}]") - - def _ar_kl(self, ar_embeds: torch.Tensor, logits: torch.Tensor): - """KL(base || lora) per position, mean. Base pass = same prefill with the LoRA inactive, no grad.""" - net = getattr(self, "_network", None) - n = logits.shape[0] - was_active = getattr(net, "is_active", None) - if net is not None: - net.is_active = False - try: - with torch.no_grad(): - _, base_hidden = self.model.ar.prefill(ar_embeds, return_hidden=True) - base_hidden = base_hidden[0, -n - 1 : -1] - finally: + def _ar_losses(self, ar_embeds: torch.Tensor, ids: torch.Tensor, prefix: torch.Tensor, total: int, whole_song: bool): + """Next-token CE over ``ids`` (and KL(base || lora) when ``ar_kl_weight`` > 0), computed in + position chunks with recomputed logits so the fp32 [n, vocab] tensors never exist at once. + Returns (ce, kl_or_None, info_str).""" + model = self.model + n = ids.shape[0] + _, hidden = model.ar.prefill(ar_embeds, return_hidden=True) + hidden = hidden[0, -n - 1 : -1] + base_hidden = None + if self.ar_kl_weight > 0: + net = getattr(self, "_network", None) + was_active = getattr(net, "is_active", None) if net is not None: - net.is_active = was_active - kl = logits.new_zeros(()) - for s0 in range(0, n, 1024): # chunked: two fp32 [n, vocab] log-softmaxes at once would be ~10 GB - with torch.no_grad(): - base_logp = torch.log_softmax(self.model.ar.model.lm_head(base_hidden[s0 : s0 + 1024]).float(), -1) - kl = kl + torch.nn.functional.kl_div(torch.log_softmax(logits[s0 : s0 + 1024], -1), base_logp, log_target=True, reduction="sum") - return kl / n + net.is_active = False + try: + with torch.no_grad(): + _, base_hidden = model.ar.prefill(ar_embeds, return_hidden=True) + base_hidden = base_hidden[0, -n - 1 : -1] + finally: + if net is not None: + net.is_active = was_active + want_info = self.debug and self._loss_log_every > 0 and (self._loss_log_step + 1) % self._loss_log_every == 0 + lm_head = model.ar.model.lm_head + ce_sum = hidden.new_zeros((), dtype=torch.float32) + kl_sum = hidden.new_zeros((), dtype=torch.float32) if base_hidden is not None else None + p_target, p_end_first = [], None + chunk = 2048 + for s0 in range(0, n, chunk): + h = hidden[s0 : s0 + chunk] + tgt = ids[s0 : s0 + chunk] + if torch.is_grad_enabled(): + logits = torch.utils.checkpoint.checkpoint(lambda x: lm_head(x).float(), h, use_reentrant=False) + else: + logits = lm_head(h).float() + ce_sum = ce_sum + torch.nn.functional.cross_entropy(logits, tgt, reduction="sum") + if base_hidden is not None: + with torch.no_grad(): + base_logp = torch.log_softmax(lm_head(base_hidden[s0 : s0 + chunk]).float(), -1) + kl_sum = kl_sum + torch.nn.functional.kl_div(torch.log_softmax(logits, -1), base_logp, log_target=True, reduction="sum") + if want_info: + with torch.no_grad(): + p = logits.softmax(-1) + if s0 == 0: + p_end_first = p[0, MUSIC_END].item() + p_target.append(p.gather(1, tgt[:, None]).squeeze(1)) + del logits + info = "" + if want_info: + p_target = torch.cat(p_target) + info = (f"[ar targets {n} unique {ids.unique().numel()} (song {total}, whole {whole_song}) " + f"prefix {prefix.shape[0]} p(END|prefix) {p_end_first:.3f} p(target) median {p_target.median().item():.3f}]") + return ce_sum / n, (kl_sum / n if kl_sum is not None else None), info - def _ar_inputs(self, prefix: torch.Tensor, tokens: torch.Tensor, end_token: bool = True): - """One item: prefix (unpadded [L, H]) + codec tokens [+ MUSIC_END] -> embeds [1, L', H], ids [n].""" - ids = tokens.long() + CODEC_OFFSET + def _ar_inputs(self, prefix: torch.Tensor, abc_ids: torch.Tensor, tokens: torch.Tensor, end_token: bool = True): + """One item: prefix head (unpadded [L, H]) + abc + [ABC_END, MUSIC_START] + codec tokens [+ MUSIC_END] + -> embeds [1, L', H] and the ids after the head [n] (the AR targets).""" + device = tokens.device + parts = [abc_ids.long().to(device), torch.tensor([ABC_END, MUSIC_START], device=device), tokens.long() + CODEC_OFFSET] if end_token: - ids = torch.cat([ids, torch.tensor([MUSIC_END], device=tokens.device)]) + parts.append(torch.tensor([MUSIC_END], device=device)) + ids = torch.cat(parts) embeds = torch.cat([prefix.to(self.model.ar.dtype), self.model.ar.embed(ids)], dim=0)[None] return embeds, ids @@ -416,6 +474,16 @@ class YuE2AudioModel(BaseAudioModel): # `batch` must be a named parameter: predict_noise only forwards it when it sees it in the signature if batch is None or not isinstance(batch.latents, DTO) or batch.latents.get("tokens") is None: raise ValueError("YuE2 training needs codec tokens in the latent cache; enable latent caching") + abc_all = batch.latents.get("abc_ids") + if self.cot != "off": + mode = batch.latents.get("abc_mode") + cached = None if mode is None else int(mode.reshape(-1)[0].item()) + if abc_all is None or cached != COT_CODES[self.cot]: + raise ValueError( + f"YuE2 cot={self.cot!r} needs a matching ABC sheet in the latent cache " + f"(cache has {'none' if cached is None else COT_NAMES.get(cached, cached)!r}): " + "delete the dataset's _latent_cache folder so it is rebuilt" + ) start, end = getattr(batch, "yue2_window", (0, latent_model_input.shape[1])) tokens_all = batch.latents.tokens model = self.model @@ -426,35 +494,48 @@ class YuE2AudioModel(BaseAudioModel): prefix_embeds = text_embeddings.text_embeds.to(device, self.torch_dtype) prefix_mask = text_embeddings.attention_mask.to(device) train_ar = self.ar_loss_weight > 0 and torch.is_grad_enabled() + captions = batch.get_caption_list() if hasattr(batch, "get_caption_list") else None preds, lm_losses, kl_losses = [], [], [] for i in range(latent_model_input.shape[0]): song = tokens_all[i].to(device) total = song.shape[0] prefix = self._train_prefix(prefix_embeds[i], prefix_mask[i]) - # NAR conditioning follows the released chunk protocol: prefix + window + MUSIC_END - embeds, ids = self._ar_inputs(prefix, song[start:end]) + if self.cot == "off" or abc_all is None: + abc = torch.zeros(0, dtype=torch.long, device=device) + elif torch.is_grad_enabled() and captions is not None and random.random() < self.abc_dropout: + # sheet dropped: the item trains exactly as an off-mode prompt (off instruction, empty ABC block) + parsed = parse_caption(captions[i]) + ids = torch.tensor([self.tokenizer.prefix_head_ids(parsed["style"], parsed["lyrics"], cot="off")], device=device) + prefix = model.ar.embed(ids)[0].to(self.torch_dtype) + abc = torch.zeros(0, dtype=torch.long, device=device) + else: + row = abc_all[i].to(device) + abc = row[row >= 0].long() + # NAR conditioning follows the released chunk protocol: prefix + abc + window + MUSIC_END + embeds, ids = self._ar_inputs(prefix, abc, song[start:end]) whole_song = start == 0 and end == total if train_ar and not whole_song: # the language-model loss must see the song from its beginning, or the AR # learns that lyrics and tokens can start anywhere relative to each other limit = total if self.ar_max_tokens <= 0 else min(total, self.ar_max_tokens) - ar_embeds, ar_ids = self._ar_inputs(prefix, song[:limit], end_token=limit == total) - _, hidden = model.ar.prefill(ar_embeds, return_hidden=True) - logits = model.ar.model.lm_head(hidden[0, -ar_ids.shape[0] - 1 : -1]).float() - lm_losses.append(torch.nn.functional.cross_entropy(logits, ar_ids)) - if self.ar_kl_weight > 0: - kl_losses.append(self._ar_kl(ar_embeds, logits)) - self._pending_aux_info = self._ar_info(logits, ar_ids, prefix, total, whole_song) + ar_embeds, ar_ids = self._ar_inputs(prefix, abc, song[:limit], end_token=limit == total) + ce, kl, info = self._ar_losses(ar_embeds, ar_ids, prefix, total, whole_song) + lm_losses.append(ce) + if kl is not None: + kl_losses.append(kl) + self._pending_aux_info = info + with torch.no_grad(): + cache, _ = model.ar.prefill(embeds) + elif train_ar: + ce, kl, info = self._ar_losses(embeds, ids, prefix, total, whole_song) + lm_losses.append(ce) + if kl is not None: + kl_losses.append(kl) + self._pending_aux_info = info with torch.no_grad(): cache, _ = model.ar.prefill(embeds) else: - cache, hidden = model.ar.prefill(embeds, return_hidden=train_ar) - if train_ar: - logits = model.ar.model.lm_head(hidden[0, -ids.shape[0] - 1 : -1]).float() - lm_losses.append(torch.nn.functional.cross_entropy(logits, ids)) - if self.ar_kl_weight > 0: - kl_losses.append(self._ar_kl(embeds, logits)) - self._pending_aux_info = self._ar_info(logits, ids, prefix, total, whole_song) + cache, _ = model.ar.prefill(embeds) # the flow loss must not train the AR through its KV cache: only next-token CE shapes the AR cache = [(k.detach(), v.detach()) for k, v in cache] pred = model.nar(latent_model_input[i : i + 1].to(device, self.torch_dtype), t[i : i + 1], cache, embeds.shape[1]) @@ -484,10 +565,28 @@ class YuE2AudioModel(BaseAudioModel): max_seconds = getattr(gen_config, "duration", None) or parsed["duration"] or self.sample_max_seconds max_tokens = max(1, int(round(max_seconds * FRAMES_PER_SECOND))) length = int((conditional_embeds.attention_mask[0] > 0.5).sum().item()) - prefix = conditional_embeds.text_embeds[:1, :length].to(self.device_torch, self.torch_dtype) + head = conditional_embeds.text_embeds[:1, :length].to(self.device_torch, self.torch_dtype) + + # stage 1 (cot full/melody): the AR writes the ABC sheet after ABC_START, as Comfy's Generate ABC does + abc_ids: List[int] = [] + if self.cot != "off": + self._status_update("YuE2: generating ABC sheet") + abc_bar = tqdm(total=self.abc_max_tokens, desc=" abc tokens", unit="tok", position=1, leave=False, dynamic_ncols=True) + try: + abc_ids = pipeline.generate_abc_tokens( + head, seed=gen_config.seed, max_tokens=self.abc_max_tokens, progress=lambda done, total: abc_bar.update(done - abc_bar.n) + ) + finally: + abc_bar.close() + tail = torch.tensor([self.tokenizer.abc_tail_ids(abc_ids)], device=self.device_torch) + prefix = torch.cat([head, self.model.ar.embed(tail).to(head.dtype)], dim=1) + # one acoustic frame needs two positions plus the boundary tokens (Comfy's budget rule) + max_tokens = min(max_tokens, CONTEXT - prefix.shape[1] - 5) + if max_tokens < 1: + raise ValueError("YuE2 prompt leaves no room for music; shorten the style, lyrics, or ABC") # nested bars under generate_images' "Generating Samples" bar (position 0) - self._status_update(f"YuE2: generating codec tokens (up to {max_seconds:.0f}s)") + self._status_update(f"YuE2: generating codec tokens (up to {max_tokens / FRAMES_PER_SECOND:.0f}s)") token_bar = tqdm(total=max_tokens, desc=" codec tokens", unit="tok", position=1, leave=False, dynamic_ncols=True) def token_progress(done, total): @@ -497,7 +596,7 @@ class YuE2AudioModel(BaseAudioModel): try: codec = pipeline.generate_codec_tokens( prefix, max_tokens=max_tokens, seed=gen_config.seed, temperature=self.sample_ar_temperature, - repetition_penalty=self.sample_ar_repetition_penalty, progress=token_progress, + repetition_penalty=self.sample_ar_repetition_penalty, progress=token_progress, legacy_off=self.cot == "off", ) finally: token_bar.close() diff --git a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py index 4f64aba..772acbf 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py @@ -14,6 +14,7 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from safetensors.torch import load_file, save_file +from toolkit.util.streamed_safetensors import save_file_streamed from toolkit.accelerator import unwrap_model from optimum.quanto import freeze @@ -502,7 +503,8 @@ class HidreamO1Model(BaseModel): if not output_path.endswith(".safetensors"): output_path += ".safetensors" meta = get_meta_for_safetensors(meta, name=self.arch) - save_file(save_dict, output_path, metadata=meta) + # sequential writer: mmap save_file fails with EINVAL on ntfs3 for large files + save_file_streamed(save_dict, output_path, metadata=meta) else: transformer.save_pretrained( save_directory=output_path, diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 50f1bfd..a928639 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -193,7 +193,7 @@ MODEL_TESTS = { }, "yue2": { # untagged test prompt = style only; keep the AR budget short for the smoke run - "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors", "quantize": True, "model_kwargs": {"sample_max_seconds": 15}}, + "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8", "model_kwargs": {"sample_max_seconds": 15, "cot": "off"}}, "sample": {"width": 512, "height": 512, "num_inference_steps": 16, "guidance_scale": 1.0, "seed": 42}, }, "f-lite": { diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index 62bbeee..18d90dd 100644 --- a/toolkit/models/registry.py +++ b/toolkit/models/registry.py @@ -55,7 +55,7 @@ ARCH_REGISTRY: Dict[str, dict] = { "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}}, - "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors", "quantize": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}}, + "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}}, } OUTPUT_EXT = {"image": "png", "video": "mp4", "audio": "mp3"} diff --git a/toolkit/util/comfy_quant_import.py b/toolkit/util/comfy_quant_import.py index 6cd16b6..4634f02 100644 --- a/toolkit/util/comfy_quant_import.py +++ b/toolkit/util/comfy_quant_import.py @@ -247,7 +247,31 @@ def import_comfy_quantized_layers( ) parent_path, _, attr = module_path.rpartition(".") parent = root.get_submodule(parent_path) if parent_path else root - setattr(parent, attr, Int8Embedding(weight, weight_scale, orig_dtype)) + if conf.get("convrot"): + # rows are stored rotated like a linear's weight: attach to a stand-in linear + # (same shape) so the convrot backend undoes the rotation, then keep the table + # as a plain embedding (per-row lookups cannot apply the group rotation) + stand_in = torch.nn.Linear( + module.embedding_dim, module.num_embeddings, bias=False, device="meta" + ) + quantizer = get_ostris_quantizer("convrot8") + stand_in.register_buffer("cr8_qdata", weight.contiguous(), persistent=False) + stand_in.register_buffer( + "cr8_scales", + weight_scale.detach().float().reshape(-1).contiguous().view(torch.uint8), + persistent=False, + ) + stand_in.cr8_rot_size = int(conf.get("convrot_groupsize", 256)) + _to_ostris(stand_in, quantizer, orig_dtype) + table = quantizer.dequantize_folded(stand_in).to(orig_dtype) + new = torch.nn.Embedding( + module.num_embeddings, module.embedding_dim, padding_idx=module.padding_idx, + device="meta", dtype=orig_dtype, + ) + new.weight = torch.nn.Parameter(table, requires_grad=False) + setattr(parent, attr, new) + else: + setattr(parent, attr, Int8Embedding(weight, weight_scale, orig_dtype)) converted += 1 continue diff --git a/toolkit/util/streamed_safetensors.py b/toolkit/util/streamed_safetensors.py new file mode 100644 index 0000000..7255cbf --- /dev/null +++ b/toolkit/util/streamed_safetensors.py @@ -0,0 +1,62 @@ +import json +import struct +from typing import Dict, Optional + +import torch + +# safetensors dtype tags +_DTYPE_TAGS = { + torch.float64: "F64", + torch.float32: "F32", + torch.float16: "F16", + torch.bfloat16: "BF16", + torch.int64: "I64", + torch.int32: "I32", + torch.int16: "I16", + torch.int8: "I8", + torch.uint8: "U8", + torch.bool: "BOOL", +} + +_CHUNK_BYTES = 256 * 1024 * 1024 + + +def save_file_streamed( + tensors: Dict[str, torch.Tensor], + path: str, + metadata: Optional[Dict[str, str]] = None, +): + """Write a safetensors file with sequential writes instead of ftruncate+mmap. + Some filesystems (kernel ntfs3) reject one large ftruncate; incremental writes succeed.""" + header = {} + offset = 0 + flat = {} + for name, t in tensors.items(): + if t.dtype not in _DTYPE_TAGS: + raise ValueError(f"Unsupported dtype {t.dtype} for tensor {name}") + t = t.detach().to("cpu").contiguous() + nbytes = t.numel() * t.element_size() + header[name] = { + "dtype": _DTYPE_TAGS[t.dtype], + "shape": list(t.shape), + "data_offsets": [offset, offset + nbytes], + } + flat[name] = t + offset += nbytes + if metadata is not None: + header["__metadata__"] = {str(k): str(v) for k, v in metadata.items()} + + header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8") + # pad header to 8-byte alignment like the reference writer + pad = (8 - len(header_bytes) % 8) % 8 + header_bytes += b" " * pad + + with open(path, "wb") as f: + f.write(struct.pack(" +

+ Training matches Comfy's Generate ABC path (cot: full in model kwargs): at cache time every song is + transcribed to an ABC lead sheet with SheetSage2 (the Comfy-Org/YuE2 audio_encoders weights, no + extra packages), the sheet is stored with the latents, and the AR learns + lyrics to sheet and sheet to tokens. abc_dropout (default 0.5) is the fraction of training steps fed + without the sheet as an off-mode prompt, so the same LoRA also works when Comfy's ABC input is empty. cot: melody uses sheets without chords; cot: off skips the sheet, which is what + Comfy runs when its ABC input is left empty. Changing the mode requires re-caching the dataset. +

Prompt format: a style line, then [Lyrics], then the lyrics with bracketed section headers such as [Verse 1] and [Chorus]. The Qwen3-Omni captioner has a YuE2 preset that writes @@ -1454,7 +1463,8 @@ export const modelArchs: ModelArch[] = [

), - disableSections: ['network.conv'], + // no separate text encoder: the prompt side is the AR expert, covered by the transformer quantization + disableSections: ['network.conv', 'model.quantize_te'], additionalSections: ['model.low_vram', 'sample.duration'], }, {