Merge pull request #1042 from ostris/yue2

YuE2

README.md

@@ -65,6 +65,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 
 ### Experimental
 - [lodestones/Zeta-Chroma](https://huggingface.co/lodestones/Zeta-Chroma) (Zeta Chroma)
+- [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B) (YuE2) - audio. The official audio-to-token encoder is unreleased; training uses the community tokenizer by Kytra ([@sin_ceriously](https://x.com/sin_ceriously)), [Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4](https://huggingface.co/Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4).
 
 ## Installation
 

extensions_built_in/audio_models/__init__.py

@@ -1,7 +1,9 @@
 from .ace_step import AceStep15Model, AceStep15XLModel
+from .yue2 import YuE2AudioModel
 
 AI_TOOLKIT_MODELS = [
     # put a list of models here
     AceStep15Model,
     AceStep15XLModel,
+    YuE2AudioModel,
 ]

extensions_built_in/audio_models/yue2/__init__.py

@@ -0,0 +1 @@
+from .yue2_model import YuE2AudioModel

extensions_built_in/audio_models/yue2/src/model.py

@@ -0,0 +1,449 @@
+"""YuE2 AR + NAR mixture-of-transformers backbone in the ComfyUI repack layout.
+
+Two 28-layer Qwen3-shaped experts share one architecture:
+  - ``YuE2AR``  (``text_encoders.*`` in the repack): token embedding, causal
+    layers, final norm and lm head. Writes ABC + semantic codec tokens.
+  - ``YuE2NAR`` (``model.diffusion_model.*``): latent in/out projections,
+    timestep embedder, sinusoidal frame position table and the same layer
+    stack. Predicts flow velocity on VAE latents while attending into the
+    AR expert's key/value cache over ``prefix + codec tokens + MUSIC_END``.
+
+Weights use the merged ``qkv_proj`` / ``gate_up_proj`` linears of the
+Comfy-Org checkpoint so LoRA keys line up with ComfyUI's loader.
+"""
+
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+import torch.utils.checkpoint as ckpt
+
+from toolkit.models.v2._mixin import OstrisModelMixin
+
+# vocabulary layout of the AR expert
+EOD = 151643
+ABC_START, ABC_END = 151847, 151848
+MUSIC_START, MUSIC_END = 151851, 151852
+CODEC_OFFSET, CODEC_SIZE = 151853, 32768
+CONTEXT = 24576
+FRAMES_PER_SECOND = 25
+LATENT_DIM = 64
+
+INSTRUCTIONS = {
+    "off": "Generate music with codec tokens from the given conditions.",
+    "melody": "Generate a melody-only ABC transcription without chord symbols, then generate music with codec tokens from the given conditions.",
+    "full": "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions.",
+}
+
+
+class YuE2Config:
+    hidden_size = 2048
+    intermediate_size = 6144
+    num_hidden_layers = 28
+    num_attention_heads = 16
+    num_key_value_heads = 8
+    head_dim = 128
+    vocab_size = 184704
+    rms_norm_eps = 1e-6
+    rope_theta = 1000000.0
+    max_position_embeddings = CONTEXT
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        out = x.float()
+        out = out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps)
+        return out.to(x.dtype) * self.weight
+
+
+def rope_cos_sin(positions: torch.Tensor, head_dim: int, theta: float):
+    """positions [B, L] -> cos, sin [B, L, head_dim // 2] in fp32."""
+    inv_freq = 1.0 / (
+        theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device) / head_dim)
+    )
+    angles = positions.float().unsqueeze(-1) * inv_freq
+    return angles.cos(), angles.sin()
+
+
+def apply_rope(x, cos, sin):
+    # x [B, H, L, D]; cos/sin [B, L, D/2]
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    cos = cos.unsqueeze(1).to(x.dtype)
+    sin = sin.unsqueeze(1).to(x.dtype)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class YuE2Attention(nn.Module):
+    def __init__(self, cfg: YuE2Config):
+        super().__init__()
+        self.num_heads = cfg.num_attention_heads
+        self.num_kv_heads = cfg.num_key_value_heads
+        self.head_dim = cfg.head_dim
+        inner = self.num_heads * self.head_dim
+        kv = self.num_kv_heads * self.head_dim
+        self.qkv_proj = nn.Linear(cfg.hidden_size, inner + 2 * kv, bias=False)
+        self.o_proj = nn.Linear(inner, cfg.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
+
+    def project(self, x, cos, sin):
+        b, l, _ = x.shape
+        inner = self.num_heads * self.head_dim
+        kv = self.num_kv_heads * self.head_dim
+        q, k, v = self.qkv_proj(x).split((inner, kv, kv), dim=-1)
+        q = q.view(b, l, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(b, l, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(b, l, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        q = apply_rope(self.q_norm(q), cos, sin)
+        k = apply_rope(self.k_norm(k), cos, sin)
+        return q, k, v
+
+    def attend(self, q, k, v, causal: bool):
+        out = F.scaled_dot_product_attention(
+            q, k, v, is_causal=causal, enable_gqa=self.num_heads != self.num_kv_heads
+        )
+        b, _, l, _ = out.shape
+        return self.o_proj(out.transpose(1, 2).reshape(b, l, -1))
+
+
+class YuE2MLP(nn.Module):
+    def __init__(self, cfg: YuE2Config):
+        super().__init__()
+        self.gate_up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size * 2, bias=False)
+        self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
+
+    def forward(self, x):
+        gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
+        return self.down_proj(F.silu(gate) * up)
+
+
+class YuE2Layer(nn.Module):
+    def __init__(self, cfg: YuE2Config):
+        super().__init__()
+        self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.self_attn = YuE2Attention(cfg)
+        self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.mlp = YuE2MLP(cfg)
+
+    def forward(
+        self,
+        x,
+        cos,
+        sin,
+        prefix_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        causal: bool = True,
+        return_kv: bool = False,
+    ):
+        q, k, v = self.self_attn.project(self.input_layernorm(x), cos, sin)
+        present = (k, v) if return_kv else None
+        if prefix_kv is not None:
+            k = torch.cat([prefix_kv[0], k], dim=2)
+            v = torch.cat([prefix_kv[1], v], dim=2)
+        x = x + self.self_attn.attend(q, k, v, causal)
+        x = x + self.mlp(self.post_attention_layernorm(x))
+        return x, present
+
+
+class YuE2Stack(nn.Module):
+    """``model.layers`` + ``model.norm`` of one expert."""
+
+    def __init__(self, cfg: YuE2Config):
+        super().__init__()
+        self.layers = nn.ModuleList([YuE2Layer(cfg) for _ in range(cfg.num_hidden_layers)])
+        self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+
+
+class YuE2Expert(nn.Module):
+    """Shared expert base: LoRA targets this class name."""
+
+    def __init__(self, cfg: YuE2Config):
+        super().__init__()
+        self.cfg = cfg
+        self.gradient_checkpointing = False
+
+    def _run_layer(self, layer, *args, **kwargs):
+        if self.gradient_checkpointing and torch.is_grad_enabled():
+            return ckpt.checkpoint(layer, *args, use_reentrant=False, **kwargs)
+        return layer(*args, **kwargs)
+
+    @property
+    def device(self):
+        return self.model.norm.weight.device
+
+    @property
+    def dtype(self):
+        # norms stay in the activation dtype when the linears are quantized
+        return self.model.norm.weight.dtype
+
+
+class YuE2AR(YuE2Expert):
+    def __init__(self, cfg: YuE2Config):
+        super().__init__(cfg)
+        self.model = YuE2Stack(cfg)
+        self.model.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.model.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+
+    def embed(self, ids: torch.Tensor):
+        # 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
+        [B, kv_heads, L, head_dim] plus (optionally) the final normed hidden."""
+        b, l, _ = inputs_embeds.shape
+        positions = torch.arange(l, device=inputs_embeds.device)[None].expand(b, -1)
+        cos, sin = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta)
+        x = inputs_embeds
+        cache: List[Tuple[torch.Tensor, torch.Tensor]] = []
+        for layer in self.model.layers:
+            x, kv = self._run_layer(layer, x, cos, sin, None, True, True)
+            cache.append(kv)
+        hidden = self.model.norm(x) if return_hidden else None
+        return cache, hidden
+
+    @torch.no_grad()
+    def decode_step(self, token_embeds: torch.Tensor, cache: "KVCache"):
+        """One new token per batch row against a preallocated cache."""
+        b = token_embeds.shape[0]
+        pos = torch.full((b, 1), cache.length, device=token_embeds.device, dtype=torch.long)
+        cos, sin = rope_cos_sin(pos, self.cfg.head_dim, self.cfg.rope_theta)
+        x = token_embeds
+        for i, layer in enumerate(self.model.layers):
+            q, k, v = layer.self_attn.project(layer.input_layernorm(x), cos, sin)
+            k, v = cache.append(i, k, v)
+            x = x + layer.self_attn.attend(q, k, v, causal=False)
+            x = x + layer.mlp(layer.post_attention_layernorm(x))
+        cache.length += 1
+        return self.model.lm_head(self.model.norm(x[:, -1]))
+
+    @torch.no_grad()
+    def prefill_into_cache(self, inputs_embeds: torch.Tensor, cache: "KVCache"):
+        b, l, _ = inputs_embeds.shape
+        positions = torch.arange(l, device=inputs_embeds.device)[None].expand(b, -1)
+        cos, sin = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta)
+        x = inputs_embeds
+        for i, layer in enumerate(self.model.layers):
+            q, k, v = layer.self_attn.project(layer.input_layernorm(x), cos, sin)
+            cache.write(i, 0, k, v)
+            x = x + layer.self_attn.attend(q, k, v, causal=True)
+            x = x + layer.mlp(layer.post_attention_layernorm(x))
+        cache.length = l
+        return self.model.lm_head(self.model.norm(x[:, -1]))
+
+
+class KVCache:
+    """Preallocated [B, kv_heads, capacity, head_dim] per layer for AR decode."""
+
+    def __init__(self, cfg: YuE2Config, batch: int, capacity: int, device, dtype):
+        shape = (batch, cfg.num_key_value_heads, capacity, cfg.head_dim)
+        self.k = [torch.empty(shape, device=device, dtype=dtype) for _ in range(cfg.num_hidden_layers)]
+        self.v = [torch.empty(shape, device=device, dtype=dtype) for _ in range(cfg.num_hidden_layers)]
+        self.capacity = capacity
+        self.length = 0
+
+    def write(self, layer, start, k, v):
+        n = k.shape[2]
+        self.k[layer][:, :, start : start + n] = k
+        self.v[layer][:, :, start : start + n] = v
+
+    def append(self, layer, k, v):
+        end = self.length + k.shape[2]
+        if end > self.capacity:
+            raise ValueError("YuE2 KV cache capacity exceeded")
+        self.write(layer, self.length, k, v)
+        return self.k[layer][:, :, :end], self.v[layer][:, :, :end]
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, freq_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(freq_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size)
+        )
+        self.freq_size = freq_size
+
+    def forward(self, t: torch.Tensor):
+        half = self.freq_size // 2
+        freqs = torch.exp(
+            -math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half
+        )
+        args = t.float().unsqueeze(-1) * freqs.unsqueeze(0)
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class AudioPositionEmbedding(nn.Module):
+    def __init__(self, max_frames, hidden_size):
+        super().__init__()
+        self.register_buffer("pe", torch.zeros(max_frames, hidden_size))
+
+    def forward(self, length: int):
+        return self.pe[:length]
+
+
+class YuE2NAR(YuE2Expert):
+    def __init__(self, cfg: YuE2Config):
+        super().__init__(cfg)
+        self.model = YuE2Stack(cfg)
+        self.vae2llm = nn.Linear(LATENT_DIM, cfg.hidden_size)
+        self.llm2vae = nn.Linear(cfg.hidden_size, LATENT_DIM)
+        self.time_embedder = TimestepEmbedder(cfg.hidden_size)
+        self.latent_pos_embed = AudioPositionEmbedding(cfg.max_position_embeddings, cfg.hidden_size)
+
+    def forward(
+        self,
+        noisy_latents: torch.Tensor,  # [B, T, 64]
+        t: torch.Tensor,  # [B] in 0..1, 1 = pure noise
+        prefix_cache: List[Tuple[torch.Tensor, torch.Tensor]],
+        prefix_len: int,
+    ):
+        b, frames, _ = noisy_latents.shape
+        n = frames + 2
+        x = F.pad(noisy_latents, (0, 0, 1, 1))  # START / END slots
+        x = self.vae2llm(x.to(self.vae2llm.weight.dtype))
+        x = x + self.time_embedder(t)[:, None]
+        x = x + self.latent_pos_embed(n)[None].to(x.dtype)
+        positions = torch.arange(prefix_len, prefix_len + n, device=x.device)[None].expand(b, -1)
+        cos, sin = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta)
+        for layer, kv in zip(self.model.layers, prefix_cache):
+            x, _ = self._run_layer(layer, x, cos, sin, kv, False, False)
+        return self.llm2vae(self.model.norm(x))[:, 1:-1]
+
+
+class YuE2Model(nn.Module, OstrisModelMixin):
+    """Both experts under one module so a single LoRA network covers them."""
+
+    aitk_comfy_repo = "Comfy-Org/YuE2"
+
+    def __init__(self, cfg: Optional[YuE2Config] = None):
+        super().__init__()
+        cfg = cfg or YuE2Config()
+        self.cfg = cfg
+        self.ar = YuE2AR(cfg)
+        self.nar = YuE2NAR(cfg)
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["ar.model.layers", "nar.model.layers"]
+
+    @classmethod
+    def get_quantization_exclude_modules(cls):
+        return ["ar.model.embed_tokens", "ar.model.lm_head", "nar.vae2llm", "nar.llm2vae", "nar.time_embedder*"]
+
+    @property
+    def device(self):
+        return self.nar.device
+
+    @property
+    def dtype(self):
+        return self.nar.dtype
+
+    def enable_gradient_checkpointing(self):
+        self.ar.gradient_checkpointing = True
+        self.nar.gradient_checkpointing = True
+
+    def disable_gradient_checkpointing(self):
+        self.ar.gradient_checkpointing = False
+        self.nar.gradient_checkpointing = False
+
+    @classmethod
+    def load_from_state_dict(cls, sd: dict, dtype=torch.bfloat16):
+        """``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)):
+            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]})")
+        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()
+def merge_nar_lora(model: YuE2Model, ckpt_path: str, scale: float = 1.0):
+    """Fold the community NAR adapter (unmerged q/k/v + gate/up LoRA pairs,
+    plus full vae2llm/llm2vae weights) into the merged-projection base."""
+    try:
+        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
+    kv = cfg.num_key_value_heads * cfg.head_dim
+    qkv_rows = {"q_proj": (0, inner), "k_proj": (inner, inner + kv), "v_proj": (inner + kv, inner + 2 * kv)}
+    gu_rows = {"gate_proj": (0, cfg.intermediate_size), "up_proj": (cfg.intermediate_size, 2 * cfg.intermediate_size)}
+    for layer in model.nar.model.layers:
+        for name in ("q_proj", "k_proj", "v_proj", "o_proj"):
+            A = next(tensors).float()
+            B = next(tensors).float()
+            delta = (B @ A) * scale
+            if name == "o_proj":
+                w = layer.self_attn.o_proj.weight
+                w.add_(delta.to(w.device, w.dtype))
+            else:
+                r0, r1 = qkv_rows[name]
+                w = layer.self_attn.qkv_proj.weight
+                w[r0:r1].add_(delta.to(w.device, w.dtype))
+        for name in ("gate_proj", "up_proj", "down_proj"):
+            A = next(tensors).float()
+            B = next(tensors).float()
+            delta = (B @ A) * scale
+            if name == "down_proj":
+                w = layer.mlp.down_proj.weight
+                w.add_(delta.to(w.device, w.dtype))
+            else:
+                r0, r1 = gu_rows[name]
+                w = layer.mlp.gate_up_proj.weight
+                w[r0:r1].add_(delta.to(w.device, w.dtype))
+    io = ck.get("io", {})
+    for name in ("vae2llm", "llm2vae"):
+        if name in io:
+            mod = getattr(model.nar, name)
+            for pname, value in io[name].items():
+                getattr(mod, pname).copy_(value.to(mod.weight.device, mod.weight.dtype))
+    return model

