Initial experimental support for yue2

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,414 @@
+"""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):
+        return self.model.embed_tokens(ids.to(self.model.embed_tokens.weight.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 the Comfy-Org all-in-one checkpoint (raw keys)."""
+        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"}
+        for name, module, part in (("NAR", model.nar, nar_sd), ("AR", model.ar, ar_sd)):
+            missing, unexpected = module.load_state_dict({k: v.to(dtype) for k, v in part.items()}, strict=False)
+            if missing:
+                raise ValueError(f"YuE2 {name} missing keys: {missing[:5]} (+{max(0, len(missing) - 5)})")
+            if unexpected:
+                print(f"    YuE2 {name} unexpected: {len(unexpected)} (first 3: {unexpected[:3]})")
+        return model.to(dtype)
+
+
+@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)
+    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,253 @@
+"""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 (
+    CODEC_OFFSET,
+    CODEC_SIZE,
+    CONTEXT,
+    MUSIC_END,
+    KVCache,
+    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,
+) -> int:
+    scores = logits.float().clone()
+    allowed = torch.full_like(scores, -torch.inf)
+    allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0
+    allowed[MUSIC_END] = 0
+    scores += allowed
+    if step < min_tokens:
+        scores[MUSIC_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
+        # released "off" mode keeps the top 3 candidates regardless of mass
+        removed[:3] = 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)
+        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 _step(self):
+        ar = self.ar
+        x = ar.model.embed_tokens(self.token)
+        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.logits.copy_(ar.model.lm_head(ar.model.norm(x[:, -1])))
+
+    def capture(self):
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for _ in range(2):
+                self._step()
+        torch.cuda.current_stream().wait_stream(stream)
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph):
+            self._step()
+
+    def step(self, token: int):
+        self.token.fill_(token)
+        if self.graph is not None:
+            self.graph.replay()
+        else:
+            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_codec_tokens(
+        self,
+        prefix_embeds: torch.Tensor,  # [1, L, H]
+        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,
+    ) -> List[int]:
+        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)
+        for step in range(max_tokens):
+            token = sample_logits(
+                logits, history, step, temperature, top_p, top_k, repetition_penalty, penalty_window, min_tokens, generator
+            )
+            if token == MUSIC_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 [t - CODEC_OFFSET for t in history]
+
+    @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/tokenizer.py

@@ -0,0 +1,154 @@
+"""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_ids(self, style: str, lyrics: str, cot: str = "off", abc: str = "") -> List[int]:
+        """Everything the AR expert sees before the first codec token."""
+        prompt = f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n"
+        abc_ids = [] if cot == "off" else self.encode(abc)
+        return [EOD] + self.encode(prompt) + [ABC_START] + abc_ids + [ABC_END, MUSIC_START]
+
+
+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
+
+
+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):
+        super().__init__()
+        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 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,470 @@
+"""YuE2 (m-a-p/YuE2-3B) for ai-toolkit.
+
+Weights: the Comfy-Org all-in-one checkpoint (``checkpoints/yue2_3b_bf16.safetensors``)
+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 is
+merged into the base NAR on load (``model_kwargs.merge_nar_lora``, default true) since the
+head's tokens are a dialect the stock renderer was not trained on.
+
+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.
+
+Captions use the ace-step tag format: <CAPTION> (style), <LYRICS>, <DURATION> (max seconds
+for samples).
+"""
+
+import os
+import random
+import re
+from typing import List, Optional
+
+import torch
+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.print import print_acc
+from toolkit.prompt_utils import PromptEmbeds
+from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
+
+from .src.model import CODEC_OFFSET, FRAMES_PER_SECOND, MUSIC_END, YuE2Model, merge_nar_lora
+from .src.pipeline import YuE2Pipeline
+from .src.tokenizer import HEAD_FILE, HEAD_REPO, NAR_LORA_FILE, SemanticTokenizer, YuE2TextTokenizer
+from .src.vae import SAMPLE_RATE, YuE2VAE
+
+DEFAULT_CHECKPOINT = "Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors"
+
+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 _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:
+        return {"style": _tag(text, "CAPTION"), "lyrics": _tag(text, "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(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 prefix -> AR token embeddings. Holds no weights of its own; the
+    embedding table lives in the AR expert (looked up on whatever device it is on)."""
+
+    def __init__(self, model: YuE2Model, tokenizer: YuE2TextTokenizer, cot: str = "off"):
+        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.model.embed_tokens.weight.dtype
+
+    def prefix_ids(self, style: str, lyrics: str) -> List[int]:
+        return self.tokenizer.prefix_ids(style, lyrics, cot=self.cot)
+
+    @torch.no_grad()
+    def forward(self, style: str, lyrics: str) -> torch.Tensor:
+        ids = torch.tensor([self.prefix_ids(style, lyrics)], 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))
+        self._loss_log_every = int(kw.get("loss_log_every", 25))
+        self._loss_log_step = 0
+        self.merge_nar_lora = bool(kw.get("merge_nar_lora", True))
+        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}")
+        self.sample_max_seconds = float(kw.get("sample_max_seconds", 120.0))
+        self.sample_ar_temperature = float(kw.get("sample_ar_temperature", 1.0))
+        self.semantic_tokenizer: Optional[SemanticTokenizer] = None
+        self._pending_aux_loss = None
+
+    @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_named_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)
+        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_named_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)
+        return self.semantic_tokenizer
+
+    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]."""
+        if device is None:
+            device = self.vae_device_torch
+        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)
+        return DTO(latents, tokens=tokens[:, :n].to(latents.device, torch.int32))
+
+    # ------------------------------------------------------------------
+    # 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 scale_loss(self, loss):
+        aux = self._pending_aux_loss
+        self._pending_aux_loss = None
+        if aux is None:
+            return loss
+        self._loss_log_step += 1
+        if self._loss_log_every > 0 and self._loss_log_step % self._loss_log_every == 0:
+            # the trainer's bar shows the sum; the two terms move on different scales
+            info = getattr(self, "_pending_aux_info", "")
+            tqdm.write(f"yue2 step {self._loss_log_step}: flow {loss.detach().float().mean().item():.4f}  ar_ce {aux.detach().float().item() / max(self.ar_loss_weight, 1e-8):.4f}  {info}")
+        return loss + 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_info(self, logits: torch.Tensor, ids: torch.Tensor, prefix: torch.Tensor, total: int, whole_song: bool) -> str:
+        """Readout diagnostics: what the AR loss was computed over."""
+        if self._loss_log_every <= 0 or (self._loss_log_step + 1) % self._loss_log_every != 0:
+            return ""
+        with torch.no_grad():
+            p = logits.softmax(-1)
+            p_target = p.gather(1, ids[:, None]).squeeze(1)
+            return (f"[ar targets {ids.shape[0]} unique {ids.unique().numel()} (song {total}, whole {whole_song}) "
+                    f"prefix {prefix.shape[0]} p(END|prefix) {p[0, MUSIC_END].item():.3f} p(target) median {p_target.median().item():.3f}]")
+
+    def _ar_inputs(self, prefix: torch.Tensor, tokens: torch.Tensor, end_token: bool = True):
+        """One item: prefix (unpadded [L, H]) + codec tokens [+ MUSIC_END] -> embeds [1, L', H], ids [n]."""
+        ids = tokens.long() + CODEC_OFFSET
+        if end_token:
+            ids = torch.cat([ids, torch.tensor([MUSIC_END], device=tokens.device)])
+        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")
+        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()
+        preds, lm_losses = [], []
+        for i in range(latent_model_input.shape[0]):
+            song = tokens_all[i].to(device)
+            total = song.shape[0]
+            prefix = self._train_prefix(prefix_embeds[i], prefix_mask[i])
+            # NAR conditioning follows the released chunk protocol: prefix + window + MUSIC_END
+            embeds, ids = self._ar_inputs(prefix, song[start:end])
+            whole_song = start == 0 and end == total
+            if train_ar and not whole_song:
+                # the language-model loss must see the song from its beginning, or the AR
+                # learns that lyrics and tokens can start anywhere relative to each other
+                limit = total if self.ar_max_tokens <= 0 else min(total, self.ar_max_tokens)
+                ar_embeds, ar_ids = self._ar_inputs(prefix, song[:limit], end_token=limit == total)
+                _, hidden = model.ar.prefill(ar_embeds, return_hidden=True)
+                logits = model.ar.model.lm_head(hidden[0, -ar_ids.shape[0] - 1 : -1]).float()
+                lm_losses.append(torch.nn.functional.cross_entropy(logits, ar_ids))
+                self._pending_aux_info = self._ar_info(logits, ar_ids, prefix, total, whole_song)
+                with torch.no_grad():
+                    cache, _ = model.ar.prefill(embeds)
+            else:
+                cache, hidden = model.ar.prefill(embeds, return_hidden=train_ar)
+                if train_ar:
+                    logits = model.ar.model.lm_head(hidden[0, -ids.shape[0] - 1 : -1]).float()
+                    lm_losses.append(torch.nn.functional.cross_entropy(logits, ids))
+                    self._pending_aux_info = self._ar_info(logits, ids, prefix, total, whole_song)
+            # the flow loss must not train the AR through its KV cache: only next-token CE shapes the AR
+            cache = [(k.detach(), v.detach()) for k, v in cache]
+            pred = model.nar(latent_model_input[i : i + 1].to(device, self.torch_dtype), t[i : i + 1], cache, embeds.shape[1])
+            preds.append(pred)
+        if lm_losses:
+            self._pending_aux_loss = torch.stack(lm_losses).mean() * self.ar_loss_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())
+        prefix = conditional_embeds.text_embeds[:1, :length].to(self.device_torch, self.torch_dtype)
+
+        # nested bars under generate_images' "Generating Samples" bar (position 0)
+        self._status_update(f"YuE2: generating codec tokens (up to {max_seconds:.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, progress=token_progress,
+            )
+        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