diff --git a/extensions_built_in/audio_models/__init__.py b/extensions_built_in/audio_models/__init__.py index b01e265..552d72a 100644 --- a/extensions_built_in/audio_models/__init__.py +++ b/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, ] diff --git a/extensions_built_in/audio_models/yue2/__init__.py b/extensions_built_in/audio_models/yue2/__init__.py new file mode 100644 index 0000000..906018b --- /dev/null +++ b/extensions_built_in/audio_models/yue2/__init__.py @@ -0,0 +1 @@ +from .yue2_model import YuE2AudioModel diff --git a/extensions_built_in/audio_models/yue2/src/__init__.py b/extensions_built_in/audio_models/yue2/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/audio_models/yue2/src/model.py b/extensions_built_in/audio_models/yue2/src/model.py new file mode 100644 index 0000000..4bbcec0 --- /dev/null +++ b/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 diff --git a/extensions_built_in/audio_models/yue2/src/pipeline.py b/extensions_built_in/audio_models/yue2/src/pipeline.py new file mode 100644 index 0000000..af516d0 --- /dev/null +++ b/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) diff --git a/extensions_built_in/audio_models/yue2/src/tokenizer.py b/extensions_built_in/audio_models/yue2/src/tokenizer.py new file mode 100644 index 0000000..8d7209a --- /dev/null +++ b/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 diff --git a/extensions_built_in/audio_models/yue2/src/vae.py b/extensions_built_in/audio_models/yue2/src/vae.py new file mode 100644 index 0000000..9e34ea6 --- /dev/null +++ b/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) diff --git a/extensions_built_in/audio_models/yue2/yue2_model.py b/extensions_built_in/audio_models/yue2/yue2_model.py new file mode 100644 index 0000000..470e448 --- /dev/null +++ b/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: (style), , (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}>(.*?)", 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 ````/````/```` tags are still accepted.""" + if not isinstance(text, str): + text = "" + if "" in text or "" 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 diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 8220d8a..3d3bcaf 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/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, diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 43bae2c..50f1bfd 100644 --- a/testing/test_model_loading.py +++ b/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_bf16.safetensors", "quantize": True, "model_kwargs": {"sample_max_seconds": 15}}, + "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}, diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index c043fd0..14f8116 100644 --- a/toolkit/config_modules.py +++ b/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 diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index 64b8fa8..62bbeee 100644 --- a/toolkit/models/registry.py +++ b/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_bf16.safetensors", "quantize": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}}, } OUTPUT_EXT = {"image": "png", "video": "mp4", "audio": "mp3"} diff --git a/ui/src/app/generate/page.tsx b/ui/src/app/generate/page.tsx index 87b2e2c..d64f07c 100644 --- a/ui/src/app/generate/page.tsx +++ b/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' ? ( -
audio
+
+ audio + {/* embedded waveform cover; falls back to the label when the file has none */} + {/* eslint-disable-next-line @next/next/no-img-element */} + { + (e.target as HTMLImageElement).style.display = 'none'; + }} + /> +
) : r.kind === 'video' ? (