extensions_built_in/audio_models/yue2/src/pipeline.py

@@ -0,0 +1,316 @@
+"""Sampling for YuE2: AR codec-token generation, then NAR flow matching.
+
+Follows the released protocol: semantic sampling with temperature 1.0,
+top-p 0.95, top-k 100, repetition penalty 1.2 over the last 50 tokens, at
+least 200 tokens before MUSIC_END; acoustic synthesis with a midpoint ODE
+solver from t=1 (noise) to t=0, one AR prefill per context chunk.
+"""
+
+from typing import Callable, List, Optional, Tuple
+
+import torch
+
+from .model import (
+    ABC_END,
+    CODEC_OFFSET,
+    CODEC_SIZE,
+    CONTEXT,
+    EOD,
+    KVCache,
+    MUSIC_END,
+    YuE2Model,
+    rope_cos_sin,
+)
+from .vae import HOP, YuE2VAE
+
+
+def chunk_ranges(frames: int, prefix_tokens: int, context: int = CONTEXT) -> List[Tuple[int, int]]:
+    size = (context - prefix_tokens - 3) // 2
+    if frames < 1 or size < 1:
+        raise ValueError("YuE2 needs codec tokens and enough context for at least one acoustic frame.")
+    return [(start, min(start + size, frames)) for start in range(0, frames, size)]
+
+
+def sample_logits(
+    logits: torch.Tensor,
+    history: List[int],
+    step: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float,
+    penalty_window: int,
+    min_tokens: int,
+    generator: torch.Generator,
+    phase: str = "semantic",
+    legacy_off: bool = False,
+) -> int:
+    """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)
+    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[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))
+        penalty = repetition_penalty ** counts
+        scores = torch.where(scores < 0, scores * penalty, scores / penalty)
+    if temperature == 0:
+        return int(scores.argmax().item())
+    scores /= temperature
+    threshold = scores.topk(min(top_k, scores.shape[-1])).values[-1]
+    scores.masked_fill_(scores < threshold, -torch.inf)
+    if top_p < 1:
+        values, indices = scores.sort(descending=True)
+        probs = values.softmax(-1)
+        removed = probs.cumsum(-1) - probs > top_p
+        removed[: 3 if legacy_off else 1] = False
+        values.masked_fill_(removed, -torch.inf)
+        scores = values.scatter(-1, indices, values)
+    probs = scores.softmax(-1)
+    return int(torch.multinomial(probs.to(generator.device), 1, generator=generator).item())
+
+
+class ARDecodeGraph:
+    """One-token AR decode step captured as a CUDA graph over static buffers:
+    token id in, logits out, preallocated [1, kv_heads, capacity, head_dim]
+    cache per layer. Attention is hand-grouped (GQA) with a length mask so the
+    whole step is capturable; the eager KVCache path stays as the fallback."""
+
+    def __init__(self, ar, capacity: int, device, dtype):
+        cfg = ar.cfg
+        self.ar = ar
+        self.capacity = capacity
+        self.kv_heads = cfg.num_key_value_heads
+        self.groups = cfg.num_attention_heads // cfg.num_key_value_heads
+        self.head_dim = cfg.head_dim
+        shape = (1, self.kv_heads, capacity, self.head_dim)
+        self.k = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(cfg.num_hidden_layers)]
+        self.v = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(cfg.num_hidden_layers)]
+        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
+
+    def load_prefill(self, cache, length: int):
+        for i, (k, v) in enumerate(cache):
+            self.k[i][:, :, :length] = k
+            self.v[i][:, :, :length] = v
+        self.pos.fill_(length)
+
+    def _layers(self):
+        ar = self.ar
+        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
+        slot = self.pos.view(-1)
+        for i, layer in enumerate(ar.model.layers):
+            q, k, v = layer.self_attn.project(layer.input_layernorm(x), cos, sin)
+            self.k[i].index_copy_(2, slot, k)
+            self.v[i].index_copy_(2, slot, v)
+            qg = q.view(1, self.kv_heads, self.groups, self.head_dim)
+            scores = torch.matmul(qg, self.k[i].transpose(-1, -2)) * scale  # [1, kv, groups, capacity]
+            scores = scores.float().masked_fill(~mask, float("-inf")).softmax(-1).to(q.dtype)
+            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.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._layers()
+        torch.cuda.current_stream().wait_stream(stream)
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self._layers()
+
+    def step(self, token: int):
+        self.token.fill_(token)
+        self._step()
+        self.pos.add_(1)
+        return self.logits[0]
+
+
+class YuE2Pipeline:
+    def __init__(self, model: YuE2Model, vae: YuE2VAE):
+        self.model = model
+        self.vae = vae
+        self.do_tiled_decoding = True
+        self.use_cuda_graphs = True
+
+    @torch.no_grad()
+    def generate_tokens(
+        self,
+        prefix_embeds: torch.Tensor,  # [1, L, H]
+        phase: str,
+        max_tokens: int,
+        seed: int,
+        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]
+        if prefix_len + max_tokens > CONTEXT:
+            raise ValueError("YuE2 prompt plus token budget exceeds the model context")
+        # reference sampler: seeded generator on the compute device, multinomial over the filtered softmax
+        generator = torch.Generator(device=device).manual_seed(seed)
+        capacity = prefix_len + max_tokens
+        graph = None
+        if self.use_cuda_graphs and device.type == "cuda":
+            try:
+                graph = ARDecodeGraph(ar, capacity, device, ar.dtype)
+                prefill_cache, hidden = ar.prefill(prefix_embeds.to(ar.dtype), return_hidden=True)
+                logits = ar.model.lm_head(hidden[:, -1])[0]
+                graph.load_prefill(prefill_cache, prefix_len)
+                del prefill_cache, hidden
+                graph.capture()
+            except Exception as e:  # capture unsupported for these ops/kernels: eager path
+                print(f"YuE2: CUDA graph decode unavailable ({type(e).__name__}: {e}); using eager decode")
+                graph = None
+                torch.cuda.empty_cache()
+        if graph is None:
+            cache = KVCache(self.model.cfg, 1, capacity, device, ar.dtype)
+            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,
+                phase=phase, legacy_off=legacy_off,
+            )
+            if token == end:
+                break
+            history.append(token)
+            if progress is not None:
+                progress(step + 1, max_tokens)
+            if step + 1 < max_tokens:
+                if graph is not None:
+                    logits = graph.step(token)
+                else:
+                    emb = ar.embed(torch.tensor([[token]], device=device))
+                    logits = ar.decode_step(emb, cache)[0]
+        if graph is None:
+            del cache
+        del graph
+        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(
+        self,
+        prefix_embeds: torch.Tensor,  # [1, L, H]
+        codec: List[int],
+        seed: int,
+        steps: int = 32,
+        step_callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+    ) -> torch.Tensor:
+        """-> latents [1, T, 64] in the NAR's dtype."""
+        ar, nar = self.model.ar, self.model.nar
+        device = prefix_embeds.device
+        dtype = nar.dtype
+        generator = torch.Generator(device="cpu").manual_seed(seed)
+        noise = torch.randn((len(codec), 64), generator=generator).to(device, dtype)
+        ranges = chunk_ranges(len(codec), prefix_embeds.shape[1])
+        outputs = []
+        total_steps = steps * len(ranges)
+        for ci, (start, end) in enumerate(ranges):
+            ids = torch.tensor([[c + CODEC_OFFSET for c in codec[start:end]] + [MUSIC_END]], device=device)
+            embeds = torch.cat([prefix_embeds.to(ar.dtype), ar.embed(ids)], dim=1)
+            cache, _ = ar.prefill(embeds)
+            prefix_len = embeds.shape[1]
+            state = noise[start:end][None]
+            dt = 1.0 / steps
+            for step in range(steps):
+                t = 1.0 - step * dt
+                v1 = nar(state, torch.full((1,), t, device=device, dtype=dtype), cache, prefix_len)
+                mid = state - v1 * (dt / 2)
+                v2 = nar(mid, torch.full((1,), t - dt / 2, device=device, dtype=dtype), cache, prefix_len)
+                state = state - v2 * dt
+                if step_callback is not None:
+                    step_callback(ci * steps + step, total_steps, state)
+            outputs.append(state)
+            del cache
+        return torch.cat(outputs, dim=1)
+
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """[1, T, 64] -> waveform [1, 2, samples] float32."""
+        z = latents.transpose(1, 2).to(self.vae.device, self.vae.dtype)
+        if self.do_tiled_decoding:
+            audio = self.vae.tiled_decode(z)
+        else:
+            audio = self.vae.decode(z)
+        return audio.float().clamp(-1, 1)

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

