diff --git a/README.md b/README.md index a98a60f..95c0760 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t ### Experimental - [lodestones/Zeta-Chroma](https://huggingface.co/lodestones/Zeta-Chroma) (Zeta Chroma) +- [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B) (YuE2) - audio. The official audio-to-token encoder is unreleased; training uses the community tokenizer by Kytra ([@sin_ceriously](https://x.com/sin_ceriously)), [Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4](https://huggingface.co/Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4). ## Installation 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..fb80541 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/model.py @@ -0,0 +1,449 @@ +"""YuE2 AR + NAR mixture-of-transformers backbone in the ComfyUI repack layout. + +Two 28-layer Qwen3-shaped experts share one architecture: + - ``YuE2AR`` (``text_encoders.*`` in the repack): token embedding, causal + layers, final norm and lm head. Writes ABC + semantic codec tokens. + - ``YuE2NAR`` (``model.diffusion_model.*``): latent in/out projections, + timestep embedder, sinusoidal frame position table and the same layer + stack. Predicts flow velocity on VAE latents while attending into the + AR expert's key/value cache over ``prefix + codec tokens + MUSIC_END``. + +Weights use the merged ``qkv_proj`` / ``gate_up_proj`` linears of the +Comfy-Org checkpoint so LoRA keys line up with ComfyUI's loader. +""" + +import math +from typing import List, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +import torch.utils.checkpoint as ckpt + +from toolkit.models.v2._mixin import OstrisModelMixin + +# vocabulary layout of the AR expert +EOD = 151643 +ABC_START, ABC_END = 151847, 151848 +MUSIC_START, MUSIC_END = 151851, 151852 +CODEC_OFFSET, CODEC_SIZE = 151853, 32768 +CONTEXT = 24576 +FRAMES_PER_SECOND = 25 +LATENT_DIM = 64 + +INSTRUCTIONS = { + "off": "Generate music with codec tokens from the given conditions.", + "melody": "Generate a melody-only ABC transcription without chord symbols, then generate music with codec tokens from the given conditions.", + "full": "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions.", +} + + +class YuE2Config: + hidden_size = 2048 + intermediate_size = 6144 + num_hidden_layers = 28 + num_attention_heads = 16 + num_key_value_heads = 8 + head_dim = 128 + vocab_size = 184704 + rms_norm_eps = 1e-6 + rope_theta = 1000000.0 + max_position_embeddings = CONTEXT + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x): + out = x.float() + out = out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps) + return out.to(x.dtype) * self.weight + + +def rope_cos_sin(positions: torch.Tensor, head_dim: int, theta: float): + """positions [B, L] -> cos, sin [B, L, head_dim // 2] in fp32.""" + inv_freq = 1.0 / ( + theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device) / head_dim) + ) + angles = positions.float().unsqueeze(-1) * inv_freq + return angles.cos(), angles.sin() + + +def apply_rope(x, cos, sin): + # x [B, H, L, D]; cos/sin [B, L, D/2] + half = x.shape[-1] // 2 + x1, x2 = x[..., :half], x[..., half:] + cos = cos.unsqueeze(1).to(x.dtype) + sin = sin.unsqueeze(1).to(x.dtype) + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) + + +class YuE2Attention(nn.Module): + def __init__(self, cfg: YuE2Config): + super().__init__() + self.num_heads = cfg.num_attention_heads + self.num_kv_heads = cfg.num_key_value_heads + self.head_dim = cfg.head_dim + inner = self.num_heads * self.head_dim + kv = self.num_kv_heads * self.head_dim + self.qkv_proj = nn.Linear(cfg.hidden_size, inner + 2 * kv, bias=False) + self.o_proj = nn.Linear(inner, cfg.hidden_size, bias=False) + self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) + + def project(self, x, cos, sin): + b, l, _ = x.shape + inner = self.num_heads * self.head_dim + kv = self.num_kv_heads * self.head_dim + q, k, v = self.qkv_proj(x).split((inner, kv, kv), dim=-1) + q = q.view(b, l, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(b, l, self.num_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(b, l, self.num_kv_heads, self.head_dim).transpose(1, 2) + q = apply_rope(self.q_norm(q), cos, sin) + k = apply_rope(self.k_norm(k), cos, sin) + return q, k, v + + def attend(self, q, k, v, causal: bool): + out = F.scaled_dot_product_attention( + q, k, v, is_causal=causal, enable_gqa=self.num_heads != self.num_kv_heads + ) + b, _, l, _ = out.shape + return self.o_proj(out.transpose(1, 2).reshape(b, l, -1)) + + +class YuE2MLP(nn.Module): + def __init__(self, cfg: YuE2Config): + super().__init__() + self.gate_up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size * 2, bias=False) + self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False) + + def forward(self, x): + gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + return self.down_proj(F.silu(gate) * up) + + +class YuE2Layer(nn.Module): + def __init__(self, cfg: YuE2Config): + super().__init__() + self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) + self.self_attn = YuE2Attention(cfg) + self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) + self.mlp = YuE2MLP(cfg) + + def forward( + self, + x, + cos, + sin, + prefix_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + causal: bool = True, + return_kv: bool = False, + ): + q, k, v = self.self_attn.project(self.input_layernorm(x), cos, sin) + present = (k, v) if return_kv else None + if prefix_kv is not None: + k = torch.cat([prefix_kv[0], k], dim=2) + v = torch.cat([prefix_kv[1], v], dim=2) + x = x + self.self_attn.attend(q, k, v, causal) + x = x + self.mlp(self.post_attention_layernorm(x)) + return x, present + + +class YuE2Stack(nn.Module): + """``model.layers`` + ``model.norm`` of one expert.""" + + def __init__(self, cfg: YuE2Config): + super().__init__() + self.layers = nn.ModuleList([YuE2Layer(cfg) for _ in range(cfg.num_hidden_layers)]) + self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) + + +class YuE2Expert(nn.Module): + """Shared expert base: LoRA targets this class name.""" + + def __init__(self, cfg: YuE2Config): + super().__init__() + self.cfg = cfg + self.gradient_checkpointing = False + + def _run_layer(self, layer, *args, **kwargs): + if self.gradient_checkpointing and torch.is_grad_enabled(): + return ckpt.checkpoint(layer, *args, use_reentrant=False, **kwargs) + return layer(*args, **kwargs) + + @property + def device(self): + return self.model.norm.weight.device + + @property + def dtype(self): + # norms stay in the activation dtype when the linears are quantized + return self.model.norm.weight.dtype + + +class YuE2AR(YuE2Expert): + def __init__(self, cfg: YuE2Config): + super().__init__(cfg) + self.model = YuE2Stack(cfg) + self.model.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) + self.model.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) + + def embed(self, ids: torch.Tensor): + # never touch embed_tokens.weight: the int8 checkpoint's table dequantizes the whole vocab on that property + return self.model.embed_tokens(ids.to(self.device)) + + def prefill(self, inputs_embeds: torch.Tensor, return_hidden: bool = False): + """Causal pass over [B, L, H]. Returns per-layer post-rope (k, v) as + [B, kv_heads, L, head_dim] plus (optionally) the final normed hidden.""" + b, l, _ = inputs_embeds.shape + positions = torch.arange(l, device=inputs_embeds.device)[None].expand(b, -1) + cos, sin = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta) + x = inputs_embeds + cache: List[Tuple[torch.Tensor, torch.Tensor]] = [] + for layer in self.model.layers: + x, kv = self._run_layer(layer, x, cos, sin, None, True, True) + cache.append(kv) + hidden = self.model.norm(x) if return_hidden else None + return cache, hidden + + @torch.no_grad() + def decode_step(self, token_embeds: torch.Tensor, cache: "KVCache"): + """One new token per batch row against a preallocated cache.""" + b = token_embeds.shape[0] + pos = torch.full((b, 1), cache.length, device=token_embeds.device, dtype=torch.long) + cos, sin = rope_cos_sin(pos, self.cfg.head_dim, self.cfg.rope_theta) + x = token_embeds + for i, layer in enumerate(self.model.layers): + q, k, v = layer.self_attn.project(layer.input_layernorm(x), cos, sin) + k, v = cache.append(i, k, v) + x = x + layer.self_attn.attend(q, k, v, causal=False) + x = x + layer.mlp(layer.post_attention_layernorm(x)) + cache.length += 1 + return self.model.lm_head(self.model.norm(x[:, -1])) + + @torch.no_grad() + def prefill_into_cache(self, inputs_embeds: torch.Tensor, cache: "KVCache"): + b, l, _ = inputs_embeds.shape + positions = torch.arange(l, device=inputs_embeds.device)[None].expand(b, -1) + cos, sin = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta) + x = inputs_embeds + for i, layer in enumerate(self.model.layers): + q, k, v = layer.self_attn.project(layer.input_layernorm(x), cos, sin) + cache.write(i, 0, k, v) + x = x + layer.self_attn.attend(q, k, v, causal=True) + x = x + layer.mlp(layer.post_attention_layernorm(x)) + cache.length = l + return self.model.lm_head(self.model.norm(x[:, -1])) + + +class KVCache: + """Preallocated [B, kv_heads, capacity, head_dim] per layer for AR decode.""" + + def __init__(self, cfg: YuE2Config, batch: int, capacity: int, device, dtype): + shape = (batch, cfg.num_key_value_heads, capacity, cfg.head_dim) + self.k = [torch.empty(shape, device=device, dtype=dtype) for _ in range(cfg.num_hidden_layers)] + self.v = [torch.empty(shape, device=device, dtype=dtype) for _ in range(cfg.num_hidden_layers)] + self.capacity = capacity + self.length = 0 + + def write(self, layer, start, k, v): + n = k.shape[2] + self.k[layer][:, :, start : start + n] = k + self.v[layer][:, :, start : start + n] = v + + def append(self, layer, k, v): + end = self.length + k.shape[2] + if end > self.capacity: + raise ValueError("YuE2 KV cache capacity exceeded") + self.write(layer, self.length, k, v) + return self.k[layer][:, :, :end], self.v[layer][:, :, :end] + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size, freq_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(freq_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size) + ) + self.freq_size = freq_size + + def forward(self, t: torch.Tensor): + half = self.freq_size // 2 + freqs = torch.exp( + -math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half + ) + args = t.float().unsqueeze(-1) * freqs.unsqueeze(0) + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + return self.mlp(emb.to(self.mlp[0].weight.dtype)) + + +class AudioPositionEmbedding(nn.Module): + def __init__(self, max_frames, hidden_size): + super().__init__() + self.register_buffer("pe", torch.zeros(max_frames, hidden_size)) + + def forward(self, length: int): + return self.pe[:length] + + +class YuE2NAR(YuE2Expert): + def __init__(self, cfg: YuE2Config): + super().__init__(cfg) + self.model = YuE2Stack(cfg) + self.vae2llm = nn.Linear(LATENT_DIM, cfg.hidden_size) + self.llm2vae = nn.Linear(cfg.hidden_size, LATENT_DIM) + self.time_embedder = TimestepEmbedder(cfg.hidden_size) + self.latent_pos_embed = AudioPositionEmbedding(cfg.max_position_embeddings, cfg.hidden_size) + + def forward( + self, + noisy_latents: torch.Tensor, # [B, T, 64] + t: torch.Tensor, # [B] in 0..1, 1 = pure noise + prefix_cache: List[Tuple[torch.Tensor, torch.Tensor]], + prefix_len: int, + ): + b, frames, _ = noisy_latents.shape + n = frames + 2 + x = F.pad(noisy_latents, (0, 0, 1, 1)) # START / END slots + x = self.vae2llm(x.to(self.vae2llm.weight.dtype)) + x = x + self.time_embedder(t)[:, None] + x = x + self.latent_pos_embed(n)[None].to(x.dtype) + positions = torch.arange(prefix_len, prefix_len + n, device=x.device)[None].expand(b, -1) + cos, sin = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta) + for layer, kv in zip(self.model.layers, prefix_cache): + x, _ = self._run_layer(layer, x, cos, sin, kv, False, False) + return self.llm2vae(self.model.norm(x))[:, 1:-1] + + +class YuE2Model(nn.Module, OstrisModelMixin): + """Both experts under one module so a single LoRA network covers them.""" + + aitk_comfy_repo = "Comfy-Org/YuE2" + + def __init__(self, cfg: Optional[YuE2Config] = None): + super().__init__() + cfg = cfg or YuE2Config() + self.cfg = cfg + self.ar = YuE2AR(cfg) + self.nar = YuE2NAR(cfg) + + @classmethod + def get_transformer_block_names(cls): + return ["ar.model.layers", "nar.model.layers"] + + @classmethod + def get_quantization_exclude_modules(cls): + return ["ar.model.embed_tokens", "ar.model.lm_head", "nar.vae2llm", "nar.llm2vae", "nar.time_embedder*"] + + @property + def device(self): + return self.nar.device + + @property + def dtype(self): + return self.nar.dtype + + def enable_gradient_checkpointing(self): + self.ar.gradient_checkpointing = True + self.nar.gradient_checkpointing = True + + def disable_gradient_checkpointing(self): + self.ar.gradient_checkpointing = False + self.nar.gradient_checkpointing = False + + @classmethod + def load_from_state_dict(cls, sd: dict, dtype=torch.bfloat16): + """``sd`` is a Comfy-Org all-in-one checkpoint (raw keys): bf16, or the int8 convrot repack whose + ``comfy_quant`` markers attach the shipped quantization directly (no requantization).""" + from toolkit.util.comfy_quant_import import Int8Embedding, import_comfy_quantized_layers + from toolkit.util.ostris_quant import OstrisLinear + + model = cls() + nar_sd = {k[len("model.diffusion_model.") :]: v for k, v in sd.items() if k.startswith("model.diffusion_model.")} + ar_sd = {k[len("text_encoders.") :]: v for k, v in sd.items() if k.startswith("text_encoders.") and k != "text_encoders.yue2_tokenizer_json"} + prequantized = any(k.endswith(".comfy_quant") for k in sd) + for name, module, part in (("NAR", model.nar, nar_sd), ("AR", model.ar, ar_sd)): + allowed_missing = set() + if prequantized: + quantized_paths = {k[: -len(".comfy_quant")] for k in part if k.endswith(".comfy_quant")} + part, n = import_comfy_quantized_layers(module, part, orig_dtype=dtype) + for mod_name, m in module.named_modules(): + if isinstance(m, OstrisLinear): + allowed_missing.add(f"{mod_name}.weight") + if m.bias is not None: + allowed_missing.add(f"{mod_name}.bias") + m.bias.data = m.bias.data.to(dtype) + elif isinstance(m, (Int8Embedding, nn.Embedding)) and mod_name in quantized_paths: + allowed_missing.add(f"{mod_name}.weight") + part = {k: (v.to(dtype) if v.is_floating_point() else v) for k, v in part.items()} + missing, unexpected = module.load_state_dict(part, strict=False) + missing = [k for k in missing if k not in allowed_missing] + if missing: + raise ValueError(f"YuE2 {name} missing keys: {missing[:5]} (+{max(0, len(missing) - 5)})") + if unexpected: + print(f" YuE2 {name} unexpected: {len(unexpected)} (first 3: {unexpected[:3]})") + if not prequantized: + return model.to(dtype) + # cast the remaining float tensors (norms, embeddings, buffers) without touching the quantized layers + for m in model.modules(): + if isinstance(m, (OstrisLinear, Int8Embedding)): + continue + for pname, prm in list(m.named_parameters(recurse=False)): + if prm.is_floating_point(): + prm.data = prm.data.to(dtype) + for bname, buf in list(m.named_buffers(recurse=False)): + if buf is not None and buf.is_floating_point(): + setattr(m, bname, buf.to(dtype)) + model.aitk_is_quantized = True # aitk_post_load keeps the shipped convrot8 layers when convrot8 is requested + return model + + +@torch.no_grad() +def merge_nar_lora(model: YuE2Model, ckpt_path: str, scale: float = 1.0): + """Fold the community NAR adapter (unmerged q/k/v + gate/up LoRA pairs, + plus full vae2llm/llm2vae weights) into the merged-projection base.""" + try: + ck = torch.load(ckpt_path, map_location="cpu", weights_only=True) + except Exception: + ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) + if getattr(model, "aitk_is_quantized", False): + raise ValueError("merge_nar_lora needs the bf16 checkpoint; the int8 convrot repack cannot take merged weights") + tensors = iter(ck["lora"]) + cfg = model.cfg + inner = cfg.num_attention_heads * cfg.head_dim + kv = cfg.num_key_value_heads * cfg.head_dim + qkv_rows = {"q_proj": (0, inner), "k_proj": (inner, inner + kv), "v_proj": (inner + kv, inner + 2 * kv)} + gu_rows = {"gate_proj": (0, cfg.intermediate_size), "up_proj": (cfg.intermediate_size, 2 * cfg.intermediate_size)} + for layer in model.nar.model.layers: + for name in ("q_proj", "k_proj", "v_proj", "o_proj"): + A = next(tensors).float() + B = next(tensors).float() + delta = (B @ A) * scale + if name == "o_proj": + w = layer.self_attn.o_proj.weight + w.add_(delta.to(w.device, w.dtype)) + else: + r0, r1 = qkv_rows[name] + w = layer.self_attn.qkv_proj.weight + w[r0:r1].add_(delta.to(w.device, w.dtype)) + for name in ("gate_proj", "up_proj", "down_proj"): + A = next(tensors).float() + B = next(tensors).float() + delta = (B @ A) * scale + if name == "down_proj": + w = layer.mlp.down_proj.weight + w.add_(delta.to(w.device, w.dtype)) + else: + r0, r1 = gu_rows[name] + w = layer.mlp.gate_up_proj.weight + w[r0:r1].add_(delta.to(w.device, w.dtype)) + io = ck.get("io", {}) + for name in ("vae2llm", "llm2vae"): + if name in io: + mod = getattr(model.nar, name) + for pname, value in io[name].items(): + getattr(mod, pname).copy_(value.to(mod.weight.device, mod.weight.dtype)) + return model 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..e17d884 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/pipeline.py @@ -0,0 +1,316 @@ +"""Sampling for YuE2: AR codec-token generation, then NAR flow matching. + +Follows the released protocol: semantic sampling with temperature 1.0, +top-p 0.95, top-k 100, repetition penalty 1.2 over the last 50 tokens, at +least 200 tokens before MUSIC_END; acoustic synthesis with a midpoint ODE +solver from t=1 (noise) to t=0, one AR prefill per context chunk. +""" + +from typing import Callable, List, Optional, Tuple + +import torch + +from .model import ( + ABC_END, + CODEC_OFFSET, + CODEC_SIZE, + CONTEXT, + EOD, + KVCache, + MUSIC_END, + YuE2Model, + rope_cos_sin, +) +from .vae import HOP, YuE2VAE + + +def chunk_ranges(frames: int, prefix_tokens: int, context: int = CONTEXT) -> List[Tuple[int, int]]: + size = (context - prefix_tokens - 3) // 2 + if frames < 1 or size < 1: + raise ValueError("YuE2 needs codec tokens and enough context for at least one acoustic frame.") + return [(start, min(start + size, frames)) for start in range(0, frames, size)] + + +def sample_logits( + logits: torch.Tensor, + history: List[int], + step: int, + temperature: float, + top_p: float, + top_k: int, + repetition_penalty: float, + penalty_window: int, + min_tokens: int, + generator: torch.Generator, + phase: str = "semantic", + legacy_off: bool = False, +) -> int: + """ComfyUI's ``distribution`` for one row: ``abc`` allows the text vocabulary + ABC_END, + ``semantic`` the codec ids + MUSIC_END. ``legacy_off`` = cot "off" (bf16 scores, keep top 3).""" + scores = logits.clone() if legacy_off else logits.float().clone() + end = ABC_END if phase == "abc" else MUSIC_END + allowed = torch.full_like(scores, -torch.inf) + if phase == "abc": + allowed[:EOD] = 0 + else: + allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 + allowed[end] = 0 + scores += allowed + if step < min_tokens: + scores[end] = -torch.inf + if repetition_penalty != 1.0 and history: + recent = torch.tensor(history[-penalty_window:], dtype=torch.long, device=scores.device) + counts = torch.zeros_like(scores).scatter_add_(0, recent, torch.ones_like(recent, dtype=scores.dtype)) + penalty = repetition_penalty ** counts + scores = torch.where(scores < 0, scores * penalty, scores / penalty) + if temperature == 0: + return int(scores.argmax().item()) + scores /= temperature + threshold = scores.topk(min(top_k, scores.shape[-1])).values[-1] + scores.masked_fill_(scores < threshold, -torch.inf) + if top_p < 1: + values, indices = scores.sort(descending=True) + probs = values.softmax(-1) + removed = probs.cumsum(-1) - probs > top_p + removed[: 3 if legacy_off else 1] = False + values.masked_fill_(removed, -torch.inf) + scores = values.scatter(-1, indices, values) + probs = scores.softmax(-1) + return int(torch.multinomial(probs.to(generator.device), 1, generator=generator).item()) + + +class ARDecodeGraph: + """One-token AR decode step captured as a CUDA graph over static buffers: + token id in, logits out, preallocated [1, kv_heads, capacity, head_dim] + cache per layer. Attention is hand-grouped (GQA) with a length mask so the + whole step is capturable; the eager KVCache path stays as the fallback.""" + + def __init__(self, ar, capacity: int, device, dtype): + cfg = ar.cfg + self.ar = ar + self.capacity = capacity + self.kv_heads = cfg.num_key_value_heads + self.groups = cfg.num_attention_heads // cfg.num_key_value_heads + self.head_dim = cfg.head_dim + shape = (1, self.kv_heads, capacity, self.head_dim) + self.k = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(cfg.num_hidden_layers)] + self.v = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(cfg.num_hidden_layers)] + self.token = torch.zeros((1, 1), device=device, dtype=torch.long) + self.pos = torch.zeros((1, 1), device=device, dtype=torch.long) + self.index = torch.arange(capacity, device=device) + # the graph covers the layer stack only; embedding lookup and lm_head run outside it + # (the int8 repack's embedding/lm_head kernels invalidate stream capture) + self.x_in = torch.zeros((1, 1, cfg.hidden_size), device=device, dtype=dtype) + self.h_out = torch.zeros((1, cfg.hidden_size), device=device, dtype=dtype) + self.logits = torch.empty((1, cfg.vocab_size), device=device, dtype=dtype) + self.graph = None + + def load_prefill(self, cache, length: int): + for i, (k, v) in enumerate(cache): + self.k[i][:, :, :length] = k + self.v[i][:, :, :length] = v + self.pos.fill_(length) + + def _layers(self): + ar = self.ar + x = self.x_in + cos, sin = rope_cos_sin(self.pos, self.head_dim, ar.cfg.rope_theta) + mask = (self.index <= self.pos)[None, None] # [1, 1, 1, capacity] + scale = self.head_dim ** -0.5 + slot = self.pos.view(-1) + for i, layer in enumerate(ar.model.layers): + q, k, v = layer.self_attn.project(layer.input_layernorm(x), cos, sin) + self.k[i].index_copy_(2, slot, k) + self.v[i].index_copy_(2, slot, v) + qg = q.view(1, self.kv_heads, self.groups, self.head_dim) + scores = torch.matmul(qg, self.k[i].transpose(-1, -2)) * scale # [1, kv, groups, capacity] + scores = scores.float().masked_fill(~mask, float("-inf")).softmax(-1).to(q.dtype) + out = torch.matmul(scores, self.v[i]).reshape(1, 1, -1) + x = x + layer.self_attn.o_proj(out) + x = x + layer.mlp(layer.post_attention_layernorm(x)) + self.h_out.copy_(ar.model.norm(x[:, -1])) + + def _step(self): + ar = self.ar + self.x_in.copy_(ar.model.embed_tokens(self.token).to(self.x_in.dtype)) + if self.graph is not None: + self.graph.replay() + else: + self._layers() + self.logits.copy_(ar.model.lm_head(self.h_out)) + + def capture(self): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(2): + self._layers() + torch.cuda.current_stream().wait_stream(stream) + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self._layers() + + def step(self, token: int): + self.token.fill_(token) + self._step() + self.pos.add_(1) + return self.logits[0] + + +class YuE2Pipeline: + def __init__(self, model: YuE2Model, vae: YuE2VAE): + self.model = model + self.vae = vae + self.do_tiled_decoding = True + self.use_cuda_graphs = True + + @torch.no_grad() + def generate_tokens( + self, + prefix_embeds: torch.Tensor, # [1, L, H] + phase: str, + max_tokens: int, + seed: int, + temperature: float, + top_p: float, + top_k: int, + repetition_penalty: float, + penalty_window: int, + min_tokens: int, + progress: Optional[Callable[[int, int], None]] = None, + legacy_off: bool = False, + ) -> List[int]: + """Sample AR tokens after ``prefix_embeds`` until the phase's end token (raw vocab ids).""" + ar = self.model.ar + device = prefix_embeds.device + prefix_len = prefix_embeds.shape[1] + if prefix_len + max_tokens > CONTEXT: + raise ValueError("YuE2 prompt plus token budget exceeds the model context") + # reference sampler: seeded generator on the compute device, multinomial over the filtered softmax + generator = torch.Generator(device=device).manual_seed(seed) + capacity = prefix_len + max_tokens + graph = None + if self.use_cuda_graphs and device.type == "cuda": + try: + graph = ARDecodeGraph(ar, capacity, device, ar.dtype) + prefill_cache, hidden = ar.prefill(prefix_embeds.to(ar.dtype), return_hidden=True) + logits = ar.model.lm_head(hidden[:, -1])[0] + graph.load_prefill(prefill_cache, prefix_len) + del prefill_cache, hidden + graph.capture() + except Exception as e: # capture unsupported for these ops/kernels: eager path + print(f"YuE2: CUDA graph decode unavailable ({type(e).__name__}: {e}); using eager decode") + graph = None + torch.cuda.empty_cache() + if graph is None: + cache = KVCache(self.model.cfg, 1, capacity, device, ar.dtype) + logits = ar.prefill_into_cache(prefix_embeds.to(ar.dtype), cache)[0] + history: List[int] = [] + min_tokens = min(min_tokens, max_tokens) + end = ABC_END if phase == "abc" else MUSIC_END + for step in range(max_tokens): + token = sample_logits( + logits, history, step, temperature, top_p, top_k, repetition_penalty, penalty_window, min_tokens, generator, + phase=phase, legacy_off=legacy_off, + ) + if token == end: + break + history.append(token) + if progress is not None: + progress(step + 1, max_tokens) + if step + 1 < max_tokens: + if graph is not None: + logits = graph.step(token) + else: + emb = ar.embed(torch.tensor([[token]], device=device)) + logits = ar.decode_step(emb, cache)[0] + if graph is None: + del cache + del graph + return history + + def generate_codec_tokens( + self, + prefix_embeds: torch.Tensor, + max_tokens: int, + seed: int, + temperature: float = 1.0, + top_p: float = 0.95, + top_k: int = 100, + repetition_penalty: float = 1.2, + penalty_window: int = 50, + min_tokens: int = 200, + progress: Optional[Callable[[int, int], None]] = None, + legacy_off: bool = False, + ) -> List[int]: + """Music phase (ComfyUI defaults). Returns codec ids without the vocab offset.""" + ids = self.generate_tokens( + prefix_embeds, "semantic", max_tokens, seed, temperature, top_p, top_k, repetition_penalty, penalty_window, + min_tokens, progress, legacy_off=legacy_off, + ) + return [t - CODEC_OFFSET for t in ids] + + def generate_abc_tokens( + self, + prefix_embeds: torch.Tensor, + seed: int, + max_tokens: int = 8192, + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 30, + repetition_penalty: float = 1.005, + penalty_window: int = 100, + progress: Optional[Callable[[int, int], None]] = None, + ) -> List[int]: + """ABC phase (ComfyUI "Generate ABC" defaults): text ids of the sheet, without ABC_END.""" + return self.generate_tokens( + prefix_embeds, "abc", max_tokens, seed, temperature, top_p, top_k, repetition_penalty, penalty_window, + min(32, max_tokens), progress, + ) + + @torch.no_grad() + def synthesize( + self, + prefix_embeds: torch.Tensor, # [1, L, H] + codec: List[int], + seed: int, + steps: int = 32, + step_callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, + ) -> torch.Tensor: + """-> latents [1, T, 64] in the NAR's dtype.""" + ar, nar = self.model.ar, self.model.nar + device = prefix_embeds.device + dtype = nar.dtype + generator = torch.Generator(device="cpu").manual_seed(seed) + noise = torch.randn((len(codec), 64), generator=generator).to(device, dtype) + ranges = chunk_ranges(len(codec), prefix_embeds.shape[1]) + outputs = [] + total_steps = steps * len(ranges) + for ci, (start, end) in enumerate(ranges): + ids = torch.tensor([[c + CODEC_OFFSET for c in codec[start:end]] + [MUSIC_END]], device=device) + embeds = torch.cat([prefix_embeds.to(ar.dtype), ar.embed(ids)], dim=1) + cache, _ = ar.prefill(embeds) + prefix_len = embeds.shape[1] + state = noise[start:end][None] + dt = 1.0 / steps + for step in range(steps): + t = 1.0 - step * dt + v1 = nar(state, torch.full((1,), t, device=device, dtype=dtype), cache, prefix_len) + mid = state - v1 * (dt / 2) + v2 = nar(mid, torch.full((1,), t - dt / 2, device=device, dtype=dtype), cache, prefix_len) + state = state - v2 * dt + if step_callback is not None: + step_callback(ci * steps + step, total_steps, state) + outputs.append(state) + del cache + return torch.cat(outputs, dim=1) + + @torch.no_grad() + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """[1, T, 64] -> waveform [1, 2, samples] float32.""" + z = latents.transpose(1, 2).to(self.vae.device, self.vae.dtype) + if self.do_tiled_decoding: + audio = self.vae.tiled_decode(z) + else: + audio = self.vae.decode(z) + return audio.float().clamp(-1, 1) diff --git a/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py b/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py new file mode 100644 index 0000000..5ae0200 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/sheetsage_decoder.py @@ -0,0 +1,110 @@ +"""SheetSage2's score decoder (BART-style, post-LN, learned positions with offset 2) in plain +torch, so the m-a-p/SheetSage2 remote code runs on current transformers without BartDecoder. +Parameter names follow the released checkpoint (decoder.embed_positions, decoder.layers.N.self_attn.*, +encoder_attn.*, fc1/fc2, *_layer_norm, layernorm_embedding). KV cache = per-layer tuples with the +batch dimension first, which the upstream generation loop can index.""" + +from types import SimpleNamespace +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +POSITION_OFFSET = 2 # BART learned positional embedding convention + + +class _Attention(nn.Module): + def __init__(self, dim: int, heads: int): + super().__init__() + self.heads = heads + self.head_dim = dim // heads + self.q_proj = nn.Linear(dim, dim) + self.k_proj = nn.Linear(dim, dim) + self.v_proj = nn.Linear(dim, dim) + self.out_proj = nn.Linear(dim, dim) + + def _split(self, x): + b, l, _ = x.shape + return x.view(b, l, self.heads, self.head_dim).transpose(1, 2) + + def forward(self, x, kv_source=None, mask=None, past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, cache_static=False): + q = self._split(self.q_proj(x)) + if cache_static and past is not None: + k, v = past # cross attention: memory keys never change + else: + src = x if kv_source is None else kv_source + k, v = self._split(self.k_proj(src)), self._split(self.v_proj(src)) + if past is not None and not cache_static: + k, v = torch.cat([past[0], k], 2), torch.cat([past[1], v], 2) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + b, _, l, _ = out.shape + return self.out_proj(out.transpose(1, 2).reshape(b, l, -1)), (k, v) + + +class _Layer(nn.Module): + def __init__(self, dim: int, heads: int, ffn: int): + super().__init__() + self.self_attn = _Attention(dim, heads) + self.self_attn_layer_norm = nn.LayerNorm(dim) + self.encoder_attn = _Attention(dim, heads) + self.encoder_attn_layer_norm = nn.LayerNorm(dim) + self.fc1 = nn.Linear(dim, ffn) + self.fc2 = nn.Linear(ffn, dim) + self.final_layer_norm = nn.LayerNorm(dim) + + def forward(self, x, memory, self_mask, past): + self_past = None if past is None else past[:2] + cross_past = None if past is None else past[2:] + h, self_kv = self.self_attn(x, mask=self_mask, past=self_past) + x = self.self_attn_layer_norm(x + h) + h, cross_kv = self.encoder_attn(x, kv_source=memory, past=cross_past, cache_static=True) + x = self.encoder_attn_layer_norm(x + h) + x = self.final_layer_norm(x + self.fc2(F.gelu(self.fc1(x)))) + return x, self_kv + cross_kv + + +class ScoreDecoder(nn.Module): + """Drop-in for the BartDecoder the upstream model builds: ``ScoreDecoder(bart_config, embed_tokens=...)``.""" + + def __init__(self, config, embed_tokens: nn.Embedding): + super().__init__() + dim, heads = config.d_model, config.decoder_attention_heads + self.embed_tokens = embed_tokens + self.embed_positions = nn.Embedding(config.max_position_embeddings + POSITION_OFFSET, dim) + self.layernorm_embedding = nn.LayerNorm(dim) + self.layers = nn.ModuleList([_Layer(dim, heads, config.decoder_ffn_dim) for _ in range(config.decoder_layers)]) + self.pad_token_id = config.pad_token_id + + def gradient_checkpointing_disable(self): + pass + + def forward(self, input_ids, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, + past_key_values=None, use_cache=False, output_hidden_states=False, return_dict=True, **_): + b, l = input_ids.shape + past_len = 0 if past_key_values is None else past_key_values[0][0].shape[2] + positions = torch.arange(past_len, past_len + l, device=input_ids.device) + POSITION_OFFSET + x = self.embed_tokens(input_ids) + self.embed_positions(positions)[None] + x = self.layernorm_embedding(x) + mask = None + if l > 1: + causal = torch.ones(l, past_len + l, dtype=torch.bool, device=x.device).tril(past_len) + mask = causal[None, None] + if attention_mask is not None: + mask = mask & attention_mask.bool()[:, None, None, :] + elif attention_mask is not None and past_len == 0: + mask = attention_mask.bool()[:, None, None, :] + if mask is not None: + mask = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, float("-inf")) + hidden_states = [x] if output_hidden_states else None + new_cache = [] + for i, layer in enumerate(self.layers): + x, kv = layer(x, encoder_hidden_states, mask, None if past_key_values is None else past_key_values[i]) + new_cache.append(kv) + if hidden_states is not None: + hidden_states.append(x) + return SimpleNamespace( + last_hidden_state=x, + past_key_values=tuple(new_cache) if use_cache else None, + hidden_states=tuple(hidden_states) if hidden_states is not None else None, + ) diff --git a/extensions_built_in/audio_models/yue2/src/shims/__init__.py b/extensions_built_in/audio_models/yue2/src/shims/__init__.py new file mode 100644 index 0000000..7f29dc2 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/shims/__init__.py @@ -0,0 +1,27 @@ +"""Stand-ins for mir_eval.chord and pretty_midi so m-a-p/SheetSage2's remote code imports without +those packages. Installed into sys.modules only when the real packages are absent.""" + +import importlib +import sys +import types + + +def install_shims(): + for name in ("mir_eval", "pretty_midi"): + try: + importlib.import_module(name) + continue # real package available; leave it alone + except ImportError: + pass + if name == "mir_eval": + from . import mir_eval_chord + + pkg = types.ModuleType("mir_eval") + pkg.chord = mir_eval_chord + pkg.__path__ = [] + sys.modules["mir_eval"] = pkg + sys.modules["mir_eval.chord"] = mir_eval_chord + else: + from . import pretty_midi + + sys.modules["pretty_midi"] = pretty_midi diff --git a/extensions_built_in/audio_models/yue2/src/shims/mir_eval_chord.py b/extensions_built_in/audio_models/yue2/src/shims/mir_eval_chord.py new file mode 100644 index 0000000..e37eeb9 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/shims/mir_eval_chord.py @@ -0,0 +1,518 @@ +"""Chord-label helpers vendored from mir_eval (mir_eval/chord.py), so the SheetSage2 remote code +runs without the mir_eval package. Unmodified except for the removed module preamble. + +The MIT License (MIT) + +Copyright (c) 2014 Colin Raffel + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +import re +import warnings + +import numpy as np + + +BITMAP_LENGTH = 12 +NO_CHORD = "N" +NO_CHORD_ENCODED = -1, np.array([0] * BITMAP_LENGTH), -1 +X_CHORD = "X" +X_CHORD_ENCODED = -1, np.array([-1] * BITMAP_LENGTH), -1 + + +class InvalidChordException(Exception): + r"""Exception class for suspect / invalid chord labels""" + + def __init__(self, message="", chord_label=None): + self.message = message + self.chord_label = chord_label + self.name = self.__class__.__name__ + super().__init__(message) + + +# --- Chord Primitives --- +def _pitch_classes(): + r"""Map from pitch class (str) to semitone (int).""" + pitch_classes = ["C", "D", "E", "F", "G", "A", "B"] + semitones = [0, 2, 4, 5, 7, 9, 11] + return {c: s for c, s in zip(pitch_classes, semitones)} + + +def _scale_degrees(): + r"""Map scale degrees (str) to semitones (int).""" + degrees = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13"] + semitones = [0, 2, 4, 5, 7, 9, 11, 12, 14, 16, 17, 19, 21] + return {d: s for d, s in zip(degrees, semitones)} + + +# Maps pitch classes (strings) to semitone indexes (ints). +PITCH_CLASSES = _pitch_classes() + + +def pitch_class_to_semitone(pitch_class): + r"""Convert a pitch class to semitone. + + Parameters + ---------- + pitch_class : str + Spelling of a given pitch class, e.g. 'C#', 'Gbb' + + Returns + ------- + semitone : int + Semitone value of the pitch class. + + """ + semitone = 0 + for idx, char in enumerate(pitch_class): + if char == "#" and idx > 0: + semitone += 1 + elif char == "b" and idx > 0: + semitone -= 1 + elif idx == 0: + semitone = PITCH_CLASSES.get(char) + else: + raise InvalidChordException( + "Pitch class improperly formed: %s" % pitch_class + ) + return semitone % 12 + + +# Maps scale degrees (strings) to semitone indexes (ints). +SCALE_DEGREES = _scale_degrees() + + +def scale_degree_to_semitone(scale_degree): + r"""Convert a scale degree to semitone. + + Parameters + ---------- + scale_degree : str + Spelling of a relative scale degree, e.g. 'b3', '7', '#5' + + Returns + ------- + semitone : int + Relative semitone of the scale degree, wrapped to a single octave + + Raises + ------ + InvalidChordException if `scale_degree` is invalid. + """ + semitone = 0 + offset = 0 + if scale_degree.startswith("#"): + offset = scale_degree.count("#") + scale_degree = scale_degree.strip("#") + elif scale_degree.startswith("b"): + offset = -1 * scale_degree.count("b") + scale_degree = scale_degree.strip("b") + + semitone = SCALE_DEGREES.get(scale_degree, None) + if semitone is None: + raise InvalidChordException( + "Scale degree improperly formed: {}, expected one of {}.".format( + scale_degree, list(SCALE_DEGREES.keys()) + ) + ) + return semitone + offset + + +def scale_degree_to_bitmap(scale_degree, modulo=False, length=BITMAP_LENGTH): + """Create a bitmap representation of a scale degree. + + Note that values in the bitmap may be negative, indicating that the + semitone is to be removed. + + Parameters + ---------- + scale_degree : str + Spelling of a relative scale degree, e.g. 'b3', '7', '#5' + modulo : bool, default=True + If a scale degree exceeds the length of the bit-vector, modulo the + scale degree back into the bit-vector; otherwise it is discarded. + length : int, default=12 + Length of the bit-vector to produce + + Returns + ------- + bitmap : np.ndarray, in [-1, 0, 1], len=`length` + Bitmap representation of this scale degree. + """ + sign = 1 + if scale_degree.startswith("*"): + sign = -1 + scale_degree = scale_degree.strip("*") + edit_map = [0] * length + sd_idx = scale_degree_to_semitone(scale_degree) + if sd_idx < length or modulo: + edit_map[sd_idx % length] = sign + return np.array(edit_map) + + +# Maps quality strings to bitmaps, corresponding to relative pitch class +# semitones, i.e. vector[0] is the tonic. +QUALITIES = { + # 1 2 3 4 5 6 7 + "maj": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0], + "min": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0], + "aug": [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0], + "dim": [1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0], + "sus4": [1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0], + "sus2": [1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], + "7": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0], + "maj7": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1], + "min7": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0], + "minmaj7": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1], + "maj6": [1, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0], + "min6": [1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0], + "dim7": [1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0], + "hdim7": [1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0], + "aug7": [1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0], + "maj9": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1], + "min9": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0], + "9": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0], + "min11": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0], + "11": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0], + "maj13": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1], + "min13": [1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0], + "13": [1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0], + "1": [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "5": [1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0], + "": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], +} + + +def quality_to_bitmap(quality): + """Return the bitmap for a given quality. + + Parameters + ---------- + quality : str + Chord quality name. + + Returns + ------- + bitmap : np.ndarray + Bitmap representation of this quality (12-dim). + + """ + if quality not in QUALITIES: + raise InvalidChordException( + "Unsupported chord quality shorthand: '%s' " + "Did you mean to reduce extended chords?" % quality + ) + return np.array(QUALITIES[quality]) + + +# Maps extended chord qualities to the subset above, translating additional +# voicings to extensions as a set of scale degrees (strings). +# TODO(ejhumphrey): Revisit how minmaj7's are mapped. This is how TMC did it, +# but MMV handles it like a separate quality (rather than an add7). +EXTENDED_QUALITY_REDUX = { + "minmaj7": ("min", {"7"}), + "maj9": ("maj7", {"9"}), + "min9": ("min7", {"9"}), + "9": ("7", {"9"}), + "b9": ("7", {"b9"}), + "#9": ("7", {"#9"}), + "11": ("7", {"9", "11"}), + "#11": ("7", {"9", "#11"}), + "13": ("7", {"9", "11", "13"}), + "b13": ("7", {"9", "11", "b13"}), + "min11": ("min7", {"9", "11"}), + "maj13": ("maj7", {"9", "11", "13"}), + "min13": ("min7", {"9", "11", "13"}), +} + + +def reduce_extended_quality(quality): + """Map an extended chord quality to a simpler one, moving upper voices to + a set of scale degree extensions. + + Parameters + ---------- + quality : str + Extended chord quality to reduce. + + Returns + ------- + base_quality : str + New chord quality. + extensions : set + Scale degrees extensions for the quality. + + """ + return EXTENDED_QUALITY_REDUX.get(quality, (quality, set())) + + +# --- Chord Label Parsing --- +# This monster regexp is pulled from the JAMS chord namespace, +# which is in turn derived from the context-free grammar of +# Harte et al., 2005. +CHORD_RE = re.compile( + r"""^((N|X)|(([A-G](b*|#*))((:(maj|min|dim|aug|1|5|sus2|sus4|maj6|min6|7|maj7|min7|dim7|hdim7|minmaj7|aug7|9|maj9|min9|11|maj11|min11|13|maj13|min13)(\((\*?((b*|#*)([1-9]|1[0-3]?))(,\*?((b*|#*)([1-9]|1[0-3]?)))*)\))?)|(:\((\*?((b*|#*)([1-9]|1[0-3]?))(,\*?((b*|#*)([1-9]|1[0-3]?)))*)\)))?((/((b*|#*)([1-9]|1[0-3]?)))?)?))$""" +) # nopep8 + + +def validate_chord_label(chord_label): + """Test for well-formedness of a chord label. + + Parameters + ---------- + chord_label : str + Chord label to validate. + """ + if not CHORD_RE.match(chord_label): + raise InvalidChordException("Invalid chord label: " "{}".format(chord_label)) + pass + + +def split(chord_label, reduce_extended_chords=False): + """Parse a chord label into its four constituent parts: + - root + - quality shorthand + - scale degrees + - bass + + Note: Chords lacking quality AND interval information are major. + - If a quality is specified, it is returned. + - If an interval is specified WITHOUT a quality, the quality field is + empty. + + Some examples:: + + 'C' -> ['C', 'maj', {}, '1'] + 'G#:min(*b3,*5)/5' -> ['G#', 'min', {'*b3', '*5'}, '5'] + 'A:(3)/6' -> ['A', '', {'3'}, '6'] + + Parameters + ---------- + chord_label : str + A chord label. + reduce_extended_chords : bool + Whether to map the upper voicings of extended chords (9's, 11's, 13's) + to semitone extensions. (Default value = False) + + Returns + ------- + chord_parts : list + Split version of the chord label. + + """ + chord_label = str(chord_label) + validate_chord_label(chord_label) + if chord_label == NO_CHORD: + return [chord_label, "", set(), ""] + + bass = "1" + if "/" in chord_label: + chord_label, bass = chord_label.split("/") + + scale_degrees = set() + omission = False + if "(" in chord_label: + chord_label, scale_degrees = chord_label.split("(") + omission = "*" in scale_degrees + scale_degrees = scale_degrees.strip(")") + scale_degrees = {i.strip() for i in scale_degrees.split(",")} + + # Note: Chords lacking quality AND added interval information are major. + # If a quality shorthand is specified, it is returned. + # If an interval is specified WITHOUT a quality, the quality field is + # empty. + # Intervals specifying omissions MUST have a quality. + if omission and ":" not in chord_label: + raise InvalidChordException( + "Intervals specifying omissions MUST have a quality." + ) + quality = "" if scale_degrees else "maj" + if ":" in chord_label: + chord_root, quality_name = chord_label.split(":") + # Extended chords (with ":"s) may not explicitly have Major qualities, + # so only overwrite the default if the string is not empty. + if quality_name: + quality = quality_name.lower() + else: + chord_root = chord_label + + if reduce_extended_chords: + quality, addl_scale_degrees = reduce_extended_quality(quality) + scale_degrees.update(addl_scale_degrees) + + return [chord_root, quality, scale_degrees, bass] + + +def join(chord_root, quality="", extensions=None, bass=""): + r"""Join the parts of a chord into a complete chord label. + + Parameters + ---------- + chord_root : str + Root pitch class of the chord, e.g. 'C', 'Eb' + quality : str + Quality of the chord, e.g. 'maj', 'hdim7' + (Default value = '') + extensions : list + Any added or absent scaled degrees for this chord, e.g. ['4', '\*3'] + (Default value = None) + bass : str + Scale degree of the bass note, e.g. '5'. + (Default value = '') + + Returns + ------- + chord_label : str + A complete chord label. + + """ + chord_label = chord_root + if quality or extensions: + chord_label += ":%s" % quality + if extensions: + chord_label += "(%s)" % ",".join(extensions) + if bass and bass != "1": + chord_label += "/%s" % bass + validate_chord_label(chord_label) + return chord_label + + +# --- Chords to Numerical Representations --- +def encode(chord_label, reduce_extended_chords=False, strict_bass_intervals=False): + """Translate a chord label to numerical representations for evaluation. + + Parameters + ---------- + chord_label : str + Chord label to encode. + reduce_extended_chords : bool + Whether to map the upper voicings of extended chords (9's, 11's, 13's) + to semitone extensions. + (Default value = False) + strict_bass_intervals : bool + Whether to require that the bass scale degree is present in the chord. + (Default value = False) + + Returns + ------- + root_number : int + Absolute semitone of the chord's root. + semitone_bitmap : np.ndarray, dtype=int + 12-dim vector of relative semitones in the chord spelling. + bass_number : int + Relative semitone of the chord's bass note, e.g. 0=root, 7=fifth, etc. + """ + if chord_label == NO_CHORD: + return NO_CHORD_ENCODED + if chord_label == X_CHORD: + return X_CHORD_ENCODED + chord_root, quality, scale_degrees, bass = split( + chord_label, reduce_extended_chords=reduce_extended_chords + ) + + root_number = pitch_class_to_semitone(chord_root) + bass_number = scale_degree_to_semitone(bass) % 12 + + semitone_bitmap = quality_to_bitmap(quality) + semitone_bitmap[0] = 1 + + for scale_degree in scale_degrees: + semitone_bitmap += scale_degree_to_bitmap(scale_degree, reduce_extended_chords) + + semitone_bitmap = (semitone_bitmap > 0).astype(np.int64) + if not semitone_bitmap[bass_number] and strict_bass_intervals: + raise InvalidChordException( + "Given bass scale degree is absent from this chord: " "%s" % chord_label, + chord_label, + ) + else: + semitone_bitmap[bass_number] = 1 + return root_number, semitone_bitmap, bass_number + + +def encode_many(chord_labels, reduce_extended_chords=False): + """Translate a set of chord labels to numerical representations for sane + evaluation. + + Parameters + ---------- + chord_labels : list + Set of chord labels to encode. + reduce_extended_chords : bool + Whether to map the upper voicings of extended chords (9's, 11's, 13's) + to semitone extensions. + (Default value = False) + + Returns + ------- + root_number : np.ndarray, dtype=int + Absolute semitone of the chord's root. + interval_bitmap : np.ndarray, dtype=int + 12-dim vector of relative semitones in the given chord quality. + bass_number : np.ndarray, dtype=int + Relative semitones of the chord's bass notes. + + """ + num_items = len(chord_labels) + roots, basses = np.zeros([2, num_items], dtype=np.int64) + semitones = np.zeros([num_items, 12], dtype=np.int64) + local_cache = dict() + for i, label in enumerate(chord_labels): + result = local_cache.get(label, None) + if result is None: + result = encode(label, reduce_extended_chords) + local_cache[label] = result + roots[i], semitones[i], basses[i] = result + return roots, semitones, basses + + +def rotate_bitmap_to_root(bitmap, chord_root): + """Circularly shift a relative bitmap to its absolute pitch classes. + + For clarity, the best explanation is an example. Given 'G:Maj', the root + and quality map are as follows:: + + root=5 + quality=[1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0] # Relative chord shape + + After rotating to the root, the resulting bitmap becomes:: + + abs_quality = [0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1] # G, B, and D + + Parameters + ---------- + bitmap : np.ndarray, shape=(12,) + Bitmap of active notes, relative to the given root. + chord_root : int + Absolute pitch class number. + + Returns + ------- + bitmap : np.ndarray, shape=(12,) + Absolute bitmap of active pitch classes. + + """ + bitmap = np.asarray(bitmap) + assert bitmap.ndim == 1, "Currently only 1D bitmaps are supported." + idxs = list(np.nonzero(bitmap)) + idxs[-1] = (idxs[-1] + chord_root) % 12 + abs_bitmap = np.zeros_like(bitmap) + abs_bitmap[tuple(idxs)] = 1 + return abs_bitmap diff --git a/extensions_built_in/audio_models/yue2/src/shims/pretty_midi.py b/extensions_built_in/audio_models/yue2/src/shims/pretty_midi.py new file mode 100644 index 0000000..7ee2eb4 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/shims/pretty_midi.py @@ -0,0 +1,89 @@ +"""Minimal in-memory stand-in for the parts of pretty_midi that SheetSage2's ABC path uses: +PrettyMIDI(resolution=...) / PrettyMIDI(bytes-like), .instruments, .resolution, .write(stream); +Instrument(program, name=...) with .notes/.name/.program/.is_drum; Note(velocity, pitch, start, end). + +"Bytes" are our own serialization (the only consumer is this shim), but note times are rounded to +MIDI ticks on write exactly as a real file round trip would (120 bpm default tempo: ticks per second += 2 * resolution), because the ABC builder is documented to follow the serialized MIDI timing.""" + +import json +from io import BytesIO + +_MAGIC = b"AITK-MIDI-JSON\n" + + +class Note: + def __init__(self, velocity, pitch, start, end): + self.velocity = int(velocity) + self.pitch = int(pitch) + self.start = float(start) + self.end = float(end) + + def __repr__(self): + return f"Note(start={self.start:.6f}, end={self.end:.6f}, pitch={self.pitch}, velocity={self.velocity})" + + +class Instrument: + def __init__(self, program, is_drum=False, name=""): + self.program = int(program) + self.is_drum = bool(is_drum) + self.name = name + self.notes = [] + + +class PrettyMIDI: + def __init__(self, midi_file=None, resolution=220, initial_tempo=120.0): + self.resolution = int(resolution) + self.initial_tempo = float(initial_tempo) + self.instruments = [] + if midi_file is not None: + if isinstance(midi_file, (bytes, bytearray)): + data = bytes(midi_file) + elif hasattr(midi_file, "read"): + data = midi_file.read() + else: + with open(str(midi_file), "rb") as f: + data = f.read() + if not data.startswith(_MAGIC): + raise ValueError("pretty_midi shim can only read MIDI written by this shim") + payload = json.loads(data[len(_MAGIC) :].decode("utf-8")) + self.resolution = int(payload["resolution"]) + self.initial_tempo = float(payload["initial_tempo"]) + for inst in payload["instruments"]: + instrument = Instrument(inst["program"], inst["is_drum"], inst["name"]) + instrument.notes = [Note(n["velocity"], n["pitch"], n["start"], n["end"]) for n in inst["notes"]] + self.instruments.append(instrument) + + def _ticks_per_second(self): + return self.resolution * self.initial_tempo / 60.0 + + def _quantize(self, seconds): + tps = self._ticks_per_second() + return int(round(seconds * tps)) / tps + + def write(self, filename): + payload = { + "resolution": self.resolution, + "initial_tempo": self.initial_tempo, + "instruments": [ + { + "program": i.program, + "is_drum": i.is_drum, + "name": i.name, + "notes": [ + {"velocity": n.velocity, "pitch": n.pitch, "start": self._quantize(n.start), "end": self._quantize(n.end)} + for n in i.notes + ], + } + for i in self.instruments + ], + } + data = _MAGIC + json.dumps(payload).encode("utf-8") + if hasattr(filename, "write"): + filename.write(data) + else: + with open(str(filename), "wb") as f: + f.write(data) + + def get_end_time(self): + return max((n.end for i in self.instruments for n in i.notes), default=0.0) diff --git a/extensions_built_in/audio_models/yue2/src/tokenizer.py b/extensions_built_in/audio_models/yue2/src/tokenizer.py new file mode 100644 index 0000000..392aab3 --- /dev/null +++ b/extensions_built_in/audio_models/yue2/src/tokenizer.py @@ -0,0 +1,230 @@ +"""Text side and audio side of YuE2 conditioning. + +- ``YuE2TextTokenizer``: the Qwen BPE that ComfyUI embeds in the checkpoint + (``text_encoders.yue2_tokenizer_json``), plus the prompt/prefix protocol. +- ``SemanticTokenizer``: the community audio -> codec-token head + (Mothersuperior v4): MERT-v2-FullSong layer-20 features at 25 Hz, + per-track instance normalised, through an 8-layer transformer classifier + over 512-frame windows. +""" + +import unicodedata +from typing import List, Optional + +import torch +import torch.nn.functional as F +from torch import nn + +from .model import ABC_END, ABC_START, CODEC_SIZE, EOD, INSTRUCTIONS, MUSIC_START + +MERT_REPO = "m-a-p/MERT-v2-FullSong" +MERT_SAMPLE_RATE = 24000 +MERT_LAYER = 20 +HEAD_REPO = "Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4" +HEAD_FILE = "tokenizer_head_joint_v4.pt" +NAR_LORA_FILE = "nar_lora_joint_v4.pt" +HEAD_WINDOW = 512 + + +class YuE2TextTokenizer: + def __init__(self, tokenizer_json: bytes): + from tokenizers import Tokenizer + + self.tokenizer = Tokenizer.from_str(tokenizer_json.decode("utf-8")) + + def encode(self, text: str) -> List[int]: + return self.tokenizer.encode(unicodedata.normalize("NFC", text)).ids + + def decode(self, ids: List[int]) -> str: + return self.tokenizer.decode(ids, skip_special_tokens=True) + + def prefix_head_ids(self, style: str, lyrics: str, cot: str = "off") -> List[int]: + """Prompt up to and including ABC_START (ComfyUI's ``prefix``); the ABC block and + ``[ABC_END, MUSIC_START]`` follow (see ``abc_tail_ids``).""" + prompt = f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n" + return [EOD] + self.encode(prompt) + [ABC_START] + + @staticmethod + def abc_tail_ids(abc_ids: List[int]) -> List[int]: + return list(abc_ids) + [ABC_END, MUSIC_START] + + def encode_abc(self, abc: str) -> List[int]: + ids = self.encode(abc) + if any(not 0 <= t < EOD for t in ids): + raise ValueError("ABC ids must stay inside the ordinary text vocabulary") + return ids + + def prefix_ids(self, style: str, lyrics: str, cot: str = "off", abc: str = "") -> List[int]: + """Everything the AR expert sees before the first codec token.""" + abc_ids = [] if cot == "off" else self.encode_abc(abc) + return self.prefix_head_ids(style, lyrics, cot) + self.abc_tail_ids(abc_ids) + + +class TokenizerHead(nn.Module): + def __init__(self, din=1024, d=512, layers=8, heads=8, vocab=CODEC_SIZE, window=HEAD_WINDOW): + super().__init__() + self.inp = nn.Linear(din, d) + self.pos = nn.Parameter(torch.zeros(1, window, d)) + layer = nn.TransformerEncoderLayer(d, heads, 4 * d, dropout=0.1, batch_first=True, norm_first=True, activation="gelu") + self.enc = nn.TransformerEncoder(layer, layers) + self.norm = nn.LayerNorm(d) + self.head = nn.Linear(d, vocab) + + def forward(self, x): + return self.head(self.norm(self.enc(self.inp(x) + self.pos[:, : x.shape[1]]))) + + +def _rebuild_rotary(model: nn.Module) -> int: + """MERT2's RotaryEmbedding keeps ``inv_freq`` as a non-persistent buffer; transformers' + meta-device loading materializes it uninitialized (zeros on a fresh GPU, garbage or NaN + otherwise). Recompute it from the module's own formula and drop the cos/sin cache.""" + n = 0 + for m in model.modules(): + if hasattr(m, "inv_freq") and hasattr(m, "head_dim") and hasattr(m, "base"): + inv = 1.0 / (m.base ** (torch.arange(0, m.head_dim, 2, dtype=torch.float32) / m.head_dim)) + m.inv_freq = inv.to(device=m.inv_freq.device) + for attr, val in (("_cos", None), ("_sin", None), ("_sequence_length", 0), ("_cache_device", None)): + if hasattr(m, attr): + setattr(m, attr, val) + n += 1 + return n + + +SHEETSAGE_REPO = "m-a-p/SheetSage2" +SHEETSAGE_FILE = "Comfy-Org/YuE2/audio_encoders/sheetsage2_bf16.safetensors" +# tied to token_embedding.weight (loaded); everything else in the Comfy file maps 1:1 onto upstream +NONPERSISTENT_OK = {"decoder.embed_tokens.weight"} +SHEETSAGE_SAMPLE_RATE = 24000 + + +class SheetSage2Transcriber: + """Audio -> ABC lead sheet with the released SheetSage2 (transformers remote code), the same + transcriber ComfyUI ships for cover mode. ``full`` keeps chords, ``melody`` drops them.""" + + def __init__(self, weights_path: str, repo: str = SHEETSAGE_REPO): + """``weights_path``: the Comfy repack ``audio_encoders/sheetsage2_bf16.safetensors`` (merged + weights); ``repo`` supplies config + code (transformers remote code) only.""" + from .shims import install_shims + + install_shims() # mir_eval.chord / pretty_midi stand-ins unless the real packages exist + import sys + + from safetensors.torch import load_file + from transformers import AutoConfig + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + from .sheetsage_decoder import ScoreDecoder + + cfg = AutoConfig.from_pretrained(repo, trust_remote_code=True) + cfg.weights_format = "merged" # the Comfy file carries the LoRA already folded into the encoder + cls = get_class_from_dynamic_module(cfg.auto_map["AutoModel"], repo) + sys.modules[cls.__module__].BartDecoder = ScoreDecoder # upstream targets transformers 4.45's BartDecoder + cls.post_init = lambda self: None # transformers-5 tied-weight bookkeeping rejects upstream's list; weights load explicitly + model = cls(cfg) + sd = {k: v.to(torch.float32) for k, v in load_file(weights_path).items()} + missing, unexpected = model.load_state_dict(sd, strict=False) + missing = [k for k in missing if k not in NONPERSISTENT_OK] + if missing or unexpected: + raise ValueError(f"SheetSage2 weights do not match: missing {missing[:5]} unexpected {unexpected[:5]}") + model.output_projection.weight = model.token_embedding.weight # keep the tie after loading + self.model = model.eval() + self.model.requires_grad_(False) + _rebuild_rotary(self.model) + + def to(self, device): + self.model.to(device) + return self + + @property + def device(self): + return next(self.model.parameters()).device + + @torch.no_grad() + def transcribe(self, waveform: torch.Tensor, sample_rate: int, cot: str = "full") -> str: + """waveform [C, samples] -> ABC text (``cot`` "full" or "melody").""" + mono = waveform.float().mean(0) if waveform.dim() == 2 else waveform.float() + result = self.model.transcribe(mono.cpu(), sampling_rate=sample_rate, melody_only=cot == "melody") + abc = result.get("abc") or "" + if not abc.strip(): + raise RuntimeError(f"SheetSage2 produced no ABC: {result.get('abc_error') or 'unknown reason'}") + return abc + + +class SemanticTokenizer(nn.Module): + """Waveform -> per-frame codec ids (0..32767) at 25 Hz.""" + + def __init__(self, head_path: str, mert_repo: str = MERT_REPO, debug: bool = False): + super().__init__() + self.debug = debug + from transformers import AutoFeatureExtractor, AutoModel + + self.processor = AutoFeatureExtractor.from_pretrained(mert_repo, trust_remote_code=True) + self.mert = AutoModel.from_pretrained(mert_repo, trust_remote_code=True).eval() + self.mert.requires_grad_(False) + _rebuild_rotary(self.mert) + self.head = TokenizerHead() + ck = torch.load(head_path, map_location="cpu", weights_only=False) + self.head.load_state_dict(ck["model"]) + self.head.eval().requires_grad_(False) + + @property + def device(self): + return self.head.head.weight.device + + @torch.no_grad() + def mert_features(self, mono24: torch.Tensor) -> torch.Tensor: + """mono24 [samples] at 24 kHz -> [T25, 1024] fp32 layer-20 features.""" + device = self.device + chunk = MERT_SAMPLE_RATE * 30 + chunks = [mono24[s : s + chunk] for s in range(0, mono24.shape[0], chunk)] + chunks = [c for c in chunks if c.shape[0] >= MERT_SAMPLE_RATE] + full = [c for c in chunks if c.shape[0] == chunk] + tail = [c for c in chunks if c.shape[0] < chunk] + feats = [] + with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + for group in ([full] if full else []) + [[c] for c in tail]: + inp = self.processor([c.cpu().numpy() for c in group], sampling_rate=MERT_SAMPLE_RATE, return_tensors="pt") + inp = {k: v.to(device) for k, v in inp.items()} + out = self.mert(**inp, output_hidden_states=True) + feats.append(out.hidden_states[MERT_LAYER].reshape(-1, 1024)) + h = torch.cat(feats, 0).float() + t25 = int(round(mono24.shape[0] / MERT_SAMPLE_RATE * 25)) + return F.interpolate(h.T[None], size=t25, mode="linear", align_corners=False)[0].T + + @torch.no_grad() + def tokens_from_features(self, feats: torch.Tensor) -> torch.Tensor: + """[T, 1024] -> [T] long. Windowed inference with edge trimming.""" + x = feats.float() + x = (x - x.mean(0)) / (x.std(0) + 1e-5) + total = x.shape[0] + win = HEAD_WINDOW + out = torch.zeros(total, dtype=torch.long, device=x.device) + starts = list(range(0, max(1, total - win + 1), win // 2)) + if starts[-1] + win < total: + starts.append(max(0, total - win)) + device = self.device + for s0 in starts: + xw = x[s0 : s0 + win] + n = xw.shape[0] + if n < win: + xw = F.pad(xw, (0, 0, 0, win - n)) + with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + pred = self.head(xw[None].to(device))[0, :n].float().argmax(-1) + lo = s0 + (0 if s0 == 0 else win // 4) + hi = s0 + n - (0 if s0 + n >= total else win // 4) + out[lo:hi] = pred[lo - s0 : hi - s0].to(out.device) + return out + + @torch.no_grad() + def tokenize(self, waveform: torch.Tensor, sample_rate: int) -> torch.Tensor: + """waveform [C, samples] -> codec ids [T] (long, no vocab offset).""" + import torchaudio + + mono = waveform.float().mean(0) if waveform.dim() == 2 else waveform.float() + if sample_rate != MERT_SAMPLE_RATE: + mono = torchaudio.functional.resample(mono, sample_rate, MERT_SAMPLE_RATE) + feats = self.mert_features(mono.to(self.device)) + out = self.tokens_from_features(feats) + if self.debug and out.unique().numel() < 8: + print(f"YuE2 tokenizer: degenerate token stream (unique {out.unique().numel()}); MERT features nan {torch.isnan(feats).sum().item()}") + return out 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..bdad93b --- /dev/null +++ b/extensions_built_in/audio_models/yue2/yue2_model.py @@ -0,0 +1,676 @@ +"""YuE2 (m-a-p/YuE2-3B) for ai-toolkit. + +Weights: the Comfy-Org all-in-one checkpoint (default ``checkpoints/yue2_3b_int8_convrot.safetensors``, +shipped int8 convrot layers attached as-is; ``yue2_3b_bf16.safetensors`` also loads) +resolved under MODELS_PATH in ComfyUI's folder layout and downloaded there when missing. +The official audio -> semantic-token tokenizer is unreleased; training conditioning comes +from the community head (MERT-v2-FullSong + classifier) in +``Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4``. Its companion NAR adapter can +be merged into the base NAR on load (``model_kwargs.merge_nar_lora``, default false). + +Training: per song the latent cache carries VAE latents plus the head's codec tokens; the +"text embedding" is the AR prompt prefix (instruction + tags + lyrics) as token embeddings. +Each step crops one window (``model_kwargs.train_window_frames``, 25 fps, 0 = whole song), +prefills the AR expert over ``prefix + tokens[window] + MUSIC_END`` and trains the NAR +flow loss on the real latents. With ``model_kwargs.ar_loss_weight`` > 0 the same prefill +also trains the AR expert with next-token cross entropy over the codec tokens, so one LoRA +network on both experts learns composition and rendering together. + +Prompt modes follow ComfyUI (``model_kwargs.cot``): "full" (default, Generate ABC path: the AR +writes a chord-annotated ABC sheet, then the codec tokens), "melody" (sheet without chords) or +"off" (no sheet; what Comfy runs when the ABC input is empty). For full/melody the training sheet +comes from SheetSage2 at cache time and rides the latent cache as ``abc_ids``. +""" + +import os +import random +import re +from typing import List, Optional + +import torch +import torch.utils.checkpoint +from safetensors.torch import load_file +from tqdm import tqdm + +from extensions_built_in.audio_models.base_audio_model import BaseAudioModel +from toolkit.basic import flush +from toolkit.config_modules import GenerateImageConfig +from toolkit.dto import DTO +from toolkit.models.v2.resolver import resolve_named_file +from toolkit.paths import MODELS_PATH +from toolkit.print import print_acc +from toolkit.prompt_utils import PromptEmbeds +from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler + +from .src.model import ABC_END, CODEC_OFFSET, CONTEXT, FRAMES_PER_SECOND, INSTRUCTIONS, MUSIC_END, MUSIC_START, YuE2Model, merge_nar_lora +from .src.pipeline import YuE2Pipeline +from .src.tokenizer import HEAD_FILE, HEAD_REPO, NAR_LORA_FILE, SHEETSAGE_FILE, SemanticTokenizer, SheetSage2Transcriber, YuE2TextTokenizer +from .src.vae import SAMPLE_RATE, YuE2VAE + +COT_CODES = {"off": 0, "melody": 1, "full": 2} +COT_NAMES = {v: k for k, v in COT_CODES.items()} + +DEFAULT_CHECKPOINT = "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors" +# community tokenizer files have no ComfyUI folder; they live under /ai_toolkit/ +AI_TOOLKIT_DIR = "ai_toolkit" + + +def resolve_ai_toolkit_file(path: str, component: str) -> str: + """Local file, or 'org/repo/file' downloaded into /ai_toolkit/.""" + if os.path.exists(path): + return path + splits = path.split("/") + if len(splits) < 3: + raise ValueError(f"Invalid {component} path: {path}. Must be a local file or 'org/repo/filename'.") + local_dir = os.path.join(MODELS_PATH, AI_TOOLKIT_DIR) + candidate = os.path.join(local_dir, splits[-1]) + if os.path.exists(candidate): + return candidate + import huggingface_hub + + return huggingface_hub.hf_hub_download(repo_id="/".join(splits[:2]), filename="/".join(splits[2:]), local_dir=local_dir) + +scheduler_config = { + "num_train_timesteps": 1000, + "shift": 1.0, + "use_dynamic_shifting": False, +} + + +def _tag(text: str, name: str) -> str: + m = re.search(rf"<{name}>(.*?)", text, re.DOTALL) + return m.group(1).strip() if m else "" + + +_SECTION = re.compile(r"^\s*\[(Tags|Lyrics|Duration)\]\s*$", re.IGNORECASE | re.MULTILINE) +_SONG_SECTION = re.compile(r"^\s*\[[^\]]+\]\s*$") + + +def _normalize_section(line: str) -> str: + """YuE2 section tags are Title case ([Verse 1], [Pre-Chorus]); uppercase the first letter and the + letter after each hyphen of a bracketed section line, leaving the rest (e.g. descriptors) alone.""" + if not _SONG_SECTION.match(line): + return line + lead, body, trail = line[: line.index("[") + 1], line[line.index("[") + 1 : line.rindex("]")], line[line.rindex("]") :] + body = " ".join(w.capitalize() if w.isupper() and len(w) > 1 else w for w in body.split(" ")) # [BRIDGE] -> [Bridge] + body = body[:1].upper() + body[1:] + body = re.sub(r"-([a-z])", lambda m: "-" + m.group(1).upper(), body) + return lead + body + trail + + +def _number(value: str): + try: + return float(value) if value.strip() else None + except ValueError: + return None + + +def parse_caption(text: str) -> dict: + """Caption / prompt -> style, lyrics, duration. + + Native form (YuE2's own prompt layout): style text, then a ``[Lyrics]`` line and the + lyrics; an optional leading ``[Tags]`` line and an optional trailing ``[Duration]`` + line (max seconds, samples only). Plain text with no sections is the style. + Legacy ````/````/```` tags are still accepted.""" + if not isinstance(text, str): + text = "" + if "" in text or "" in text: + lyrics = "\n".join(_normalize_section(l) for l in _tag(text, "LYRICS").splitlines()) + return {"style": _tag(text, "CAPTION"), "lyrics": lyrics, "duration": _number(_tag(text, "DURATION"))} + sections = {"tags": []} + current = "tags" + for line in text.splitlines(): + m = _SECTION.match(line) + if m: + current = m.group(1).lower() + sections.setdefault(current, []) + continue + # no [Lyrics] header: the first song-section line ([Intro], [Verse 1], ...) starts the lyrics + if current == "tags" and "lyrics" not in sections and _SONG_SECTION.match(line): + current = "lyrics" + sections["lyrics"] = [] + sections.setdefault(current, []).append(_normalize_section(line) if current == "lyrics" else line) + style = "\n".join(sections["tags"]).strip() + lyrics = "\n".join(sections.get("lyrics", [])).strip() + duration = _number("\n".join(sections.get("duration", []))) + return {"style": style, "lyrics": lyrics, "duration": duration} + + +class YuE2TextEncoder(torch.nn.Module): + """Prompt -> AR token embeddings of the prefix head (instruction, tags, lyrics, ABC_START). + Holds no weights of its own; the embedding table lives in the AR expert.""" + + def __init__(self, model: YuE2Model, tokenizer: YuE2TextTokenizer, cot: str = "full"): + super().__init__() + self._model = [model] + self.tokenizer = tokenizer + self.cot = cot + self.register_buffer("_anchor", torch.zeros(1), persistent=False) + + @property + def device(self): + return self._anchor.device + + @property + def dtype(self): + return self._model[0].ar.dtype + + @torch.no_grad() + def forward(self, style: str, lyrics: str) -> torch.Tensor: + ids = torch.tensor([self.tokenizer.prefix_head_ids(style, lyrics, cot=self.cot)], dtype=torch.long) + return self._model[0].ar.embed(ids) + + +class YuE2AudioModel(BaseAudioModel): + arch = "yue2" + sample_rate = SAMPLE_RATE + + def __init__(self, device, model_config, dtype="bf16", custom_pipeline=None, noise_scheduler=None, **kwargs): + super().__init__(device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs) + self.is_flow_matching = True + self.is_transformer = True + self.target_lora_modules = ["YuE2AR", "YuE2NAR"] + kw = self.model_config.model_kwargs + self.train_window_frames = int(kw.get("train_window_frames", 1500)) + self.ar_loss_weight = float(kw.get("ar_loss_weight", 1.0)) + # AR next-token loss always runs on the song from its start (lyric order); 0 = whole song + self.ar_max_tokens = int(kw.get("ar_max_tokens", 0)) + # optional separate learning rate for the AR expert's LoRA (Adam ignores loss scale) + self.ar_lr_multiplier = float(kw.get("ar_lr_multiplier", 1.0)) + # trust region: KL(base || lora) on the AR next-token distributions, base = LoRA switched off + self.ar_kl_weight = float(kw.get("ar_kl_weight", 0.0)) + self._loss_log_every = int(kw.get("loss_log_every", 25)) + # extra per-readout AR diagnostics (targets, unique tokens, p(END), p(target)) + self.debug = bool(kw.get("debug", False)) + self._loss_log_step = 0 + self.merge_nar_lora = bool(kw.get("merge_nar_lora", False)) + self.nar_lora_path = kw.get("nar_lora_path", f"{HEAD_REPO}/{NAR_LORA_FILE}") + self.semantic_head_path = kw.get("semantic_head_path", f"{HEAD_REPO}/{HEAD_FILE}") + # ComfyUI's Generate ABC path: "full" (chords) or "melody"; "off" skips the sheet (empty ABC in Comfy) + self.cot = str(kw.get("cot", "full")) + if self.cot not in INSTRUCTIONS: + raise ValueError(f"model_kwargs.cot must be one of {list(INSTRUCTIONS)}") + self.abc_max_tokens = int(kw.get("abc_max_tokens", 8192)) + # fraction of training items fed without the sheet (off-mode prefix), so one LoRA serves both Comfy paths + self.abc_dropout = float(kw.get("abc_dropout", 0.5)) + self.sheetsage_path = kw.get("sheetsage_path", SHEETSAGE_FILE) + self.transcriber: Optional[SheetSage2Transcriber] = None + # encoders are only needed while caching; once training steps run without encode calls they are released + self._train_calls = 0 + self._last_encode_call = 0 + self.sample_max_seconds = float(kw.get("sample_max_seconds", 120.0)) + self.sample_ar_temperature = float(kw.get("sample_ar_temperature", 1.0)) + # reference sampler value; 1.0 lets a memorized song replay (the penalty knocks it off path) + self.sample_ar_repetition_penalty = float(kw.get("sample_ar_repetition_penalty", 1.2)) + self.semantic_tokenizer: Optional[SemanticTokenizer] = None + self._pending_aux_loss = None + self._pending_ar_ce = None + self._pending_ar_kl = None + self.additional_loss_logs = {} + + @staticmethod + def get_train_scheduler(): + return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config) + + # ------------------------------------------------------------------ + # loading + # ------------------------------------------------------------------ + def load_model(self): + dtype = self.torch_dtype + device = self.device_torch + name_or_path = self.model_config.name_or_path or DEFAULT_CHECKPOINT + ckpt_path = resolve_named_file(name_or_path, component="yue2 checkpoint") + self.print_and_status_update(f"Loading YuE2 from {ckpt_path}") + sd = load_file(ckpt_path) + tok_json = sd.pop("text_encoders.yue2_tokenizer_json", None) + if tok_json is None: + raise ValueError("YuE2 checkpoint has no embedded tokenizer (expected the Comfy-Org repack)") + self.tokenizer = YuE2TextTokenizer(tok_json.numpy().tobytes()) + + if self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0: + raise NotImplementedError("Layer offloading not yet implemented for YuE2") + + self.model = YuE2Model.load_from_state_dict(sd, dtype=dtype) + if self.merge_nar_lora: + lora_path = resolve_ai_toolkit_file(self.nar_lora_path, component="yue2 nar lora") + self.print_and_status_update(f"Merging NAR adapter {os.path.basename(lora_path)}") + merge_nar_lora(self.model, lora_path) + vae_sd = {k[len("vae.") :]: v for k, v in sd.items() if k.startswith("vae.")} + # reference runs the VAE in fp32; keep the trainer's dtype check from downcasting it + self.vae_torch_dtype = torch.float32 + self.vae = YuE2VAE.load_from_state_dict(vae_sd, dtype=torch.float32) + del sd + flush() + + self.model.aitk_post_load(**self.component_load_kwargs("transformer")) + flush() + self.model.to(device) + self.vae.to(self.vae_device_torch) + self.text_encoder = YuE2TextEncoder(self.model, self.tokenizer, cot=self.cot) + self.pipeline = YuE2Pipeline(self.model, self.vae) + self.pipeline.do_tiled_decoding = True + self.pipeline.use_cuda_graphs = bool(self.model_config.model_kwargs.get("ar_cuda_graphs", True)) + + def _get_semantic_tokenizer(self) -> SemanticTokenizer: + if self.semantic_tokenizer is None: + head_path = resolve_ai_toolkit_file(self.semantic_head_path, component="yue2 semantic head") + # log only: this can run mid-training and a status update would replace "Training" in the UI + print_acc("Loading YuE2 semantic tokenizer (MERT-v2-FullSong + head)") + self.semantic_tokenizer = SemanticTokenizer(head_path, debug=self.debug) + return self.semantic_tokenizer + + def _get_transcriber(self) -> SheetSage2Transcriber: + if self.transcriber is None: + print_acc("Loading SheetSage2 (audio -> ABC)") + self.transcriber = SheetSage2Transcriber(resolve_named_file(self.sheetsage_path, component="sheetsage2")) + return self.transcriber + + def get_transformer_block_names(self) -> Optional[List[str]]: + return ["model.layers"] + + def get_model_has_grad(self): + return False + + def get_te_has_grad(self): + return False + + # ------------------------------------------------------------------ + # conditioning + # ------------------------------------------------------------------ + def get_prompt_embeds(self, prompt) -> PromptEmbeds: + prompts = [prompt] if isinstance(prompt, str) else list(prompt) + embeds, masks = [], [] + for p in prompts: + parsed = parse_caption(p) + e = self.text_encoder(parsed["style"], parsed["lyrics"]) # [1, L, H] + embeds.append(e) + masks.append(torch.ones(1, e.shape[1], dtype=torch.long, device=e.device)) + max_len = max(e.shape[1] for e in embeds) + embeds = [torch.nn.functional.pad(e, (0, 0, 0, max_len - e.shape[1])) for e in embeds] + masks = [torch.nn.functional.pad(m, (0, max_len - m.shape[1])) for m in masks] + return PromptEmbeds(torch.cat(embeds, 0), attention_mask=torch.cat(masks, 0)) + + def encode_audio(self, audio_tensor: torch.Tensor, device=None, dtype=None): + """[B, 2, samples] at 48 kHz -> DTO [B, T, 64] with codec ``tokens`` [B, T] and, unless cot is + "off", the SheetSage2 sheet as ``abc_ids`` [B, N] (text ids, -1 padded).""" + if device is None: + device = self.vae_device_torch + self._last_encode_call = self._train_calls + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + latents = self.vae.encode(audio_tensor.to(device=device, dtype=self.vae.dtype)) # [B, 64, T] + latents = latents.transpose(1, 2).contiguous() + tokenizer = self._get_semantic_tokenizer() + if tokenizer.device != self.device_torch: + tokenizer.to(self.device_torch) + tokens = torch.stack([tokenizer.tokenize(wav, SAMPLE_RATE) for wav in audio_tensor]) # [B, T'] + n = min(latents.shape[1], tokens.shape[1]) + latents = latents[:, :n].to(dtype=self.torch_dtype if dtype is None else dtype) + extras = {"tokens": tokens[:, :n].to(latents.device, torch.int32)} + if self.cot != "off": + transcriber = self._get_transcriber() + if transcriber.device != self.device_torch: + transcriber.to(self.device_torch) + sheets = [self.tokenizer.encode_abc(transcriber.transcribe(wav, SAMPLE_RATE, cot=self.cot)) for wav in audio_tensor] + width = max(len(x) for x in sheets) + abc = torch.full((len(sheets), width), -1, dtype=torch.int32) + for i, ids in enumerate(sheets): + abc[i, : len(ids)] = torch.tensor(ids, dtype=torch.int32) + extras["abc_ids"] = abc.to(latents.device) + # which sheet flavor the cache holds; training refuses a mismatch instead of silently using it + extras["abc_mode"] = torch.full((len(sheets),), COT_CODES[self.cot], dtype=torch.int32, device=latents.device) + return DTO(latents, **extras) + + # ------------------------------------------------------------------ + # training + # ------------------------------------------------------------------ + def condition_noisy_latents(self, latents: torch.Tensor, batch): + frames = latents.shape[1] + window = self.train_window_frames + if window <= 0 or frames <= window: + batch.yue2_window = (0, frames) + return latents + start = random.randint(0, frames - window) + batch.yue2_window = (start, start + window) + return latents[:, start : start + window] + + def get_loss_target(self, *args, **kwargs): + noise = kwargs.get("noise") + batch = kwargs.get("batch") + start, end = getattr(batch, "yue2_window", (0, noise.shape[1])) + return (noise - batch.latents)[:, start:end].detach() + + def get_additional_loss(self, pred: torch.Tensor, target: torch.Tensor): + """AR next-token CE from the last forward, added to the flow loss by the trainer and + logged as its own term (`loss/ar_ce`, unweighted).""" + aux = self._pending_aux_loss + self._pending_aux_loss = None + if aux is None: + self.additional_loss_logs = {} + return None + ce = self._pending_ar_ce.detach().float().item() + self.additional_loss_logs = {"loss/ar_ce": ce} + kl = self._pending_ar_kl + if kl is not None: + self.additional_loss_logs["loss/ar_kl"] = kl.detach().float().item() + self._loss_log_step += 1 + if self.debug and self._loss_log_every > 0 and self._loss_log_step % self._loss_log_every == 0: + info = getattr(self, "_pending_aux_info", "") + kl_txt = f" ar_kl {kl.detach().float().item():.4f}" if kl is not None else "" + tqdm.write(f"yue2 step {self._loss_log_step}: ar_ce {ce:.4f}{kl_txt}" + (f" {info}" if info else "")) + return aux + + @staticmethod + def _train_prefix(prefix_embeds: torch.Tensor, prefix_mask: torch.Tensor) -> torch.Tensor: + """Unpadded prefix embeds for one item.""" + # the mask may arrive cast to bf16; count entries, never sum them + length = int((prefix_mask > 0.5).sum().item()) + return prefix_embeds[:length] + + def _ar_losses(self, ar_embeds: torch.Tensor, ids: torch.Tensor, prefix: torch.Tensor, total: int, whole_song: bool): + """Next-token CE over ``ids`` (and KL(base || lora) when ``ar_kl_weight`` > 0), computed in + position chunks with recomputed logits so the fp32 [n, vocab] tensors never exist at once. + Returns (ce, kl_or_None, info_str).""" + model = self.model + n = ids.shape[0] + _, hidden = model.ar.prefill(ar_embeds, return_hidden=True) + hidden = hidden[0, -n - 1 : -1] + base_hidden = None + if self.ar_kl_weight > 0: + net = getattr(self, "_network", None) + was_active = getattr(net, "is_active", None) + if net is not None: + net.is_active = False + try: + with torch.no_grad(): + _, base_hidden = model.ar.prefill(ar_embeds, return_hidden=True) + base_hidden = base_hidden[0, -n - 1 : -1] + finally: + if net is not None: + net.is_active = was_active + want_info = self.debug and self._loss_log_every > 0 and (self._loss_log_step + 1) % self._loss_log_every == 0 + lm_head = model.ar.model.lm_head + want_kl = base_hidden is not None + + def chunk_losses(h, tgt, bh): + # everything vocab-sized lives inside this checkpointed function, so autograd keeps only + # the [chunk, hidden] inputs and recomputes the fp32 logits/log-probs in backward + logits = lm_head(h).float() + ce = torch.nn.functional.cross_entropy(logits, tgt, reduction="sum") + if bh is not None: + with torch.no_grad(): + base_logp = torch.log_softmax(lm_head(bh).float(), -1) + kl = torch.nn.functional.kl_div(torch.log_softmax(logits, -1), base_logp, log_target=True, reduction="sum") + else: + kl = ce.new_zeros(()) + if want_info: + with torch.no_grad(): + p = logits.softmax(-1) + p_first_end = p[0, MUSIC_END] + p_tgt = p.gather(1, tgt[:, None]).squeeze(1) + else: + p_first_end, p_tgt = ce.new_zeros(()), ce.new_zeros((0,)) + return ce, kl, p_first_end, p_tgt + + ce_sum = hidden.new_zeros((), dtype=torch.float32) + kl_sum = hidden.new_zeros((), dtype=torch.float32) + p_target, p_end_first = [], None + chunk = 512 # vocab-sized fp32 transients scale with this (~1.9 GB per 512 positions incl. KL) + for s0 in range(0, n, chunk): + h = hidden[s0 : s0 + chunk] + tgt = ids[s0 : s0 + chunk] + bh = base_hidden[s0 : s0 + chunk] if want_kl else None + if torch.is_grad_enabled(): + ce, kl, p_first, p_tgt = torch.utils.checkpoint.checkpoint(chunk_losses, h, tgt, bh, use_reentrant=False) + else: + ce, kl, p_first, p_tgt = chunk_losses(h, tgt, bh) + ce_sum = ce_sum + ce + kl_sum = kl_sum + kl + if want_info: + if s0 == 0: + p_end_first = p_first.item() + p_target.append(p_tgt) + info = "" + if want_info: + p_target = torch.cat(p_target) + info = (f"[ar targets {n} unique {ids.unique().numel()} (song {total}, whole {whole_song}) " + f"prefix {prefix.shape[0]} p(END|prefix) {p_end_first:.3f} p(target) median {p_target.median().item():.3f}]") + return ce_sum / n, (kl_sum / n if want_kl else None), info + + def _ar_inputs(self, prefix: torch.Tensor, abc_ids: torch.Tensor, tokens: torch.Tensor, end_token: bool = True): + """One item: prefix head (unpadded [L, H]) + abc + [ABC_END, MUSIC_START] + codec tokens [+ MUSIC_END] + -> embeds [1, L', H] and the ids after the head [n] (the AR targets).""" + device = tokens.device + parts = [abc_ids.long().to(device), torch.tensor([ABC_END, MUSIC_START], device=device), tokens.long() + CODEC_OFFSET] + if end_token: + parts.append(torch.tensor([MUSIC_END], device=device)) + ids = torch.cat(parts) + embeds = torch.cat([prefix.to(self.model.ar.dtype), self.model.ar.embed(ids)], dim=0)[None] + return embeds, ids + + # ------------------------------------------------------------------ + # per-expert learning rate: the trainer hands us the LoRA network before it + # builds optimizer groups, so split the AR expert's modules into their own group + # ------------------------------------------------------------------ + @property + def network(self): + return self._network + + @network.setter + def network(self, value): + self._network = value + mult = getattr(self, "ar_lr_multiplier", 1.0) + if value is None or mult == 1.0 or getattr(value, "_yue2_lr_split", False): + return + orig = value.prepare_optimizer_params + + def prepare_optimizer_params(text_encoder_lr=None, unet_lr=None, default_lr=None): + groups = orig(text_encoder_lr, unet_lr, default_lr) + ar_ids = set() + for lora in getattr(value, "unet_loras", []): + if ".ar." in lora.lora_name.replace("$$", "."): + ar_ids.update(id(p) for p in lora.parameters()) + out = [] + n_ar = 0 + for group in groups: + base_lr = group.get("lr", default_lr) + ar = [p for p in group["params"] if id(p) in ar_ids] + rest = [p for p in group["params"] if id(p) not in ar_ids] + if rest: + out.append({**group, "params": rest}) + if ar: + n_ar += len(ar) + if base_lr is not None: + out.append({**group, "params": ar, "lr": base_lr * mult}) + else: + out.append({**group, "params": ar}) + if n_ar: + self.print_and_status_update(f"YuE2: AR expert LoRA learning rate x{mult} ({n_ar} tensors)") + return out + + value.prepare_optimizer_params = prepare_optimizer_params + value._yue2_lr_split = True + + def get_noise_prediction(self, latent_model_input: torch.Tensor, timestep: torch.Tensor, text_embeddings: PromptEmbeds, batch=None, **kwargs): + # `batch` must be a named parameter: predict_noise only forwards it when it sees it in the signature + if batch is None or not isinstance(batch.latents, DTO) or batch.latents.get("tokens") is None: + raise ValueError("YuE2 training needs codec tokens in the latent cache; enable latent caching") + abc_all = batch.latents.get("abc_ids") + if self.cot != "off": + mode = batch.latents.get("abc_mode") + cached = None if mode is None else int(mode.reshape(-1)[0].item()) + if abc_all is None or cached != COT_CODES[self.cot]: + raise ValueError( + f"YuE2 cot={self.cot!r} needs a matching ABC sheet in the latent cache " + f"(cache has {'none' if cached is None else COT_NAMES.get(cached, cached)!r}): " + "delete the dataset's _latent_cache folder so it is rebuilt" + ) + start, end = getattr(batch, "yue2_window", (0, latent_model_input.shape[1])) + tokens_all = batch.latents.tokens + model = self.model + if model.device == torch.device("cpu"): + model.to(self.device_torch) + device = self.device_torch + t = (timestep.to(device, dtype=torch.float32) / 1000.0).to(self.torch_dtype) + prefix_embeds = text_embeddings.text_embeds.to(device, self.torch_dtype) + prefix_mask = text_embeddings.attention_mask.to(device) + train_ar = self.ar_loss_weight > 0 and torch.is_grad_enabled() + if torch.is_grad_enabled(): + self._train_calls += 1 + if self._train_calls - self._last_encode_call > 2 and (self.semantic_tokenizer is not None or self.transcriber is not None): + # latents are cached: MERT + SheetSage2 (~4 GB) are dead weight for the rest of training + self.semantic_tokenizer = None + self.transcriber = None + flush() + captions = batch.get_caption_list() if hasattr(batch, "get_caption_list") else None + preds, lm_losses, kl_losses = [], [], [] + for i in range(latent_model_input.shape[0]): + song = tokens_all[i].to(device) + total = song.shape[0] + prefix = self._train_prefix(prefix_embeds[i], prefix_mask[i]) + if self.cot == "off" or abc_all is None: + abc = torch.zeros(0, dtype=torch.long, device=device) + elif torch.is_grad_enabled() and captions is not None and random.random() < self.abc_dropout: + # sheet dropped: the item trains exactly as an off-mode prompt (off instruction, empty ABC block) + parsed = parse_caption(captions[i]) + ids = torch.tensor([self.tokenizer.prefix_head_ids(parsed["style"], parsed["lyrics"], cot="off")], device=device) + prefix = model.ar.embed(ids)[0].to(self.torch_dtype) + abc = torch.zeros(0, dtype=torch.long, device=device) + else: + row = abc_all[i].to(device) + abc = row[row >= 0].long() + # NAR conditioning follows the released chunk protocol: prefix + abc + window + MUSIC_END + embeds, ids = self._ar_inputs(prefix, abc, song[start:end]) + whole_song = start == 0 and end == total + if train_ar and not whole_song: + # the language-model loss must see the song from its beginning, or the AR + # learns that lyrics and tokens can start anywhere relative to each other + limit = total if self.ar_max_tokens <= 0 else min(total, self.ar_max_tokens) + ar_embeds, ar_ids = self._ar_inputs(prefix, abc, song[:limit], end_token=limit == total) + ce, kl, info = self._ar_losses(ar_embeds, ar_ids, prefix, total, whole_song) + lm_losses.append(ce) + if kl is not None: + kl_losses.append(kl) + self._pending_aux_info = info + with torch.no_grad(): + cache, _ = model.ar.prefill(embeds) + elif train_ar: + ce, kl, info = self._ar_losses(embeds, ids, prefix, total, whole_song) + lm_losses.append(ce) + if kl is not None: + kl_losses.append(kl) + self._pending_aux_info = info + with torch.no_grad(): + cache, _ = model.ar.prefill(embeds) + else: + cache, _ = model.ar.prefill(embeds) + # the flow loss must not train the AR through its KV cache: only next-token CE shapes the AR + cache = [(k.detach(), v.detach()) for k, v in cache] + pred = model.nar(latent_model_input[i : i + 1].to(device, self.torch_dtype), t[i : i + 1], cache, embeds.shape[1]) + preds.append(pred) + if lm_losses: + self._pending_ar_ce = torch.stack(lm_losses).mean() + self._pending_ar_kl = torch.stack(kl_losses).mean() if kl_losses else None + self._pending_aux_loss = self._pending_ar_ce * self.ar_loss_weight + if self._pending_ar_kl is not None: + self._pending_aux_loss = self._pending_aux_loss + self._pending_ar_kl * self.ar_kl_weight + return torch.cat(preds, 0) + + # ------------------------------------------------------------------ + # sampling + # ------------------------------------------------------------------ + def get_generation_pipeline(self): + return self.pipeline + + + def generate_single_audio(self, pipeline: YuE2Pipeline, gen_config: GenerateImageConfig, conditional_embeds: PromptEmbeds, unconditional_embeds, generator, extra): + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + if gen_config.output_ext not in ["mp3", "wav", "flac"]: + gen_config.output_ext = "mp3" + parsed = parse_caption(gen_config.prompt) + # sample-config duration field first, then a [Duration] prompt line, then the model default + max_seconds = getattr(gen_config, "duration", None) or parsed["duration"] or self.sample_max_seconds + max_tokens = max(1, int(round(max_seconds * FRAMES_PER_SECOND))) + length = int((conditional_embeds.attention_mask[0] > 0.5).sum().item()) + head = conditional_embeds.text_embeds[:1, :length].to(self.device_torch, self.torch_dtype) + + # stage 1 (cot full/melody): the AR writes the ABC sheet after ABC_START, as Comfy's Generate ABC does + abc_ids: List[int] = [] + if self.cot != "off": + self._status_update("YuE2: generating ABC sheet") + abc_bar = tqdm(total=self.abc_max_tokens, desc=" abc tokens", unit="tok", position=1, leave=False, dynamic_ncols=True) + try: + abc_ids = pipeline.generate_abc_tokens( + head, seed=gen_config.seed, max_tokens=self.abc_max_tokens, progress=lambda done, total: abc_bar.update(done - abc_bar.n) + ) + finally: + abc_bar.close() + tail = torch.tensor([self.tokenizer.abc_tail_ids(abc_ids)], device=self.device_torch) + prefix = torch.cat([head, self.model.ar.embed(tail).to(head.dtype)], dim=1) + # one acoustic frame needs two positions plus the boundary tokens (Comfy's budget rule) + max_tokens = min(max_tokens, CONTEXT - prefix.shape[1] - 5) + if max_tokens < 1: + raise ValueError("YuE2 prompt leaves no room for music; shorten the style, lyrics, or ABC") + + # nested bars under generate_images' "Generating Samples" bar (position 0) + self._status_update(f"YuE2: generating codec tokens (up to {max_tokens / FRAMES_PER_SECOND:.0f}s)") + token_bar = tqdm(total=max_tokens, desc=" codec tokens", unit="tok", position=1, leave=False, dynamic_ncols=True) + + def token_progress(done, total): + token_bar.update(done - token_bar.n) + token_bar.set_postfix_str(f"{done / FRAMES_PER_SECOND:.1f}s", refresh=False) + + try: + codec = pipeline.generate_codec_tokens( + prefix, max_tokens=max_tokens, seed=gen_config.seed, temperature=self.sample_ar_temperature, + repetition_penalty=self.sample_ar_repetition_penalty, progress=token_progress, legacy_off=self.cot == "off", + ) + finally: + token_bar.close() + if len(codec) == 0: + codec = [0] + self._status_update(f"YuE2: rendering {len(codec) / FRAMES_PER_SECOND:.1f}s") + render_bar = tqdm(desc=f" render {len(codec) / FRAMES_PER_SECOND:.0f}s", unit="step", position=1, leave=False, dynamic_ncols=True) + + def render_step(i, n, x): + if render_bar.total != n: + render_bar.total = n + render_bar.update(i + 1 - render_bar.n) + self._emit_sample_step(x, i, n) + + try: + latents = pipeline.synthesize( + prefix, + codec, + seed=gen_config.seed, + steps=gen_config.num_inference_steps, + step_callback=render_step, + ) + finally: + render_bar.close() + audio = pipeline.decode(latents) # [1, 2, samples] + return audio.cpu() + + # ------------------------------------------------------------------ + # LoRA key layout: transformer.nar.* -> diffusion_model.*, transformer.ar.* -> text_encoders.* + # ------------------------------------------------------------------ + def convert_lora_weights_before_save(self, state_dict): + out = {} + for k, v in state_dict.items(): + if k.startswith("transformer.nar."): + k = "diffusion_model." + k[len("transformer.nar.") :] + elif k.startswith("transformer.ar."): + k = "text_encoders." + k[len("transformer.ar.") :] + out[k] = v + return out + + def convert_lora_weights_before_load(self, state_dict): + out = {} + for k, v in state_dict.items(): + if k.startswith("diffusion_model."): + k = "transformer.nar." + k[len("diffusion_model.") :] + elif k.startswith("text_encoders."): + k = "transformer.ar." + k[len("text_encoders.") :] + out[k] = v + return out diff --git a/extensions_built_in/captioner/Qwen3OmniCaptioner.py b/extensions_built_in/captioner/Qwen3OmniCaptioner.py index 38f2ce5..c4ad548 100644 --- a/extensions_built_in/captioner/Qwen3OmniCaptioner.py +++ b/extensions_built_in/captioner/Qwen3OmniCaptioner.py @@ -29,6 +29,8 @@ VIDEO_FPS = 2 # still-image files caption through the image pipeline (no audio, no frames) IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "bmp", "webp"} +# audio-only files caption through the audio pipeline (no frames) +AUDIO_EXTENSIONS = {"mp3", "wav", "flac", "ogg", "m4a", "aac"} # fixed generation ceiling under compiled decode: a constant max_length keeps # the static kv cache (and so the compiled decode graph) at one shape for @@ -660,9 +662,15 @@ class Qwen3OmniCaptioner(BaseCaptioner): def _is_image_file(file_path: str) -> bool: return os.path.splitext(file_path)[1].lower().lstrip(".") in IMAGE_EXTENSIONS + @staticmethod + def _is_audio_file(file_path: str) -> bool: + return os.path.splitext(file_path)[1].lower().lstrip(".") in AUDIO_EXTENSIONS + def _build_messages(self, _file_path: str): if self._is_image_file(_file_path): media = {"type": "image", "image": _file_path} + elif self._is_audio_file(_file_path): + media = {"type": "audio", "audio": _file_path} else: media = {"type": "video", "video": _file_path} return [ @@ -695,6 +703,10 @@ class Qwen3OmniCaptioner(BaseCaptioner): image = Image.open(file_path).convert("RGB") item = {"file": file_path, "kind": "image", "image": image, "audio": None} + elif self._is_audio_file(file_path): + from transformers.audio_utils import load_audio + + item = {"file": file_path, "kind": "audio", "audio": load_audio(file_path, sampling_rate=16000)} else: from transformers.video_utils import load_video from transformers.audio_utils import load_audio @@ -738,6 +750,13 @@ class Qwen3OmniCaptioner(BaseCaptioner): padding=True, size=self._size_kwargs(), ) + if kind == "audio": + return self.processor( + text=[it["text"] for it in items], + audio=[it["audio"] for it in items], + return_tensors="pt", + padding=True, + ) use_audio = kind == "video_audio" return self.processor( text=[it["text"] for it in items], @@ -770,10 +789,14 @@ class Qwen3OmniCaptioner(BaseCaptioner): # under static cache, generate hands the forward a prepared 4D mask; # the true 2D padding mask is needed for the prefill rope index self.model._pad_mask_2d = inputs.get("attention_mask", None) + gen_kwargs = self._gen_kwargs(inputs["input_ids"].shape[1]) + if items[0]["kind"] == "audio": + # greedy lyric transcription loops on syllables; mild penalty (the model's own talker default) + gen_kwargs["repetition_penalty"] = 1.05 generated_ids = self.model.generate( **inputs, use_audio_in_video=use_audio, - **self._gen_kwargs(inputs["input_ids"].shape[1]), + **gen_kwargs, ) trimmed = generated_ids[:, inputs["input_ids"].shape[1] :] captions = self.processor.batch_decode( @@ -876,7 +899,7 @@ class Qwen3OmniCaptioner(BaseCaptioner): futures.append((path, executor.submit(self._prep_media, path))) # batches must be homogeneous: the processor call differs per kind - buckets = {"image": [], "video_audio": [], "video_silent": []} + buckets = {"image": [], "audio": [], "video_audio": [], "video_silent": []} while futures: if self.is_ui_captioner: self.maybe_stop() diff --git a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py index 4f64aba..772acbf 100644 --- a/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py +++ b/extensions_built_in/diffusion_models/hidream/hidream_o1_model.py @@ -14,6 +14,7 @@ from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from safetensors.torch import load_file, save_file +from toolkit.util.streamed_safetensors import save_file_streamed from toolkit.accelerator import unwrap_model from optimum.quanto import freeze @@ -502,7 +503,8 @@ class HidreamO1Model(BaseModel): if not output_path.endswith(".safetensors"): output_path += ".safetensors" meta = get_meta_for_safetensors(meta, name=self.arch) - save_file(save_dict, output_path, metadata=meta) + # sequential writer: mmap save_file fails with EINVAL on ntfs3 for large files + save_file_streamed(save_dict, output_path, metadata=meta) else: transformer.save_pretrained( save_directory=output_path, diff --git a/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..a928639 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_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8", "model_kwargs": {"sample_max_seconds": 15, "cot": "off"}}, + "sample": {"width": 512, "height": 512, "num_inference_steps": 16, "guidance_scale": 1.0, "seed": 42}, + }, "f-lite": { "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 25, "guidance_scale": 4.0, "seed": 42}, 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..18d90dd 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_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}}, } OUTPUT_EXT = {"image": "png", "video": "mp4", "audio": "mp3"} diff --git a/toolkit/util/comfy_quant_import.py b/toolkit/util/comfy_quant_import.py index 6cd16b6..4634f02 100644 --- a/toolkit/util/comfy_quant_import.py +++ b/toolkit/util/comfy_quant_import.py @@ -247,7 +247,31 @@ def import_comfy_quantized_layers( ) parent_path, _, attr = module_path.rpartition(".") parent = root.get_submodule(parent_path) if parent_path else root - setattr(parent, attr, Int8Embedding(weight, weight_scale, orig_dtype)) + if conf.get("convrot"): + # rows are stored rotated like a linear's weight: attach to a stand-in linear + # (same shape) so the convrot backend undoes the rotation, then keep the table + # as a plain embedding (per-row lookups cannot apply the group rotation) + stand_in = torch.nn.Linear( + module.embedding_dim, module.num_embeddings, bias=False, device="meta" + ) + quantizer = get_ostris_quantizer("convrot8") + stand_in.register_buffer("cr8_qdata", weight.contiguous(), persistent=False) + stand_in.register_buffer( + "cr8_scales", + weight_scale.detach().float().reshape(-1).contiguous().view(torch.uint8), + persistent=False, + ) + stand_in.cr8_rot_size = int(conf.get("convrot_groupsize", 256)) + _to_ostris(stand_in, quantizer, orig_dtype) + table = quantizer.dequantize_folded(stand_in).to(orig_dtype) + new = torch.nn.Embedding( + module.num_embeddings, module.embedding_dim, padding_idx=module.padding_idx, + device="meta", dtype=orig_dtype, + ) + new.weight = torch.nn.Parameter(table, requires_grad=False) + setattr(parent, attr, new) + else: + setattr(parent, attr, Int8Embedding(weight, weight_scale, orig_dtype)) converted += 1 continue diff --git a/toolkit/util/streamed_safetensors.py b/toolkit/util/streamed_safetensors.py new file mode 100644 index 0000000..7255cbf --- /dev/null +++ b/toolkit/util/streamed_safetensors.py @@ -0,0 +1,62 @@ +import json +import struct +from typing import Dict, Optional + +import torch + +# safetensors dtype tags +_DTYPE_TAGS = { + torch.float64: "F64", + torch.float32: "F32", + torch.float16: "F16", + torch.bfloat16: "BF16", + torch.int64: "I64", + torch.int32: "I32", + torch.int16: "I16", + torch.int8: "I8", + torch.uint8: "U8", + torch.bool: "BOOL", +} + +_CHUNK_BYTES = 256 * 1024 * 1024 + + +def save_file_streamed( + tensors: Dict[str, torch.Tensor], + path: str, + metadata: Optional[Dict[str, str]] = None, +): + """Write a safetensors file with sequential writes instead of ftruncate+mmap. + Some filesystems (kernel ntfs3) reject one large ftruncate; incremental writes succeed.""" + header = {} + offset = 0 + flat = {} + for name, t in tensors.items(): + if t.dtype not in _DTYPE_TAGS: + raise ValueError(f"Unsupported dtype {t.dtype} for tensor {name}") + t = t.detach().to("cpu").contiguous() + nbytes = t.numel() * t.element_size() + header[name] = { + "dtype": _DTYPE_TAGS[t.dtype], + "shape": list(t.shape), + "data_offsets": [offset, offset + nbytes], + } + flat[name] = t + offset += nbytes + if metadata is not None: + header["__metadata__"] = {str(k): str(v) for k, v in metadata.items()} + + header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8") + # pad header to 8-byte alignment like the reference writer + pad = (8 - len(header_bytes) % 8) % 8 + header_bytes += b" " * pad + + with open(path, "wb") as f: + f.write(struct.pack(" {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' ? (