Load sage sheet and use it on training with 50% dropout so it generates sheet music and handles without it.

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

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
+    if phase == "abc":
-    allowed[MUSIC_END] = 0
+        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 if legacy_off else 1] = False
-        removed[:3] = 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._step()
-            self.graph.replay()
-        else:
-            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,
+        temperature: float,
-        top_p: float = 0.95,
+        top_p: float,
-        top_k: int = 100,
+        top_k: int,
-        repetition_penalty: float = 1.2,
+        repetition_penalty: float,
-        penalty_window: int = 50,
+        penalty_window: int,
-        min_tokens: int = 200,
+        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(

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,
+        )

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

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

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)

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(abc)
-        abc_ids = [] if cot == "off" else self.encode(abc)
+        return self.prefix_head_ids(style, lyrics, cot) + self.abc_tail_ids(abc_ids)
-        return [EOD] + self.encode(prompt) + [ABC_START] + abc_ids + [ABC_END, MUSIC_START]
 
 
 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."""
 

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: <CAPTION> (style), <LYRICS>, <DURATION> (max seconds
+Prompt modes follow ComfyUI (``model_kwargs.cot``): "full" (default, Generate ABC path: the AR
-for samples).
+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 <MODELS_PATH>/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
+    """Prompt -> AR token embeddings of the prefix head (instruction, tags, lyrics, ABC_START).
-    embedding table lives in the AR expert (looked up on whatever device it is on)."""
+    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
+        return self._model[0].ar.dtype
-
-    def prefix_ids(self, style: str, lyrics: str) -> List[int]:
-        return self.tokenizer.prefix_ids(style, lyrics, cot=self.cot)
 
     @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:
+    def _ar_losses(self, ar_embeds: torch.Tensor, ids: torch.Tensor, prefix: torch.Tensor, total: int, whole_song: bool):
-        """Readout diagnostics: what the AR loss was computed over."""
+        """Next-token CE over ``ids`` (and KL(base || lora) when ``ar_kl_weight`` > 0), computed in
-        if not self.debug or self._loss_log_every <= 0 or (self._loss_log_step + 1) % self._loss_log_every != 0:
+        position chunks with recomputed logits so the fp32 [n, vocab] tensors never exist at once.
-            return ""
+        Returns (ce, kl_or_None, info_str)."""
-        with torch.no_grad():
+        model = self.model
-            p = logits.softmax(-1)
+        n = ids.shape[0]
-            p_target = p.gather(1, ids[:, None]).squeeze(1)
+        _, hidden = model.ar.prefill(ar_embeds, return_hidden=True)
-            return (f"[ar targets {ids.shape[0]} unique {ids.unique().numel()} (song {total}, whole {whole_song}) "
+        hidden = hidden[0, -n - 1 : -1]
-                    f"prefix {prefix.shape[0]} p(END|prefix) {p[0, MUSIC_END].item():.3f} p(target) median {p_target.median().item():.3f}]")
+        base_hidden = None
-
+        if self.ar_kl_weight > 0:
-    def _ar_kl(self, ar_embeds: torch.Tensor, logits: torch.Tensor):
+            net = getattr(self, "_network", None)
-        """KL(base || lora) per position, mean. Base pass = same prefill with the LoRA inactive, no grad."""
+            was_active = getattr(net, "is_active", None)
-        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:
             if net is not None:
-                net.is_active = was_active
+                net.is_active = False
-        kl = logits.new_zeros(())
+            try:
-        for s0 in range(0, n, 1024):  # chunked: two fp32 [n, vocab] log-softmaxes at once would be ~10 GB
+                with torch.no_grad():
-            with torch.no_grad():
+                    _, base_hidden = model.ar.prefill(ar_embeds, return_hidden=True)
-                base_logp = torch.log_softmax(self.model.ar.model.lm_head(base_hidden[s0 : s0 + 1024]).float(), -1)
+                    base_hidden = base_hidden[0, -n - 1 : -1]
-            kl = kl + torch.nn.functional.kl_div(torch.log_softmax(logits[s0 : s0 + 1024], -1), base_logp, log_target=True, reduction="sum")
+            finally:
-        return kl / n
+                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):
+    def _ar_inputs(self, prefix: torch.Tensor, abc_ids: 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]."""
+        """One item: prefix head (unpadded [L, H]) + abc + [ABC_END, MUSIC_START] + codec tokens [+ MUSIC_END]
-        ids = tokens.long() + CODEC_OFFSET
+        -> 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
+            if self.cot == "off" or abc_all is None:
-            embeds, ids = self._ar_inputs(prefix, song[start:end])
+                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)
+                ar_embeds, ar_ids = self._ar_inputs(prefix, abc, song[:limit], end_token=limit == total)
-                _, hidden = model.ar.prefill(ar_embeds, return_hidden=True)
+                ce, kl, info = self._ar_losses(ar_embeds, ar_ids, prefix, total, whole_song)
-                logits = model.ar.model.lm_head(hidden[0, -ar_ids.shape[0] - 1 : -1]).float()
+                lm_losses.append(ce)
-                lm_losses.append(torch.nn.functional.cross_entropy(logits, ar_ids))
+                if kl is not None:
-                if self.ar_kl_weight > 0:
+                    kl_losses.append(kl)
-                    kl_losses.append(self._ar_kl(ar_embeds, logits))
+                self._pending_aux_info = info
-                self._pending_aux_info = self._ar_info(logits, ar_ids, prefix, total, whole_song)
+                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)
+                cache, _ = model.ar.prefill(embeds)
-                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)
             # 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()

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,

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": {

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

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
 

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("<Q", len(header_bytes)))
+        f.write(header_bytes)
+        for name, t in flat.items():
+            if t.numel() == 0:
+                continue
+            buf = memoryview(t.reshape(-1).view(torch.uint8).numpy())
+            for start in range(0, len(buf), _CHUNK_BYTES):
+                f.write(buf[start:start + _CHUNK_BYTES])

ui/src/app/jobs/new/options.tsx

@@ -1385,7 +1385,7 @@ export const modelArchs: ModelArch[] = [
     defaults: {
       // default updates when [selected, unselected] in the UI
       'config.process[0].model.name_or_path': [
-        'Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors',
+        'Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors',
         defaultNameOrPath,
       ],
       'config.process[0].model.quantize': [true, false],
@@ -1394,12 +1394,13 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].train.unload_text_encoder': [false, false],
       'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
       'config.process[0].train.timestep_type': ['sigmoid', 'sigmoid'],
-      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+      // the int8 repack ships convrot8 layers; requesting convrot8 keeps them as-is (no requantization)
+      'config.process[0].model.qtype': ['convrot8', 'qfloat8'],
       'config.process[0].sample': [defaultYue2SampleConfig, defaultSampleConfig],
       'config.process[0].datasets[x].cache_latents_to_disk': [true, true],
       // blank captions break lyric following; the AR must always see the prefix
       'config.process[0].datasets[x].caption_dropout_rate': [0, 0.05],
-      'config.process[0].model.model_kwargs': [{ sample_ar_repetition_penalty: 1.2, ar_kl_weight: 0.2 }, {}],
+      'config.process[0].model.model_kwargs': [{ cot: 'full', abc_dropout: 0.5, sample_ar_repetition_penalty: 1.2, ar_kl_weight: 0.2 }, {}],
     },
     // native YuE2 prompt: style text, a [Lyrics] line, the lyrics
     hasMultiLinePrompts: true,
@@ -1447,6 +1448,14 @@ export const modelArchs: ModelArch[] = [
           not be bit-faithful to the training audio even when the AR replays a song exactly. Though it is extremely
           close.
         </p>
+        <p>
+          Training matches Comfy's Generate ABC path (<code>cot: full</code> in model kwargs): at cache time every song is
+          transcribed to an ABC lead sheet with SheetSage2 (the Comfy-Org/YuE2 <code>audio_encoders</code> weights, no
+          extra packages), the sheet is stored with the latents, and the AR learns
+          lyrics to sheet and sheet to tokens. <code>abc_dropout</code> (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. <code>cot: melody</code> uses sheets without chords; <code>cot: off</code> skips the sheet, which is what
+          Comfy runs when its ABC input is left empty. Changing the mode requires re-caching the dataset.
+        </p>
         <p>
           Prompt format: a style line, then <code>[Lyrics]</code>, then the lyrics with bracketed section headers such
           as <code>[Verse 1]</code> and <code>[Chorus]</code>. The Qwen3-Omni captioner has a YuE2 preset that writes
@@ -1454,7 +1463,8 @@ export const modelArchs: ModelArch[] = [
         </p>
       </div>
     ),
-    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'],
   },
   {