@@ -0,0 +1,230 @@
+"""Text side and audio side of YuE2 conditioning.
+
+- ``YuE2TextTokenizer``: the Qwen BPE that ComfyUI embeds in the checkpoint
+  (``text_encoders.yue2_tokenizer_json``), plus the prompt/prefix protocol.
+- ``SemanticTokenizer``: the community audio -> codec-token head
+  (Mothersuperior v4): MERT-v2-FullSong layer-20 features at 25 Hz,
+  per-track instance normalised, through an 8-layer transformer classifier
+  over 512-frame windows.
+"""
+
+import unicodedata
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .model import ABC_END, ABC_START, CODEC_SIZE, EOD, INSTRUCTIONS, MUSIC_START
+
+MERT_REPO = "m-a-p/MERT-v2-FullSong"
+MERT_SAMPLE_RATE = 24000
+MERT_LAYER = 20
+HEAD_REPO = "Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4"
+HEAD_FILE = "tokenizer_head_joint_v4.pt"
+NAR_LORA_FILE = "nar_lora_joint_v4.pt"
+HEAD_WINDOW = 512
+
+
+class YuE2TextTokenizer:
+    def __init__(self, tokenizer_json: bytes):
+        from tokenizers import Tokenizer
+
+        self.tokenizer = Tokenizer.from_str(tokenizer_json.decode("utf-8"))
+
+    def encode(self, text: str) -> List[int]:
+        return self.tokenizer.encode(unicodedata.normalize("NFC", text)).ids
+
+    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."""
+        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):
+    def __init__(self, din=1024, d=512, layers=8, heads=8, vocab=CODEC_SIZE, window=HEAD_WINDOW):
+        super().__init__()
+        self.inp = nn.Linear(din, d)
+        self.pos = nn.Parameter(torch.zeros(1, window, d))
+        layer = nn.TransformerEncoderLayer(d, heads, 4 * d, dropout=0.1, batch_first=True, norm_first=True, activation="gelu")
+        self.enc = nn.TransformerEncoder(layer, layers)
+        self.norm = nn.LayerNorm(d)
+        self.head = nn.Linear(d, vocab)
+
+    def forward(self, x):
+        return self.head(self.norm(self.enc(self.inp(x) + self.pos[:, : x.shape[1]])))
+
+
+def _rebuild_rotary(model: nn.Module) -> int:
+    """MERT2's RotaryEmbedding keeps ``inv_freq`` as a non-persistent buffer; transformers'
+    meta-device loading materializes it uninitialized (zeros on a fresh GPU, garbage or NaN
+    otherwise). Recompute it from the module's own formula and drop the cos/sin cache."""
+    n = 0
+    for m in model.modules():
+        if hasattr(m, "inv_freq") and hasattr(m, "head_dim") and hasattr(m, "base"):
+            inv = 1.0 / (m.base ** (torch.arange(0, m.head_dim, 2, dtype=torch.float32) / m.head_dim))
+            m.inv_freq = inv.to(device=m.inv_freq.device)
+            for attr, val in (("_cos", None), ("_sin", None), ("_sequence_length", 0), ("_cache_device", None)):
+                if hasattr(m, attr):
+                    setattr(m, attr, val)
+            n += 1
+    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."""
+
+    def __init__(self, head_path: str, mert_repo: str = MERT_REPO, debug: bool = False):
+        super().__init__()
+        self.debug = debug
+        from transformers import AutoFeatureExtractor, AutoModel
+
+        self.processor = AutoFeatureExtractor.from_pretrained(mert_repo, trust_remote_code=True)
+        self.mert = AutoModel.from_pretrained(mert_repo, trust_remote_code=True).eval()
+        self.mert.requires_grad_(False)
+        _rebuild_rotary(self.mert)
+        self.head = TokenizerHead()
+        ck = torch.load(head_path, map_location="cpu", weights_only=False)
+        self.head.load_state_dict(ck["model"])
+        self.head.eval().requires_grad_(False)
+
+    @property
+    def device(self):
+        return self.head.head.weight.device
+
+    @torch.no_grad()
+    def mert_features(self, mono24: torch.Tensor) -> torch.Tensor:
+        """mono24 [samples] at 24 kHz -> [T25, 1024] fp32 layer-20 features."""
+        device = self.device
+        chunk = MERT_SAMPLE_RATE * 30
+        chunks = [mono24[s : s + chunk] for s in range(0, mono24.shape[0], chunk)]
+        chunks = [c for c in chunks if c.shape[0] >= MERT_SAMPLE_RATE]
+        full = [c for c in chunks if c.shape[0] == chunk]
+        tail = [c for c in chunks if c.shape[0] < chunk]
+        feats = []
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
+            for group in ([full] if full else []) + [[c] for c in tail]:
+                inp = self.processor([c.cpu().numpy() for c in group], sampling_rate=MERT_SAMPLE_RATE, return_tensors="pt")
+                inp = {k: v.to(device) for k, v in inp.items()}
+                out = self.mert(**inp, output_hidden_states=True)
+                feats.append(out.hidden_states[MERT_LAYER].reshape(-1, 1024))
+        h = torch.cat(feats, 0).float()
+        t25 = int(round(mono24.shape[0] / MERT_SAMPLE_RATE * 25))
+        return F.interpolate(h.T[None], size=t25, mode="linear", align_corners=False)[0].T
+
+    @torch.no_grad()
+    def tokens_from_features(self, feats: torch.Tensor) -> torch.Tensor:
+        """[T, 1024] -> [T] long. Windowed inference with edge trimming."""
+        x = feats.float()
+        x = (x - x.mean(0)) / (x.std(0) + 1e-5)
+        total = x.shape[0]
+        win = HEAD_WINDOW
+        out = torch.zeros(total, dtype=torch.long, device=x.device)
+        starts = list(range(0, max(1, total - win + 1), win // 2))
+        if starts[-1] + win < total:
+            starts.append(max(0, total - win))
+        device = self.device
+        for s0 in starts:
+            xw = x[s0 : s0 + win]
+            n = xw.shape[0]
+            if n < win:
+                xw = F.pad(xw, (0, 0, 0, win - n))
+            with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
+                pred = self.head(xw[None].to(device))[0, :n].float().argmax(-1)
+            lo = s0 + (0 if s0 == 0 else win // 4)
+            hi = s0 + n - (0 if s0 + n >= total else win // 4)
+            out[lo:hi] = pred[lo - s0 : hi - s0].to(out.device)
+        return out
+
+    @torch.no_grad()
+    def tokenize(self, waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:
+        """waveform [C, samples] -> codec ids [T] (long, no vocab offset)."""
+        import torchaudio
+
+        mono = waveform.float().mean(0) if waveform.dim() == 2 else waveform.float()
+        if sample_rate != MERT_SAMPLE_RATE:
+            mono = torchaudio.functional.resample(mono, sample_rate, MERT_SAMPLE_RATE)
+        feats = self.mert_features(mono.to(self.device))
+        out = self.tokens_from_features(feats)
+        if self.debug and out.unique().numel() < 8:
+            print(f"YuE2 tokenizer: degenerate token stream (unique {out.unique().numel()}); MERT features nan {torch.isnan(feats).sum().item()}")
+        return out

extensions_built_in/audio_models/yue2/src/vae.py

@@ -0,0 +1,178 @@
+"""YuE2 Oobleck-style waveform VAE (48 kHz stereo, 64-ch latents at 25 Hz).
+
+Matches the ``vae.*`` tensors of the Comfy-Org repack: channels 64, channel
+multipliers (1, 2, 4, 8, 16, 32), strides (2, 2, 4, 4, 5, 6) -> hop 1920.
+Encoding is deterministic (mean half of the bottleneck). Checkpoints carry
+old-style ``weight_g`` / ``weight_v`` weight-norm keys; they are remapped to
+torch parametrizations on load.
+"""
+
+import math
+
+import torch
+from torch import nn
+
+from toolkit.models.v2._mixin import OstrisModelMixin
+
+SAMPLE_RATE = 48000
+HOP = 1920
+LATENT_RATE = SAMPLE_RATE // HOP  # 25
+
+
+def WNConv1d(*args, **kwargs):
+    return torch.nn.utils.parametrizations.weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvT1d(*args, **kwargs):
+    return torch.nn.utils.parametrizations.weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class SnakeBeta(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.zeros(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        a = self.alpha.unsqueeze(0).unsqueeze(-1).exp()
+        b = self.beta.unsqueeze(0).unsqueeze(-1).exp()
+        return x + (1.0 / (b + 1e-9)) * torch.sin(x * a).pow(2)
+
+
+class ResUnit(nn.Module):
+    def __init__(self, ch, dilation):
+        super().__init__()
+        self.layers = nn.Sequential(
+            SnakeBeta(ch),
+            WNConv1d(ch, ch, 7, dilation=dilation, padding=(dilation * 6) // 2),
+            SnakeBeta(ch),
+            WNConv1d(ch, ch, 1),
+        )
+
+    def forward(self, x):
+        return x + self.layers(x)
+
+
+class EncBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, stride):
+        super().__init__()
+        self.layers = nn.Sequential(
+            ResUnit(in_ch, 1),
+            ResUnit(in_ch, 3),
+            ResUnit(in_ch, 9),
+            SnakeBeta(in_ch),
+            WNConv1d(in_ch, out_ch, 2 * stride, stride=stride, padding=math.ceil(stride / 2)),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class DecBlock(nn.Module):
+    def __init__(self, in_ch, out_ch, stride):
+        super().__init__()
+        self.layers = nn.Sequential(
+            SnakeBeta(in_ch),
+            # odd strides need output_padding to land on exactly stride * T
+            WNConvT1d(in_ch, out_ch, 2 * stride, stride=stride, padding=math.ceil(stride / 2), output_padding=stride % 2),
+            ResUnit(out_ch, 1),
+            ResUnit(out_ch, 3),
+            ResUnit(out_ch, 9),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class _SeqWrap(nn.Module):
+    def __init__(self, *modules):
+        super().__init__()
+        self.layers = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+def convert_weight_norm_keys(sd: dict) -> dict:
+    out = {}
+    for k, v in sd.items():
+        if k.endswith(".weight_g"):
+            k = k[: -len(".weight_g")] + ".parametrizations.weight.original0"
+        elif k.endswith(".weight_v"):
+            k = k[: -len(".weight_v")] + ".parametrizations.weight.original1"
+        out[k] = v
+    return out
+
+
+class YuE2VAE(nn.Module, OstrisModelMixin):
+    @classmethod
+    def load_from_state_dict(cls, state_dict, dtype=torch.float32):
+        vae = cls()
+        sd = convert_weight_norm_keys(state_dict)
+        missing, unexpected = vae.load_state_dict({k: v.to(dtype) for k, v in sd.items()}, strict=False)
+        if missing:
+            raise ValueError(f"YuE2 VAE missing keys: {missing[:5]} (+{max(0, len(missing) - 5)})")
+        if unexpected:
+            print(f"    YuE2 VAE unexpected: {len(unexpected)} (first 3: {unexpected[:3]})")
+        return vae.to(dtype)
+
+    def __init__(
+        self,
+        in_ch=2,
+        channels=64,
+        latent_dim=64,
+        c_mults=(1, 2, 4, 8, 16, 32),
+        strides=(2, 2, 4, 4, 5, 6),
+    ):
+        super().__init__()
+        cm = [1] + list(c_mults)
+        enc = [WNConv1d(in_ch, cm[0] * channels, 7, padding=3)]
+        for i in range(len(cm) - 1):
+            enc.append(EncBlock(cm[i] * channels, cm[i + 1] * channels, strides[i]))
+        enc += [SnakeBeta(cm[-1] * channels), WNConv1d(cm[-1] * channels, latent_dim * 2, 3, padding=1)]
+        self.encoder = _SeqWrap(*enc)
+        dec = [WNConv1d(latent_dim, cm[-1] * channels, 7, padding=3)]
+        for i in range(len(cm) - 1, 0, -1):
+            dec.append(DecBlock(cm[i] * channels, cm[i - 1] * channels, strides[i - 1]))
+        dec += [SnakeBeta(cm[0] * channels), WNConv1d(cm[0] * channels, in_ch, 7, padding=3, bias=False)]
+        self.decoder = _SeqWrap(*dec)
+        self.latent_dim = latent_dim
+        self.upscale_factor = math.prod(strides)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def encode(self, x: torch.Tensor, chunk_seconds: float = 60.0) -> torch.Tensor:
+        """[B, 2, samples] -> [B, 64, T]. Encodes in whole-second chunks that
+        are multiples of the hop, matching the community tokenizer's prep."""
+        chunk = int(chunk_seconds * SAMPLE_RATE)
+        chunk -= chunk % HOP
+        outs = []
+        for s in range(0, x.shape[-1], chunk):
+            seg = x[..., s : s + chunk]
+            if seg.shape[-1] < HOP:
+                break
+            outs.append(self.encoder(seg).chunk(2, dim=1)[0])
+        return torch.cat(outs, dim=-1)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(z)
+
+    def tiled_decode(self, z: torch.Tensor, core_frames: int = 750, halo_frames: int = 16) -> torch.Tensor:
+        """Halo-crop tiled decode: decode core + halo, keep the core."""
+        total = z.shape[-1]
+        if total <= core_frames + 2 * halo_frames:
+            return self.decoder(z)
+        pieces = []
+        for start in range(0, total, core_frames):
+            end = min(start + core_frames, total)
+            lo = max(0, start - halo_frames)
+            hi = min(total, end + halo_frames)
+            audio = self.decoder(z[..., lo:hi])
+            pieces.append(audio[..., (start - lo) * HOP : (start - lo + end - start) * HOP])
+        return torch.cat(pieces, dim=-1)

extensions_built_in/audio_models/yue2/yue2_model.py

@@ -0,0 +1,676 @@
+"""YuE2 (m-a-p/YuE2-3B) for ai-toolkit.
+
+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
+``Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4``. Its companion NAR adapter can
+be merged into the base NAR on load (``model_kwargs.merge_nar_lora``, default false).
+
+Training: per song the latent cache carries VAE latents plus the head's codec tokens; the
+"text embedding" is the AR prompt prefix (instruction + tags + lyrics) as token embeddings.
+Each step crops one window (``model_kwargs.train_window_frames``, 25 fps, 0 = whole song),
+prefills the AR expert over ``prefix + tokens[window] + MUSIC_END`` and trains the NAR
+flow loss on the real latents. With ``model_kwargs.ar_loss_weight`` > 0 the same prefill
+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.
+
+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
+import random
+import re
+from typing import List, Optional
+
+import torch
+import torch.utils.checkpoint
+from safetensors.torch import load_file
+from tqdm import tqdm
+
+from extensions_built_in.audio_models.base_audio_model import BaseAudioModel
+from toolkit.basic import flush
+from toolkit.config_modules import GenerateImageConfig
+from toolkit.dto import DTO
+from toolkit.models.v2.resolver import resolve_named_file
+from toolkit.paths import MODELS_PATH
+from toolkit.print import print_acc
+from toolkit.prompt_utils import PromptEmbeds
+from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
+
+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, SHEETSAGE_FILE, SemanticTokenizer, SheetSage2Transcriber, YuE2TextTokenizer
+from .src.vae import SAMPLE_RATE, YuE2VAE
+
+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"
+
+
+def resolve_ai_toolkit_file(path: str, component: str) -> str:
+    """Local file, or 'org/repo/file' downloaded into <MODELS_PATH>/ai_toolkit/<file>."""
+    if os.path.exists(path):
+        return path
+    splits = path.split("/")
+    if len(splits) < 3:
+        raise ValueError(f"Invalid {component} path: {path}. Must be a local file or 'org/repo/filename'.")
+    local_dir = os.path.join(MODELS_PATH, AI_TOOLKIT_DIR)
+    candidate = os.path.join(local_dir, splits[-1])
+    if os.path.exists(candidate):
+        return candidate
+    import huggingface_hub
+
+    return huggingface_hub.hf_hub_download(repo_id="/".join(splits[:2]), filename="/".join(splits[2:]), local_dir=local_dir)
+
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "shift": 1.0,
+    "use_dynamic_shifting": False,
+}
+
+
+def _tag(text: str, name: str) -> str:
+    m = re.search(rf"<{name}>(.*?)</{name}>", text, re.DOTALL)
+    return m.group(1).strip() if m else ""
+
+
+_SECTION = re.compile(r"^\s*\[(Tags|Lyrics|Duration)\]\s*$", re.IGNORECASE | re.MULTILINE)
+_SONG_SECTION = re.compile(r"^\s*\[[^\]]+\]\s*$")
+
+
+def _normalize_section(line: str) -> str:
+    """YuE2 section tags are Title case ([Verse 1], [Pre-Chorus]); uppercase the first letter and the
+    letter after each hyphen of a bracketed section line, leaving the rest (e.g. descriptors) alone."""
+    if not _SONG_SECTION.match(line):
+        return line
+    lead, body, trail = line[: line.index("[") + 1], line[line.index("[") + 1 : line.rindex("]")], line[line.rindex("]") :]
+    body = " ".join(w.capitalize() if w.isupper() and len(w) > 1 else w for w in body.split(" "))  # [BRIDGE] -> [Bridge]
+    body = body[:1].upper() + body[1:]
+    body = re.sub(r"-([a-z])", lambda m: "-" + m.group(1).upper(), body)
+    return lead + body + trail
+
+
+def _number(value: str):
+    try:
+        return float(value) if value.strip() else None
+    except ValueError:
+        return None
+
+
+def parse_caption(text: str) -> dict:
+    """Caption / prompt -> style, lyrics, duration.
+
+    Native form (YuE2's own prompt layout): style text, then a ``[Lyrics]`` line and the
+    lyrics; an optional leading ``[Tags]`` line and an optional trailing ``[Duration]``
+    line (max seconds, samples only). Plain text with no sections is the style.
+    Legacy ``<CAPTION>``/``<LYRICS>``/``<DURATION>`` tags are still accepted."""
+    if not isinstance(text, str):
+        text = ""
+    if "<CAPTION>" in text or "<LYRICS>" in text:
+        lyrics = "\n".join(_normalize_section(l) for l in _tag(text, "LYRICS").splitlines())
+        return {"style": _tag(text, "CAPTION"), "lyrics": lyrics, "duration": _number(_tag(text, "DURATION"))}
+    sections = {"tags": []}
+    current = "tags"
+    for line in text.splitlines():
+        m = _SECTION.match(line)
+        if m:
+            current = m.group(1).lower()
+            sections.setdefault(current, [])
+            continue
+        # no [Lyrics] header: the first song-section line ([Intro], [Verse 1], ...) starts the lyrics
+        if current == "tags" and "lyrics" not in sections and _SONG_SECTION.match(line):
+            current = "lyrics"
+            sections["lyrics"] = []
+        sections.setdefault(current, []).append(_normalize_section(line) if current == "lyrics" else line)
+    style = "\n".join(sections["tags"]).strip()
+    lyrics = "\n".join(sections.get("lyrics", [])).strip()
+    duration = _number("\n".join(sections.get("duration", [])))
+    return {"style": style, "lyrics": lyrics, "duration": duration}
+
+
+class YuE2TextEncoder(torch.nn.Module):
+    """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 = "full"):
+        super().__init__()
+        self._model = [model]
+        self.tokenizer = tokenizer
+        self.cot = cot
+        self.register_buffer("_anchor", torch.zeros(1), persistent=False)
+
+    @property
+    def device(self):
+        return self._anchor.device
+
+    @property
+    def dtype(self):
+        return self._model[0].ar.dtype
+
+    @torch.no_grad()
+    def forward(self, style: str, lyrics: str) -> torch.Tensor:
+        ids = torch.tensor([self.tokenizer.prefix_head_ids(style, lyrics, cot=self.cot)], dtype=torch.long)
+        return self._model[0].ar.embed(ids)
+
+
+class YuE2AudioModel(BaseAudioModel):
+    arch = "yue2"
+    sample_rate = SAMPLE_RATE
+
+    def __init__(self, device, model_config, dtype="bf16", custom_pipeline=None, noise_scheduler=None, **kwargs):
+        super().__init__(device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs)
+        self.is_flow_matching = True
+        self.is_transformer = True
+        self.target_lora_modules = ["YuE2AR", "YuE2NAR"]
+        kw = self.model_config.model_kwargs
+        self.train_window_frames = int(kw.get("train_window_frames", 1500))
+        self.ar_loss_weight = float(kw.get("ar_loss_weight", 1.0))
+        # AR next-token loss always runs on the song from its start (lyric order); 0 = whole song
+        self.ar_max_tokens = int(kw.get("ar_max_tokens", 0))
+        # optional separate learning rate for the AR expert's LoRA (Adam ignores loss scale)
+        self.ar_lr_multiplier = float(kw.get("ar_lr_multiplier", 1.0))
+        # trust region: KL(base || lora) on the AR next-token distributions, base = LoRA switched off
+        self.ar_kl_weight = float(kw.get("ar_kl_weight", 0.0))
+        self._loss_log_every = int(kw.get("loss_log_every", 25))
+        # extra per-readout AR diagnostics (targets, unique tokens, p(END), p(target))
+        self.debug = bool(kw.get("debug", False))
+        self._loss_log_step = 0
+        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
+        # encoders are only needed while caching; once training steps run without encode calls they are released
+        self._train_calls = 0
+        self._last_encode_call = 0
+        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)
+        self.sample_ar_repetition_penalty = float(kw.get("sample_ar_repetition_penalty", 1.2))
+        self.semantic_tokenizer: Optional[SemanticTokenizer] = None
+        self._pending_aux_loss = None
+        self._pending_ar_ce = None
+        self._pending_ar_kl = None
+        self.additional_loss_logs = {}
+
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config)
+
+    # ------------------------------------------------------------------
+    # loading
+    # ------------------------------------------------------------------
+    def load_model(self):
+        dtype = self.torch_dtype
+        device = self.device_torch
+        name_or_path = self.model_config.name_or_path or DEFAULT_CHECKPOINT
+        ckpt_path = resolve_named_file(name_or_path, component="yue2 checkpoint")
+        self.print_and_status_update(f"Loading YuE2 from {ckpt_path}")
+        sd = load_file(ckpt_path)
+        tok_json = sd.pop("text_encoders.yue2_tokenizer_json", None)
+        if tok_json is None:
+            raise ValueError("YuE2 checkpoint has no embedded tokenizer (expected the Comfy-Org repack)")
+        self.tokenizer = YuE2TextTokenizer(tok_json.numpy().tobytes())
+
+        if self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0:
+            raise NotImplementedError("Layer offloading not yet implemented for YuE2")
+
+        self.model = YuE2Model.load_from_state_dict(sd, dtype=dtype)
+        if self.merge_nar_lora:
+            lora_path = resolve_ai_toolkit_file(self.nar_lora_path, component="yue2 nar lora")
+            self.print_and_status_update(f"Merging NAR adapter {os.path.basename(lora_path)}")
+            merge_nar_lora(self.model, lora_path)
+        vae_sd = {k[len("vae.") :]: v for k, v in sd.items() if k.startswith("vae.")}
+        # reference runs the VAE in fp32; keep the trainer's dtype check from downcasting it
+        self.vae_torch_dtype = torch.float32
+        self.vae = YuE2VAE.load_from_state_dict(vae_sd, dtype=torch.float32)
+        del sd
+        flush()
+
+        self.model.aitk_post_load(**self.component_load_kwargs("transformer"))
+        flush()
+        self.model.to(device)
+        self.vae.to(self.vae_device_torch)
+        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))
+
+    def _get_semantic_tokenizer(self) -> SemanticTokenizer:
+        if self.semantic_tokenizer is None:
+            head_path = resolve_ai_toolkit_file(self.semantic_head_path, component="yue2 semantic head")
+            # log only: this can run mid-training and a status update would replace "Training" in the UI
+            print_acc("Loading YuE2 semantic tokenizer (MERT-v2-FullSong + head)")
+            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"]
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    # ------------------------------------------------------------------
+    # conditioning
+    # ------------------------------------------------------------------
+    def get_prompt_embeds(self, prompt) -> PromptEmbeds:
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        embeds, masks = [], []
+        for p in prompts:
+            parsed = parse_caption(p)
+            e = self.text_encoder(parsed["style"], parsed["lyrics"])  # [1, L, H]
+            embeds.append(e)
+            masks.append(torch.ones(1, e.shape[1], dtype=torch.long, device=e.device))
+        max_len = max(e.shape[1] for e in embeds)
+        embeds = [torch.nn.functional.pad(e, (0, 0, 0, max_len - e.shape[1])) for e in embeds]
+        masks = [torch.nn.functional.pad(m, (0, max_len - m.shape[1])) for m in masks]
+        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] 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
+        self._last_encode_call = self._train_calls
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        latents = self.vae.encode(audio_tensor.to(device=device, dtype=self.vae.dtype))  # [B, 64, T]
+        latents = latents.transpose(1, 2).contiguous()
+        tokenizer = self._get_semantic_tokenizer()
+        if tokenizer.device != self.device_torch:
+            tokenizer.to(self.device_torch)
+        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)
+        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
+    # ------------------------------------------------------------------
+    def condition_noisy_latents(self, latents: torch.Tensor, batch):
+        frames = latents.shape[1]
+        window = self.train_window_frames
+        if window <= 0 or frames <= window:
+            batch.yue2_window = (0, frames)
+            return latents
+        start = random.randint(0, frames - window)
+        batch.yue2_window = (start, start + window)
+        return latents[:, start : start + window]
+
+    def get_loss_target(self, *args, **kwargs):
+        noise = kwargs.get("noise")
+        batch = kwargs.get("batch")
+        start, end = getattr(batch, "yue2_window", (0, noise.shape[1]))
+        return (noise - batch.latents)[:, start:end].detach()
+
+    def get_additional_loss(self, pred: torch.Tensor, target: torch.Tensor):
+        """AR next-token CE from the last forward, added to the flow loss by the trainer and
+        logged as its own term (`loss/ar_ce`, unweighted)."""
+        aux = self._pending_aux_loss
+        self._pending_aux_loss = None
+        if aux is None:
+            self.additional_loss_logs = {}
+            return None
+        ce = self._pending_ar_ce.detach().float().item()
+        self.additional_loss_logs = {"loss/ar_ce": ce}
+        kl = self._pending_ar_kl
+        if kl is not None:
+            self.additional_loss_logs["loss/ar_kl"] = kl.detach().float().item()
+        self._loss_log_step += 1
+        if self.debug and self._loss_log_every > 0 and self._loss_log_step % self._loss_log_every == 0:
+            info = getattr(self, "_pending_aux_info", "")
+            kl_txt = f"  ar_kl {kl.detach().float().item():.4f}" if kl is not None else ""
+            tqdm.write(f"yue2 step {self._loss_log_step}: ar_ce {ce:.4f}{kl_txt}" + (f"  {info}" if info else ""))
+        return aux
+
+    @staticmethod
+    def _train_prefix(prefix_embeds: torch.Tensor, prefix_mask: torch.Tensor) -> torch.Tensor:
+        """Unpadded prefix embeds for one item."""
+        # the mask may arrive cast to bf16; count entries, never sum them
+        length = int((prefix_mask > 0.5).sum().item())
+        return prefix_embeds[:length]
+
+    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 = 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
+        want_kl = base_hidden is not None
+
+        def chunk_losses(h, tgt, bh):
+            # everything vocab-sized lives inside this checkpointed function, so autograd keeps only
+            # the [chunk, hidden] inputs and recomputes the fp32 logits/log-probs in backward
+            logits = lm_head(h).float()
+            ce = torch.nn.functional.cross_entropy(logits, tgt, reduction="sum")
+            if bh is not None:
+                with torch.no_grad():
+                    base_logp = torch.log_softmax(lm_head(bh).float(), -1)
+                kl = torch.nn.functional.kl_div(torch.log_softmax(logits, -1), base_logp, log_target=True, reduction="sum")
+            else:
+                kl = ce.new_zeros(())
+            if want_info:
+                with torch.no_grad():
+                    p = logits.softmax(-1)
+                    p_first_end = p[0, MUSIC_END]
+                    p_tgt = p.gather(1, tgt[:, None]).squeeze(1)
+            else:
+                p_first_end, p_tgt = ce.new_zeros(()), ce.new_zeros((0,))
+            return ce, kl, p_first_end, p_tgt
+
+        ce_sum = hidden.new_zeros((), dtype=torch.float32)
+        kl_sum = hidden.new_zeros((), dtype=torch.float32)
+        p_target, p_end_first = [], None
+        chunk = 512  # vocab-sized fp32 transients scale with this (~1.9 GB per 512 positions incl. KL)
+        for s0 in range(0, n, chunk):
+            h = hidden[s0 : s0 + chunk]
+            tgt = ids[s0 : s0 + chunk]
+            bh = base_hidden[s0 : s0 + chunk] if want_kl else None
+            if torch.is_grad_enabled():
+                ce, kl, p_first, p_tgt = torch.utils.checkpoint.checkpoint(chunk_losses, h, tgt, bh, use_reentrant=False)
+            else:
+                ce, kl, p_first, p_tgt = chunk_losses(h, tgt, bh)
+            ce_sum = ce_sum + ce
+            kl_sum = kl_sum + kl
+            if want_info:
+                if s0 == 0:
+                    p_end_first = p_first.item()
+                p_target.append(p_tgt)
+        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 want_kl else None), info
+
+    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:
+            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
+
+    # ------------------------------------------------------------------
+    # per-expert learning rate: the trainer hands us the LoRA network before it
+    # builds optimizer groups, so split the AR expert's modules into their own group
+    # ------------------------------------------------------------------
+    @property
+    def network(self):
+        return self._network
+
+    @network.setter
+    def network(self, value):
+        self._network = value
+        mult = getattr(self, "ar_lr_multiplier", 1.0)
+        if value is None or mult == 1.0 or getattr(value, "_yue2_lr_split", False):
+            return
+        orig = value.prepare_optimizer_params
+
+        def prepare_optimizer_params(text_encoder_lr=None, unet_lr=None, default_lr=None):
+            groups = orig(text_encoder_lr, unet_lr, default_lr)
+            ar_ids = set()
+            for lora in getattr(value, "unet_loras", []):
+                if ".ar." in lora.lora_name.replace("$$", "."):
+                    ar_ids.update(id(p) for p in lora.parameters())
+            out = []
+            n_ar = 0
+            for group in groups:
+                base_lr = group.get("lr", default_lr)
+                ar = [p for p in group["params"] if id(p) in ar_ids]
+                rest = [p for p in group["params"] if id(p) not in ar_ids]
+                if rest:
+                    out.append({**group, "params": rest})
+                if ar:
+                    n_ar += len(ar)
+                    if base_lr is not None:
+                        out.append({**group, "params": ar, "lr": base_lr * mult})
+                    else:
+                        out.append({**group, "params": ar})
+            if n_ar:
+                self.print_and_status_update(f"YuE2: AR expert LoRA learning rate x{mult} ({n_ar} tensors)")
+            return out
+
+        value.prepare_optimizer_params = prepare_optimizer_params
+        value._yue2_lr_split = True
+
+    def get_noise_prediction(self, latent_model_input: torch.Tensor, timestep: torch.Tensor, text_embeddings: PromptEmbeds, batch=None, **kwargs):
+        # `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
+        if model.device == torch.device("cpu"):
+            model.to(self.device_torch)
+        device = self.device_torch
+        t = (timestep.to(device, dtype=torch.float32) / 1000.0).to(self.torch_dtype)
+        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()
+        if torch.is_grad_enabled():
+            self._train_calls += 1
+            if self._train_calls - self._last_encode_call > 2 and (self.semantic_tokenizer is not None or self.transcriber is not None):
+                # latents are cached: MERT + SheetSage2 (~4 GB) are dead weight for the rest of training
+                self.semantic_tokenizer = None
+                self.transcriber = None
+                flush()
+        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])
+            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, 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, _ = 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])
+            preds.append(pred)
+        if lm_losses:
+            self._pending_ar_ce = torch.stack(lm_losses).mean()
+            self._pending_ar_kl = torch.stack(kl_losses).mean() if kl_losses else None
+            self._pending_aux_loss = self._pending_ar_ce * self.ar_loss_weight
+            if self._pending_ar_kl is not None:
+                self._pending_aux_loss = self._pending_aux_loss + self._pending_ar_kl * self.ar_kl_weight
+        return torch.cat(preds, 0)
+
+    # ------------------------------------------------------------------
+    # sampling
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return self.pipeline
+
+
+    def generate_single_audio(self, pipeline: YuE2Pipeline, gen_config: GenerateImageConfig, conditional_embeds: PromptEmbeds, unconditional_embeds, generator, extra):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+        if gen_config.output_ext not in ["mp3", "wav", "flac"]:
+            gen_config.output_ext = "mp3"
+        parsed = parse_caption(gen_config.prompt)
+        # sample-config duration field first, then a [Duration] prompt line, then the model default
+        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())
+        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_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):
+            token_bar.update(done - token_bar.n)
+            token_bar.set_postfix_str(f"{done / FRAMES_PER_SECOND:.1f}s", refresh=False)
+
+        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, legacy_off=self.cot == "off",
+            )
+        finally:
+            token_bar.close()
+        if len(codec) == 0:
+            codec = [0]
+        self._status_update(f"YuE2: rendering {len(codec) / FRAMES_PER_SECOND:.1f}s")
+        render_bar = tqdm(desc=f"  render {len(codec) / FRAMES_PER_SECOND:.0f}s", unit="step", position=1, leave=False, dynamic_ncols=True)
+
+        def render_step(i, n, x):
+            if render_bar.total != n:
+                render_bar.total = n
+            render_bar.update(i + 1 - render_bar.n)
+            self._emit_sample_step(x, i, n)
+
+        try:
+            latents = pipeline.synthesize(
+                prefix,
+                codec,
+                seed=gen_config.seed,
+                steps=gen_config.num_inference_steps,
+                step_callback=render_step,
+            )
+        finally:
+            render_bar.close()
+        audio = pipeline.decode(latents)  # [1, 2, samples]
+        return audio.cpu()
+
+    # ------------------------------------------------------------------
+    # LoRA key layout: transformer.nar.* -> diffusion_model.*, transformer.ar.* -> text_encoders.*
+    # ------------------------------------------------------------------
+    def convert_lora_weights_before_save(self, state_dict):
+        out = {}
+        for k, v in state_dict.items():
+            if k.startswith("transformer.nar."):
+                k = "diffusion_model." + k[len("transformer.nar.") :]
+            elif k.startswith("transformer.ar."):
+                k = "text_encoders." + k[len("transformer.ar.") :]
+            out[k] = v
+        return out
+
+    def convert_lora_weights_before_load(self, state_dict):
+        out = {}
+        for k, v in state_dict.items():
+            if k.startswith("diffusion_model."):
+                k = "transformer.nar." + k[len("diffusion_model.") :]
+            elif k.startswith("text_encoders."):
+                k = "transformer.ar." + k[len("text_encoders.") :]
+            out[k] = v
+        return out

extensions_built_in/captioner/Qwen3OmniCaptioner.py

@@ -29,6 +29,8 @@ VIDEO_FPS = 2
 
 # still-image files caption through the image pipeline (no audio, no frames)
 IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "bmp", "webp"}
+# audio-only files caption through the audio pipeline (no frames)
+AUDIO_EXTENSIONS = {"mp3", "wav", "flac", "ogg", "m4a", "aac"}
 
 # fixed generation ceiling under compiled decode: a constant max_length keeps
 # the static kv cache (and so the compiled decode graph) at one shape for
@@ -660,9 +662,15 @@ class Qwen3OmniCaptioner(BaseCaptioner):
     def _is_image_file(file_path: str) -> bool:
         return os.path.splitext(file_path)[1].lower().lstrip(".") in IMAGE_EXTENSIONS
 
+    @staticmethod
+    def _is_audio_file(file_path: str) -> bool:
+        return os.path.splitext(file_path)[1].lower().lstrip(".") in AUDIO_EXTENSIONS
+
     def _build_messages(self, _file_path: str):
         if self._is_image_file(_file_path):
             media = {"type": "image", "image": _file_path}
+        elif self._is_audio_file(_file_path):
+            media = {"type": "audio", "audio": _file_path}
         else:
             media = {"type": "video", "video": _file_path}
         return [
@@ -695,6 +703,10 @@ class Qwen3OmniCaptioner(BaseCaptioner):
 
             image = Image.open(file_path).convert("RGB")
             item = {"file": file_path, "kind": "image", "image": image, "audio": None}
+        elif self._is_audio_file(file_path):
+            from transformers.audio_utils import load_audio
+
+            item = {"file": file_path, "kind": "audio", "audio": load_audio(file_path, sampling_rate=16000)}
         else:
             from transformers.video_utils import load_video
             from transformers.audio_utils import load_audio
@@ -738,6 +750,13 @@ class Qwen3OmniCaptioner(BaseCaptioner):
                 padding=True,
                 size=self._size_kwargs(),
             )
+        if kind == "audio":
+            return self.processor(
+                text=[it["text"] for it in items],
+                audio=[it["audio"] for it in items],
+                return_tensors="pt",
+                padding=True,
+            )
         use_audio = kind == "video_audio"
         return self.processor(
             text=[it["text"] for it in items],
@@ -770,10 +789,14 @@ class Qwen3OmniCaptioner(BaseCaptioner):
         # under static cache, generate hands the forward a prepared 4D mask;
         # the true 2D padding mask is needed for the prefill rope index
         self.model._pad_mask_2d = inputs.get("attention_mask", None)
+        gen_kwargs = self._gen_kwargs(inputs["input_ids"].shape[1])
+        if items[0]["kind"] == "audio":
+            # greedy lyric transcription loops on syllables; mild penalty (the model's own talker default)
+            gen_kwargs["repetition_penalty"] = 1.05
         generated_ids = self.model.generate(
             **inputs,
             use_audio_in_video=use_audio,
-            **self._gen_kwargs(inputs["input_ids"].shape[1]),
+            **gen_kwargs,
         )
         trimmed = generated_ids[:, inputs["input_ids"].shape[1] :]
         captions = self.processor.batch_decode(
@@ -876,7 +899,7 @@ class Qwen3OmniCaptioner(BaseCaptioner):
                 futures.append((path, executor.submit(self._prep_media, path)))
 
             # batches must be homogeneous: the processor call differs per kind
-            buckets = {"image": [], "video_audio": [], "video_silent": []}
+            buckets = {"image": [], "audio": [], "video_audio": [], "video_silent": []}
             while futures:
                 if self.is_ui_captioner:
                     self.maybe_stop()

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,

jobs/process/BaseSDTrainProcess.py

@@ -354,6 +354,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 logger=self.logger,
                 num_frames=sample_item.num_frames,
                 fps=sample_item.fps,
+                duration=sample_item.duration,
                 ctrl_img=sample_item.ctrl_img,
                 ctrl_idx=sample_item.ctrl_idx,
                 ctrl_img_1=sample_item.ctrl_img_1,

testing/test_model_loading.py

@@ -191,6 +191,11 @@ MODEL_TESTS = {
         "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, "seed": 42},
     },
+    "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_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": {
         "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True},
         "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42},

toolkit/config_modules.py

@@ -57,6 +57,8 @@ class SampleItem:
         self.sample_steps: int = kwargs.get('sample_steps', sample_config.sample_steps)
         self.fps: int = kwargs.get('fps', sample_config.fps)
         self.num_frames: int = kwargs.get('num_frames', sample_config.num_frames)
+        # audio models: max seconds to generate
+        self.duration: Optional[float] = kwargs.get('duration', sample_config.duration)
         self.ctrl_img: Optional[str] = kwargs.get('ctrl_img', None)
         self.ctrl_idx: int = kwargs.get('ctrl_idx', 0)
         # for multi control image models
@@ -97,6 +99,7 @@ class SampleConfig:
         self.extra_values = kwargs.get('extra_values', [])
         self.num_frames = kwargs.get('num_frames', 1)
         self.fps: int = kwargs.get('fps', 16)
+        self.duration: Optional[float] = kwargs.get('duration', None)
         if self.num_frames > 1 and self.ext not in ['webp']:
             print("Changing sample extention to animated webp")
             self.ext = 'webp'
@@ -1155,6 +1158,7 @@ class GenerateImageConfig:
             ctrl_img_3: Optional[str] = None,  # third control image for multi control model
             num_frames: int = 1,
             fps: int = 15,
+            duration: Optional[float] = None,  # audio models: max seconds
             ctrl_idx: int = 0,
             do_cfg_norm: bool = False,
     ):
@@ -1187,6 +1191,7 @@ class GenerateImageConfig:
         self.extra_values = extra_values if extra_values is not None else []
         self.num_frames = num_frames
         self.fps = fps
+        self.duration = duration
         self.ctrl_img = ctrl_img
         self.ctrl_idx = ctrl_idx
         

toolkit/models/registry.py

@@ -55,6 +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_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/generate/page.tsx

@@ -13,7 +13,7 @@ import { startQueue } from '@/utils/queue';
 import useGPUInfo from '@/hooks/useGPUInfo';
 import usePollLoop from '@/hooks/usePollLoop';
 import { defaultInferenceJobConfig } from '@/helpers/inferenceJobConfig';
-import { encodeFilePathForUrl } from '@/utils/basic';
+import { encodeFilePathForUrl, objToTags, tagsToObj } from '@/utils/basic';
 import { latentToImage, payloadToFloat32, readEngineFrames, PreviewInfo } from '@/utils/engineStream';
 import { isMac } from '@/helpers/basic';
 import { modelArchs, getGenerateDefaults, GenerateDefaults } from '@/app/jobs/new/options';
@@ -842,7 +842,19 @@ function GeneratePageInner() {
                           className={`h-16 w-16 rounded-md overflow-hidden border-2 ${isSel ? 'border-blue-500' : 'border-transparent hover:border-gray-600'} bg-gray-800 block`}
                         >
                           {r.kind === 'audio' ? (
-                            <div className="w-full h-full flex items-center justify-center text-gray-400 text-[10px]">audio</div>
+                            <div className="relative w-full h-full flex items-center justify-center text-gray-400 text-[10px]">
+                              audio
+                              {/* embedded waveform cover; falls back to the label when the file has none */}
+                              {/* eslint-disable-next-line @next/next/no-img-element */}
+                              <img
+                                src={`/api/audio/art/${encodeFilePathForUrl(r.path)}`}
+                                alt=""
+                                className="absolute inset-0 w-full h-full object-cover"
+                                onError={e => {
+                                  (e.target as HTMLImageElement).style.display = 'none';
+                                }}
+                              />
+                            </div>
                           ) : r.kind === 'video' ? (
                             <video src={resultUrl(r)} muted className="w-full h-full object-cover" />
                           ) : (
@@ -1025,16 +1037,35 @@ function GeneratePageInner() {
                   </div>
                 </Card>
                 <Card title="Prompt" open={cardOpen('prompt')} onToggle={() => toggleCard('prompt')}>
-                  <TextAreaInput label="Prompt" value={sample.prompt || ''} onChange={v => setSample(s => ({ ...s, prompt: v }))} placeholder={modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'} />
-                  {(sample.guidance_scale ?? 4) > 1 && (
+                  {entry?.sampleTags ? (
+                    // tagged prompt archs (audio): one field per tag, stored as the tagged prompt string
+                    Object.entries(entry.sampleTags).map(([tagKey, tag]) => {
+                      const tags = tagsToObj(sample.prompt || '');
+                      const setTag = (v: any) => setSample(s => ({ ...s, prompt: objToTags({ ...tagsToObj(s.prompt || ''), [tagKey]: v }) }));
+                      const value = tags[tagKey] ?? '';
+                      if (tag.type === 'multiline') return <TextAreaInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                      if (tag.type === 'number') return <NumberInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                      return <TextInput key={tagKey} label={tag.title} value={value} onChange={setTag} placeholder={`Enter ${tag.title.toLowerCase()}`} />;
+                    })
+                  ) : (
+                    <TextAreaInput label="Prompt" value={sample.prompt || ''} onChange={v => setSample(s => ({ ...s, prompt: v }))} placeholder={modality === 'audio' ? 'song description, lyrics, bpm…' : 'describe what to generate'} />
+                  )}
+                  {(sample.guidance_scale ?? 4) > 1 && modality !== 'audio' && (
                     <TextAreaInput label="Negative prompt" value={sample.negative_prompt || ''} onChange={v => setSample(s => ({ ...s, negative_prompt: v }))} />
                   )}
                   <div className="grid grid-cols-3 gap-2">
-                    <NumberInput label="Width" value={sample.width ?? 1024} onChange={v => setSample(s => ({ ...s, width: v }))} min={64} max={4096} />
-                    <NumberInput label="Height" value={sample.height ?? 1024} onChange={v => setSample(s => ({ ...s, height: v }))} min={64} max={4096} />
+                    {modality !== 'audio' && (
+                      <>
+                        <NumberInput label="Width" value={sample.width ?? 1024} onChange={v => setSample(s => ({ ...s, width: v }))} min={64} max={4096} />
+                        <NumberInput label="Height" value={sample.height ?? 1024} onChange={v => setSample(s => ({ ...s, height: v }))} min={64} max={4096} />
+                      </>
+                    )}
                     <NumberInput label="Steps" value={sample.num_inference_steps ?? 25} onChange={v => setSample(s => ({ ...s, num_inference_steps: v }))} min={1} max={200} />
                     <NumberInput label="Guidance" value={sample.guidance_scale ?? 4} onChange={v => setSample(s => ({ ...s, guidance_scale: v }))} min={1} max={30} />
                     <NumberInput label="Seed (-1 random)" value={sample.seed ?? -1} onChange={v => setSample(s => ({ ...s, seed: v }))} min={-1} max={4294967295} />
+                    {modality === 'audio' && sample.duration !== undefined && (
+                      <NumberInput label="Duration (s)" value={sample.duration} onChange={v => setSample(s => ({ ...s, duration: v }))} min={1} max={600} />
+                    )}
                     {modality === 'video' && (
                       <>
                         <NumberInput label="Frames" value={sample.num_frames ?? 33} onChange={v => setSample(s => ({ ...s, num_frames: v }))} min={1} max={1000} />

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

@@ -1420,6 +1420,14 @@ export default function SimpleJob({
                   // automaticallt add the controls for a new dataset
                   const controls = modelArch?.controls ?? [];
                   newDataset.controls = controls;
+                  // arch dataset defaults (datasets[x].*) apply to added datasets too, not just at arch switch
+                  for (const key in modelArch?.defaults ?? {}) {
+                    const marker = 'datasets[x].';
+                    const idx = key.indexOf(marker);
+                    if (idx !== -1) {
+                      (newDataset as any)[key.slice(idx + marker.length)] = modelArch!.defaults![key][0];
+                    }
+                  }
                   setJobConfig([...jobConfig.config.process[0].datasets, newDataset], 'config.process[0].datasets');
                 }}
                 className="w-full px-4 py-2 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors"
@@ -1533,6 +1541,17 @@ export default function SimpleJob({
                   min={0}
                   required
                 />
+                {modelArch?.additionalSections?.includes('sample.duration') && (
+                  <NumberInput
+                    label="Duration (seconds)"
+                    value={jobConfig.config.process[0].sample.duration ?? 120}
+                    onChange={value => setJobConfig(value, 'config.process[0].sample.duration')}
+                    placeholder="eg. 120"
+                    className="pt-2"
+                    min={1}
+                    required
+                  />
+                )}
                 <Checkbox
                   label="Walk Seed"
                   className="pt-4 pl-2"

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

@@ -2,7 +2,12 @@ import React from 'react';
 import Link from 'next/link';
 import { GroupedSelectOption, SelectOption, JobConfig, ConfigDoc } from '@/types';
 import { defaultSliderConfig } from './jobConfig';
-import { defaultAudioSampleConfig, defaultSampleConfig, defaultIdeogramSamplesConfig } from '@/helpers/defaultSamples';
+import {
+  defaultAudioSampleConfig,
+  defaultSampleConfig,
+  defaultIdeogramSamplesConfig,
+  defaultYue2SampleConfig,
+} from '@/helpers/defaultSamples';
 
 type Control = 'depth' | 'line' | 'pose' | 'inpaint';
 
@@ -27,6 +32,7 @@ type AdditionalSections =
   | 'datasets.auto_frame_count'
   | 'sample.ctrl_img'
   | 'sample.multi_ctrl_imgs'
+  | 'sample.duration'
   | 'train.audio_loss_multiplier'
   | 'datasets.num_frames'
   | 'model.multistage'
@@ -80,6 +86,8 @@ export interface GenerateDefaults {
   sample: { [key: string]: any };
   needsControlImage: boolean;
   sizeLocked: boolean;
+  /** structured prompt fields (audio models): the prompt is their tagged form */
+  sampleTags?: SampleTags;
 }
 
 export interface ModelArch {
@@ -1370,6 +1378,109 @@ export const modelArchs: ModelArch[] = [
     disableSections: ['network.conv'],
     additionalSections: ['sample.multi_ctrl_imgs', 'model.low_vram', 'model.layer_offloading'],
   },
+  {
+    name: 'yue2',
+    label: 'YuE2',
+    group: 'audio',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': [
+        'Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors',
+        defaultNameOrPath,
+      ],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [false, false],
+      'config.process[0].model.low_vram': [false, false],
+      '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'],
+      // 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],
+      // audio has no resolution; every bucket would duplicate the whole dataset
+      'config.process[0].datasets[x].resolution': [[512], [512, 768, 1024]],
+      // 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': [{ 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,
+    modelNotes: (
+      <div className="space-y-2">
+        <p className="font-semibold text-amber-400">
+          Experimental. The AR (composition) model memorizes quickly and does not work well on small datasets. A handful
+          of songs is enough for it to learn the exact token sequence of each song; after that it stops generalizing and
+          free-running samples drift away from the training material. Expect to need a large, varied dataset for the AR
+          side to learn a style rather than the songs themselves.
+        </p>
+        <p>
+          YuE2 is two experts on one backbone. The <b>AR expert</b> reads the style line and lyrics and writes the song
+          as a sequence of semantic codec tokens (25 per second). The <b>NAR expert</b> then renders those tokens into
+          audio latents with flow matching. Training a LoRA here trains both: next-token loss on the AR over the whole
+          song from its start, flow loss on the NAR over a random window.
+        </p>
+        <p>
+          The AR loss is the one to watch (<code>loss/ar_ce</code>). It starts near 5 and, on a small dataset, falls
+          toward 0 within a few hundred steps, which is memorization. <code>ar_kl_weight</code> in model kwargs anchors
+          the AR to the base model so it cannot collapse onto the training songs; <code>ar_lr_multiplier</code> and a
+          smaller AR rank slow it further. Style comes mostly from the NAR, lyric following from the AR. Do not use
+          caption dropout: a blank prompt breaks lyric following.
+        </p>
+        <p>
+          The official audio-to-token encoder is unreleased. Training uses the community tokenizer by Kytra (
+          <a
+            href="https://x.com/sin_ceriously"
+            target="_blank"
+            rel="noreferrer"
+            className="text-blue-400 hover:underline"
+          >
+            @sin_ceriously
+          </a>
+          ), a MERT-v2-FullSong head that maps real audio to YuE2 codec tokens:{' '}
+          <a
+            href="https://huggingface.co/Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4"
+            target="_blank"
+            rel="noreferrer"
+            className="text-blue-400 hover:underline"
+          >
+            Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4
+          </a>
+          . It is downloaded on first use. Its tokens are an approximation of the model's own, so rendered samples will
+          not be bit-faithful to the training audio even when the AR replays a song exactly. Though it is extremely
+          close.
+        </p>
+        <p>
+          <b>ABC generation.</b> Generation is two stages: the AR first writes a lead sheet of the whole song in ABC
+          notation (sections, chords, vocal and instrumental melody), then writes the codec tokens conditioned on that
+          sheet. It can also run without a sheet ("off" mode), which uses a different instruction line. Samples here do
+          the two stages with <code>cot: full</code> (chords) or <code>cot: melody</code> (melody only), or the single
+          stage with <code>cot: off</code>.
+        </p>
+        <p>
+          <b>Training for both.</b> Every song is transcribed to an ABC sheet with SheetSage2 and the AR is trained on
+          lyrics to sheet and sheet to tokens. <code>abc_dropout</code> (default 0.5) is the fraction of training items
+          fed without the sheet instead, as an off-mode prompt, so one LoRA works in both modes. Set it to 0 to train
+          the sheet path only, or 1 to train off mode only.
+        </p>
+        <p>
+          <b>Caching is required.</b> The sheet is produced at latent-cache time and stored with the latents and codec
+          tokens, so Cache Latents to Disk must stay on: without it, SheetSage2 (about 12 s per song), the MERT
+          tokenizer and the VAE would run again on every training step. The cache records which mode built it; changing{' '}
+          <code>cot</code> means deleting the dataset's <code>_latent_cache</code> folder so the sheets are rebuilt.
+          Audio has no resolution, so keep a single resolution bucket per dataset or every bucket duplicates the songs.
+        </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
+          captions in this layout. Sample length is set by the Duration field in the sample section.
+        </p>
+      </div>
+    ),
+    // 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'],
+  },
   {
     name: 'ace_step_15',
     label: 'ACE-Step 1.5',
@@ -1802,7 +1913,6 @@ export const jobTypeOptions: JobTypeOption[] = [
   },
 ];
 
-
 const MODEL_PREFIX = 'config.process[0].model.';
 const SAMPLE_PREFIX = 'config.process[0].sample';
 // training-only model settings: the training adapter (e.g. Z-Image Turbo's
@@ -1832,13 +1942,15 @@ export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => {
       if (sc.guidance_scale !== undefined) sample.guidance_scale = sc.guidance_scale;
       if (sc.num_frames) sample.num_frames = sc.num_frames;
       if (sc.fps) sample.fps = sc.fps;
+      if (sc.duration) sample.duration = sc.duration;
       if (sc.neg) sample.negative_prompt = sc.neg;
     } else if (key.startsWith(SAMPLE_PREFIX + '.')) {
       const field = key.slice(SAMPLE_PREFIX.length + 1);
       if (value === undefined || value === null || value === '') continue;
       if (field === 'sample_steps') sample.num_inference_steps = value;
       else if (field === 'neg') sample.negative_prompt = value;
-      else if (['width', 'height', 'guidance_scale', 'num_frames', 'fps'].includes(field)) sample[field] = value;
+      else if (['width', 'height', 'guidance_scale', 'num_frames', 'fps', 'duration'].includes(field))
+        sample[field] = value;
     }
   }
   // the engine defaults to convrot8; the training default qtype is the
@@ -1847,7 +1959,8 @@ export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => {
   if (model.quantize_te && (!model.qtype_te || model.qtype_te === 'qfloat8')) model.qtype_te = 'convrot8';
   const sections = arch.additionalSections || [];
   const gen = arch.generate || {};
-  const modality: GenerateModality = gen.modality || (arch.group === 'audio' ? 'audio' : arch.isVideoModel ? 'video' : 'image');
+  const modality: GenerateModality =
+    gen.modality || (arch.group === 'audio' ? 'audio' : arch.isVideoModel ? 'video' : 'image');
   if (modality !== 'video') {
     delete sample.num_frames;
     delete sample.fps;
@@ -1859,7 +1972,9 @@ export const getGenerateDefaults = (arch: ModelArch): GenerateDefaults => {
     modality,
     model: { ...model, ...(gen.model || {}) },
     sample: { ...sample, ...(gen.sample || {}) },
-    needsControlImage: gen.needsControlImage ?? (sections.includes('sample.ctrl_img') || sections.includes('sample.multi_ctrl_imgs')),
+    needsControlImage:
+      gen.needsControlImage ?? (sections.includes('sample.ctrl_img') || sections.includes('sample.multi_ctrl_imgs')),
     sizeLocked: gen.sizeLocked ?? false,
+    sampleTags: arch.sampleTags,
   };
 };

ui/src/helpers/captionJobConfig.ts

@@ -17,7 +17,7 @@ export const defaultCaptionJobConfig: CaptionJobConfig = {
           quantize: true,
           qtype: 'convrot8',
           low_vram: false,
-          extensions: ['mp4', 'mov', 'webm', 'mkv', 'avi', 'jpg', 'jpeg', 'png', 'bmp', 'webp'],
+          extensions: ['mp4', 'mov', 'webm', 'mkv', 'avi', 'jpg', 'jpeg', 'png', 'bmp', 'webp', 'mp3', 'wav', 'flac', 'ogg'],
           path_to_caption: '',
           recaption: false,
           compile: true,

ui/src/helpers/captionOptions.ts

@@ -67,6 +67,17 @@ overall_soundscape and non_diegetic_music are always exactly N/A for a still ima
 
 Describe only what is actually visible. Be decisive. No preamble and no extra text - output only the three fields.`;
 
+// Captions a song as a YuE2 training caption: one line of comma-separated style
+// tags, then a [Lyrics] line and the verbatim lyrics with bracketed section
+// headers - the exact prefix layout the yue2 arch's parse_caption expects.
+const yue2CaptionPrompt = `Listen to this song and write a YuE2 training caption for it. Output exactly two parts and nothing else.
+
+Part 1, the first line only: comma-separated style tags describing the music. Cover, in this order where applicable: genre and sub-genre, mood, vocal type (male vocal, female vocal, duet, choir, rap, or instrumental), vocal delivery (breathy, belted, falsetto, spoken, whispered, harmonized), lead instruments and production (acoustic guitar, piano, synth pads, 808 bass, live drums, drum machine, strings, lo-fi, reverb-heavy), tempo feel (slow ballad, mid-tempo, upbeat, fast), and era or scene (90s alt rock, modern trap, bedroom pop). Use concrete lowercase tags, no sentences, no hedging, no artist or song names. Example: alternative rock, melancholic, male vocal, falsetto, electric guitar, piano, live drums, slow ballad, 2000s
+
+Part 2: on the next line write [Lyrics] and then the complete lyrics transcribed verbatim, one sung line per line. Split the song into sections with a bracketed header on its own line before each one, using these names: [Intro], [Verse 1], [Verse 2], [Pre-Chorus], [Chorus], [Bridge], [Instrumental], [Solo], [Outro]. A short descriptor may follow the name inside the brackets, e.g. [Intro choir] or [Chorus harmonized]. Leave one blank line between sections. Repeat a chorus each time it is sung, but write each repeated line once per time it is sung, never more. Wordless vocalizing (na, la, oh, ah, hums, vocal chops) is never transcribed syllable by syllable: describe it once in the section header instead, e.g. [Intro wordless vocals] or [Bridge oohs], and if you cannot make out real words in a passage, treat it as wordless. Do not add timestamps, speaker labels, quotation marks, translations, or commentary, and do not annotate ad-libs beyond the section header. If the song has no vocals at all, Part 2 is [Lyrics] followed by a single [Instrumental] line.
+
+Transcribe only what is actually sung. Be decisive. No preamble, no explanations, no markdown - output only the tag line and the lyrics block.`;
+
 // Editable ADDITIONAL INSTRUCTIONS block injected into the Ideogram system prompt.
 // Users can tweak this for dataset-specific guidance without altering the fixed
 // output contract, element/background rules, or bbox format.
@@ -127,7 +138,7 @@ export const captionerTypes: CaptionOption[] = [
         group: 'image/video/sound',
         defaults: {
             'config.process[0].caption.model_name_or_path': ['ai-toolkit/Qwen3-Omni-30B-A3B-Thinking', defaultNameOrPath],
-            'config.process[0].caption.extensions': [[...extensionsVideo, ...extensionsImage], defaultExtensions],
+            'config.process[0].caption.extensions': [[...extensionsVideo, ...extensionsImage, ...extensionsAudio], defaultExtensions],
             'config.process[0].caption.caption_prompt': [defaultVideoCaptionPrompt, undefined],
             'config.process[0].caption.max_res': [512, undefined],
             'config.process[0].caption.max_new_tokens': [512, undefined],
@@ -143,6 +154,7 @@ export const captionerTypes: CaptionOption[] = [
             'General': defaultVideoCaptionPrompt,
             'MiniMax H4 T2V': minimaxT2VCaptionPrompt,
             'MiniMax H4 Image': minimaxImageCaptionPrompt,
+            'YuE2': yue2CaptionPrompt,
         },
         additionalSections: [
             'caption.caption_prompt',

ui/src/helpers/defaultSamples.ts

@@ -977,3 +977,22 @@ export const defaultIdeogramSamplesConfig: SampleConfig = {
   num_frames: 1,
   fps: 1,
 }
+
+// YuE2 takes its native prompt layout (style text, a [Lyrics] line, the lyrics):
+// derive it from the shared audio samples instead of duplicating the lyrics
+const tagValue = (prompt: string, tag: string) => {
+  const m = prompt.match(new RegExp(`<${tag}>([\\s\\S]*?)</${tag}>`));
+  return m ? m[1].trim() : '';
+};
+
+export const defaultYue2SampleConfig: SampleConfig = {
+  ...defaultAudioSampleConfig,
+  samples: defaultAudioSampleConfig.samples.map(s => ({
+    ...s,
+    prompt: `${tagValue(s.prompt, 'CAPTION')}\n[Lyrics]\n${tagValue(s.prompt, 'LYRICS')}\n`,
+  })),
+  guidance_scale: 1,
+  sample_steps: 32,
+  // max seconds per sample; the AR stops earlier when the song ends
+  duration: 120,
+};

ui/src/types.ts

@@ -229,6 +229,7 @@ export interface SampleItem {
   sample_steps?: number;
   fps?: number;
   num_frames?: number;
+  duration?: number;
   ctrl_img?: string | null;
   ctrl_idx?: number;
   network_multiplier?: number;
@@ -252,6 +253,7 @@ export interface SampleConfig {
   sample_steps: number;
   num_frames: number;
   fps: number;
+  duration?: number;
 }
 
 export interface LoggingConfig {

version.py

@@ -1 +1 @@
-VERSION = "0.13.10"
+VERSION = "0.13.11"