diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 9de244087..c6660846d 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -180,6 +180,7 @@ parser.add_argument("--disable-async-offload", action="store_true", help="Disabl parser.add_argument("--disable-dynamic-vram", action="store_true", help="Disable dynamic VRAM and use estimate based model loading.") parser.add_argument("--enable-dynamic-vram", action="store_true", help="Enable dynamic VRAM on systems where it's not enabled by default.") parser.add_argument("--fast-disk", action="store_true", help="Prefer disk-backed dynamic loading and offload over unpinned RAM. Can be faster for users with fast NVME disks.") +parser.add_argument("--disable-cuda-graphs", action="store_true", help="Disable CUDA graphs.") parser.add_argument("--force-non-blocking", action="store_true", help="Force ComfyUI to use non-blocking operations for all applicable tensors. This may improve performance on some non-Nvidia systems but can cause issues with some workflows.") diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index c4270022b..dc737fc7d 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -957,6 +957,11 @@ class ACEAudio15(LatentFormat): latent_dimensions = 1 temporal_downscale_ratio = 1764 +class MiniMaxMusic3(LatentFormat): + latent_channels = 128 + latent_dimensions = 1 + temporal_downscale_ratio = 512 + class ChromaRadiance(LatentFormat): latent_channels = 3 spacial_downscale_ratio = 1 diff --git a/comfy/ldm/minimax_music/__init__.py b/comfy/ldm/minimax_music/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/comfy/ldm/minimax_music/ar.py b/comfy/ldm/minimax_music/ar.py new file mode 100644 index 000000000..28215cb58 --- /dev/null +++ b/comfy/ldm/minimax_music/ar.py @@ -0,0 +1,337 @@ +import dataclasses +import hashlib + +import torch +from torch import nn + +import comfy.model_management +import comfy.model_prefetch +import comfy.ops +import comfy.utils +from comfy.ldm.modules.attention import optimized_attention_for_device +from comfy.text_encoders.llama import Llama2_, Qwen3_8BConfig + +from .prompt import AUDIO_CODE_OFFSET, SPECIAL_TOKEN_IDS + + +CFG_SCALE = 1.5 +CFG_TOP_K = 50 +C0_VOCAB_SIZE = 16384 +MAX_PROMPT_TOKENS = 5000 +MAX_AUDIO_FRAMES = 9000 +AUDIO_FRAMES_PER_SECOND = 25 + + +def derive_seed(seed, *parts): + digest = hashlib.blake2b(digest_size=8, person=b"minimax-ttm") + digest.update(int(seed).to_bytes(8, "little", signed=False)) + for part in parts: + value = str(part).encode("utf-8") + digest.update(len(value).to_bytes(4, "little")) + digest.update(value) + return int.from_bytes(digest.digest(), "little") & ((1 << 63) - 1) + + +def sample_topk(logits, top_k, generator): + values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9) + top_k = min(top_k, values.shape[-1]) + threshold = torch.topk(values, top_k, dim=-1).values[..., -1, None] + values = values.masked_fill(values < threshold, -float("inf")) + probabilities = torch.nan_to_num(torch.softmax(values, dim=-1), nan=0.0) + probabilities = probabilities / probabilities.sum(dim=-1, keepdim=True).clamp_min(1e-12) + return torch.multinomial(probabilities, 1, generator=generator).squeeze(-1) + + +class RVQAttention(nn.Module): + def __init__(self, hidden_size, num_heads, dtype, device, operations): + super().__init__() + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.merged_qkv = None + self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device) + self.q_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) + self.k_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) + self.v_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) + self.o_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) + + def forward(self, x): + batch, length, hidden_size = x.shape + if self.merged_qkv: + q, k, v = self.qkv_proj(x).chunk(3, dim=-1) + else: + q = self.q_proj(x) + k = self.k_proj(x) + v = self.v_proj(x) + q = q.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2) + k = k.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2) + v = v.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2) + mask = torch.full((length, length), torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype).triu_(1) + attention = optimized_attention_for_device(q.device, mask=True, small_input=True) + out = attention(q, k, v, self.num_heads, mask=mask, skip_reshape=True) + return self.o_proj(out) + + +class RVQRMSNorm(nn.Module): + def __init__(self, hidden_size, dtype, device): + super().__init__() + self.weight = nn.Parameter(torch.empty(hidden_size, dtype=dtype, device=device)) + + def forward(self, x): + return torch.nn.functional.rms_norm(x, (x.shape[-1],), comfy.ops.cast_to_input(self.weight, x), 1e-6) + + +class RVQMLP(nn.Module): + def __init__(self, hidden_size, intermediate_size, dtype, device, operations): + super().__init__() + self.merged_mlp = None + self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device) + self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device) + self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device) + self.down_proj = operations.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype, device=device) + + def forward(self, x): + if self.merged_mlp: + return comfy.ops.linear_input_act(self.down_proj, self.gate_up_proj(x), "swiglu") + return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class RVQDecoderBlock(nn.Module): + def __init__(self, hidden_size, num_heads, intermediate_size, dtype, device, operations): + super().__init__() + self.input_layernorm = RVQRMSNorm(hidden_size, dtype, device) + self.self_attn = RVQAttention(hidden_size, num_heads, dtype, device, operations) + self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device) + self.mlp = RVQMLP(hidden_size, intermediate_size, dtype, device, operations) + + def forward(self, x): + x = x + self.self_attn(self.input_layernorm(x)) + return x + self.mlp(self.post_attention_layernorm(x)) + + +class RVQDepthDecoder(nn.Module): + def __init__(self, config, dtype, device, operations): + super().__init__() + hidden_size = int(config["hidden_size"]) + audio_vocab_size = int(config["audio_vocab_size"]) + num_codebooks = int(config["audio_num_codebooks"]) + self.projection = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device) + self.pos_embedding = operations.Embedding(16, hidden_size, dtype=dtype, device=device) + self.audio_heads = nn.ModuleList([ + operations.Linear(hidden_size, audio_vocab_size, bias=False, dtype=dtype, device=device) + for _ in range(num_codebooks - 1) + ]) + self.layers = nn.ModuleList([ + RVQDecoderBlock( + hidden_size, + int(config["decoder_num_heads"]), + int(config["decoder_intermediate_size"]), + dtype, + device, + operations, + ) + for _ in range(int(config["decoder_num_layers"])) + ]) + self.norm = RVQRMSNorm(hidden_size, dtype, device) + + def forward(self, sequence): + positions = torch.arange(sequence.shape[1], device=sequence.device) + x = sequence + self.pos_embedding(positions, out_dtype=sequence.dtype).unsqueeze(0) + for layer in self.layers: + x = layer(x) + return self.norm(x) + + +class MiniMaxMusic3AR(nn.Module): + def __init__(self, config, dtype, device, operations): + super().__init__() + config_fields = {field.name for field in dataclasses.fields(Qwen3_8BConfig)} + qwen_config = Qwen3_8BConfig(**{key: value for key, value in config.items() if key in config_fields}) + qwen_config.lm_head = False + qwen_config.fixed_kv = True + qwen_config.merged_qkv = None + qwen_config.merged_mlp = None + self.model = Llama2_(qwen_config, device=device, dtype=dtype, ops=operations) + self.model.prefetch_dynamic_vbars = True + self.model.graph_dynamic_vbar_blocks = True + self.model.lm_head = operations.Linear(qwen_config.hidden_size, qwen_config.vocab_size, bias=False, dtype=dtype, device=device) + self.model.lm_head_pruned = operations.Linear(qwen_config.hidden_size, C0_VOCAB_SIZE + 1, bias=False, dtype=dtype, device=device) + self.model.embed_tokens_prefill = operations.Embedding(AUDIO_CODE_OFFSET, qwen_config.hidden_size, dtype=dtype, device=device) + self.model.embed_tokens_audio = operations.Embedding(C0_VOCAB_SIZE, qwen_config.hidden_size, dtype=dtype, device=device) + self.model.pruned_lm_head = None + self.model.pruned_embedding = None + self.model.audio_extra_embedding = operations.Embedding( + int(config["audio_vocab_size"]) * (int(config["audio_num_codebooks"]) - 1), + qwen_config.hidden_size, + dtype=dtype, + device=device, + ) + self.model.audio_decoder = RVQDepthDecoder(config, dtype, device, operations) + self.audio_vocab_size = int(config["audio_vocab_size"]) + self.num_codebooks = int(config["audio_num_codebooks"]) + self.embedding_scale = self.num_codebooks ** -0.5 + + def _guided_c0(self, logits, cfg_scale, top_k): + conditioned = logits[0:1].float() + unconditioned = logits[1:2].float() + guided = unconditioned + (conditioned - unconditioned) * cfg_scale + threshold = torch.topk(conditioned, top_k, dim=-1).values[..., -1, None] + return guided.masked_fill(conditioned < threshold, -float("inf")) + + def _depth_codes(self, hidden, c0, c0_embed, generator, execution_dtype, cfg_scale, top_k): + decoder = self.model.audio_decoder + sequence = [decoder.projection(hidden).unsqueeze(1)] + sequence.append(decoder.projection(c0_embed).unsqueeze(1)) + codes = [c0] + hidden_parts = [] + for index in range(1, self.num_codebooks): + out = decoder(torch.cat(sequence, dim=1))[:, -1] + hidden_parts.append(out[:1].detach()) + logits = decoder.audio_heads[index - 1](out) + conditioned = logits[:1].float() + unconditioned = logits[1:2].float() + code = sample_topk(unconditioned + (conditioned - unconditioned) * cfg_scale, top_k, generator).repeat(2) + codes.append(code) + if index < self.num_codebooks - 1: + embedding = self.model.audio_extra_embedding( + code + (index - 1) * self.audio_vocab_size, + out_dtype=execution_dtype, + ) + sequence.append(decoder.projection(embedding).unsqueeze(1)) + return torch.stack(codes, dim=1), torch.cat(hidden_parts, dim=-1) + + def _embed_c0(self, codes, execution_dtype): + if self.model.pruned_embedding: + return self.model.embed_tokens_audio(codes, out_dtype=execution_dtype) + return self.model.embed_tokens(codes + AUDIO_CODE_OFFSET, out_dtype=execution_dtype) + + def _embed_audio_frame(self, codes, execution_dtype): + c0 = self._embed_c0(codes[:, 0], execution_dtype) + offsets = torch.arange(self.num_codebooks - 1, device=codes.device) * self.audio_vocab_size + extra = self.model.audio_extra_embedding(codes[:, 1:] + offsets.unsqueeze(0), out_dtype=execution_dtype).sum(dim=1) + return ((c0 + extra) * self.embedding_scale).unsqueeze(1) + + def _sample_c0(self, hidden, cfg_scale, top_k, generator, vocab_mask): + if self.model.pruned_lm_head: + guided = self._guided_c0(self.model.lm_head_pruned(hidden).float(), cfg_scale, top_k) + code = sample_topk(guided, top_k, generator) + stop_token = 0 + offset = 1 + else: + logits = self.model.lm_head(hidden).float() + stop_token = SPECIAL_TOKEN_IDS["<|audio_end|>"] + logits = logits.masked_fill(vocab_mask, -float("inf")) + guided = self._guided_c0(logits, cfg_scale, top_k).masked_fill(vocab_mask, -float("inf")) + code = sample_topk(guided, top_k, generator) + offset = AUDIO_CODE_OFFSET + return torch.where(code == stop_token, 0, code - offset), code, stop_token + + def generate(self, input_ids, seed, max_audio_frames, device, cfg_scale=CFG_SCALE, top_k=CFG_TOP_K): + prompt_tokens = int(input_ids.shape[1]) + if prompt_tokens > MAX_PROMPT_TOKENS: + raise ValueError(f"MiniMax Music3 prompt has {prompt_tokens} tokens; maximum is {MAX_PROMPT_TOKENS}") + + input_ids = input_ids.to(device) + if comfy.model_management.should_use_bf16(device): + execution_dtype = torch.bfloat16 + else: + execution_dtype = torch.float32 + unconditioned = input_ids.clone() + unconditioned[:, 1:-2] = SPECIAL_TOKEN_IDS["<|audio_cfg|>"] + text_ids = torch.cat((input_ids, unconditioned), dim=0) + if self.model.pruned_embedding: + text_embeds = self.model.embed_tokens_prefill(text_ids, out_dtype=execution_dtype) + else: + text_embeds = self.model.embed_tokens(text_ids, out_dtype=execution_dtype) + decode_limit = min(int(max_audio_frames), MAX_AUDIO_FRAMES) + past = self.model.init_kv_cache(2, prompt_tokens + decode_limit + 1, device, execution_dtype) + output = self.model(None, embeds=text_embeds, past_key_values=past, dtype=execution_dtype) + last_hidden = output[0][:, -1] + past = output[2] + + generator = torch.Generator(device=device).manual_seed(derive_seed(seed, "ar")) + decoder = self.model.audio_decoder + depth_io = { + "hidden": torch.empty_like(last_hidden), + "c0": torch.empty((last_hidden.shape[0],), dtype=torch.long, device=device), + "c0_embed": torch.empty_like(last_hidden), + "codes": torch.empty((last_hidden.shape[0], self.num_codebooks), dtype=torch.long, device=device), + "depth_hidden": torch.empty((1, last_hidden.shape[-1] * (self.num_codebooks - 1)), dtype=execution_dtype, device=device), + } + decoder._comfy_cross_step_state = depth_io + comfy.model_management._register_cross_step(decoder) + hidden_frames = [] + pending_code = None + stop_token = None + pending_event = None + pending_hidden = None + progress = comfy.utils.ProgressBar(decode_limit) + cuda_device = torch.device(device).type == "cuda" + vocab_mask = None + if not self.model.pruned_lm_head: + vocab_mask = torch.ones(self.model.vocab_size, dtype=torch.bool, device=device) + vocab_mask[AUDIO_CODE_OFFSET:AUDIO_CODE_OFFSET + C0_VOCAB_SIZE] = False + vocab_mask[SPECIAL_TOKEN_IDS["<|audio_end|>"]] = False + + for frame_index in comfy.utils.model_trange(decode_limit + 1, desc="AR sampling"): + comfy.model_management.throw_exception_if_processing_interrupted() + if pending_code is not None: + if pending_event is not None: + pending_event.synchronize() + if int(pending_code.item()) == stop_token: + pending_hidden = None + break + if pending_hidden is not None: + hidden_frames.append(pending_hidden) + progress.update_absolute(len(hidden_frames)) + if len(hidden_frames) >= decode_limit: + break + + c0, code_or_stop, stop_token = self._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask) + if pending_code is None: + pending_code = torch.empty_like(code_or_stop, device="cpu", pin_memory=cuda_device) + if cuda_device: + pending_event = torch.cuda.Event() + pending_code.copy_(code_or_stop, non_blocking=cuda_device) + if pending_event is not None: + pending_event.record() + + c0 = c0.repeat(2) + c0_embed = self._embed_c0(c0, execution_dtype) + depth_io["hidden"].copy_(last_hidden) + depth_io["c0"].copy_(c0) + depth_io["c0_embed"].copy_(c0_embed) + + def depth_core(): + codes, depth_hidden = self._depth_codes( + depth_io["hidden"], depth_io["c0"], depth_io["c0_embed"], generator, execution_dtype, cfg_scale, top_k + ) + depth_io["codes"].copy_(codes) + depth_io["depth_hidden"].copy_(depth_hidden) + + depth_queue = comfy.model_prefetch.make_prefetch_queue( + [[decoder, self.model.audio_extra_embedding]], device, {"prefetch_dynamic_vbars": True} + ) + comfy.model_prefetch.prefetch_queue_pop( + depth_queue, device, decoder, execution_dtype, core=depth_core, enable_graph=True, generator=generator + ) + comfy.model_prefetch.prefetch_queue_pop(depth_queue, device, None) + feedback_codes = depth_io["codes"] + depth_hidden = depth_io["depth_hidden"] + frame_hidden = torch.cat((last_hidden[:1].detach(), depth_hidden), dim=-1) + if frame_index > 0: + pending_hidden = frame_hidden[0].clone() + + feedback = self._embed_audio_frame(feedback_codes, execution_dtype) + output = self.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype) + last_hidden = output[0][:, -1] + past = output[2] + + if pending_hidden is not None and len(hidden_frames) < decode_limit: + if pending_event is not None: + pending_event.synchronize() + if int(pending_code.item()) != stop_token: + hidden_frames.append(pending_hidden) + + if not hidden_frames: + raise ValueError("MiniMax Music3 generated zero audio frames") + return torch.stack(hidden_frames).to(device="cpu") diff --git a/comfy/ldm/minimax_music/dav.py b/comfy/ldm/minimax_music/dav.py new file mode 100644 index 000000000..d442559f4 --- /dev/null +++ b/comfy/ldm/minimax_music/dav.py @@ -0,0 +1,137 @@ +import math + +import torch +from torch import nn + +import comfy.ops + + +def snake(x, alpha): + shape = x.shape + flat = x.reshape(shape[0], shape[1], -1) + alpha = comfy.ops.cast_to_input(alpha, flat) + flat = flat + (alpha + 1e-9).reciprocal() * torch.sin(alpha * flat).pow(2) + return flat.reshape(shape) + + +class Snake1d(nn.Module): + def __init__(self, channels, dtype, device): + super().__init__() + self.alpha = nn.Parameter(torch.empty(1, channels, 1, dtype=dtype, device=device)) + + def forward(self, x): + return snake(x, self.alpha) + + +def _weight_norm_conv(operations, *args, **kwargs): + return nn.utils.parametrizations.weight_norm(operations.Conv1d(*args, **kwargs)) + + +def _weight_norm_conv_transpose(operations, *args, **kwargs): + return nn.utils.parametrizations.weight_norm(operations.ConvTranspose1d(*args, **kwargs)) + + +class ResidualUnit(nn.Module): + def __init__(self, dim, dilation, dtype, device, operations): + super().__init__() + padding = 3 * dilation + self.block = nn.Sequential( + Snake1d(dim, dtype, device), + _weight_norm_conv( + operations, + dim, + dim, + kernel_size=7, + dilation=dilation, + padding=padding, + dtype=dtype, + device=device, + ), + Snake1d(dim, dtype, device), + _weight_norm_conv(operations, dim, dim, kernel_size=1, dtype=dtype, device=device), + ) + + def forward(self, x): + residual = self.block(x) + if residual.shape[-1] != x.shape[-1]: + padding = (x.shape[-1] - residual.shape[-1]) // 2 + x = x[..., padding:x.shape[-1] - padding] + return x + residual + + +class DecoderBlock(nn.Module): + def __init__(self, input_dim, output_dim, stride, dtype, device, operations): + super().__init__() + self.block = nn.Sequential( + Snake1d(input_dim, dtype, device), + _weight_norm_conv_transpose( + operations, + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + dtype=dtype, + device=device, + ), + ResidualUnit(output_dim, 1, dtype, device, operations), + ResidualUnit(output_dim, 3, dtype, device, operations), + ResidualUnit(output_dim, 9, dtype, device, operations), + ) + + def forward(self, x): + return self.block(x) + + +class Decoder(nn.Module): + def __init__(self, dtype, device, operations): + super().__init__() + layers = [ + _weight_norm_conv( + operations, + 1024, + 1536, + kernel_size=7, + padding=3, + dtype=dtype, + device=device, + ) + ] + channels = 1536 + output_dim = channels + for index, stride in enumerate((8, 8, 4, 2)): + input_dim = channels // (2 ** index) + output_dim = channels // (2 ** (index + 1)) + layers.append(DecoderBlock(input_dim, output_dim, stride, dtype, device, operations)) + layers.extend(( + Snake1d(output_dim, dtype, device), + _weight_norm_conv( + operations, + output_dim, + 1, + kernel_size=7, + padding=3, + dtype=dtype, + device=device, + ), + nn.Tanh(), + )) + self.model = nn.Sequential(*layers) + + def forward(self, x): + return self.model(x) + + +class MiniMaxMusic3DAV(nn.Module): + def __init__(self, dtype=None, device=None, operations=None): + super().__init__() + self.dec_in_proj = operations.Conv1d(64, 1024, kernel_size=1, dtype=dtype, device=device) + self.decoder = Decoder(dtype, device, operations) + + def decode(self, latent): + batch, _, frames = latent.shape + folded = latent.reshape(batch * 2, 64, frames) + waveform = self.decoder(self.dec_in_proj(folded)) + return waveform.reshape(batch, 2, -1) + + forward = decode diff --git a/comfy/ldm/minimax_music/dit.py b/comfy/ldm/minimax_music/dit.py new file mode 100644 index 000000000..211e0d7db --- /dev/null +++ b/comfy/ldm/minimax_music/dit.py @@ -0,0 +1,213 @@ +import math + +import torch +from torch import nn + +import comfy.model_management +import comfy.ops +import comfy.quant_ops +from comfy.ldm.modules.attention import optimized_attention_for_device + + +MAX_CONDITION_FRAMES = 200 +CONDITION_HOP_FRAMES = 100 + + +def latent_length(audio_frames): + return max(1, int(audio_frames * 44100 / 24000 * 960 / 512)) + + +class FourierFeatures(nn.Module): + def __init__(self, in_features, out_features, dtype, device): + super().__init__() + self.weight = nn.Parameter(torch.empty(out_features // 2, in_features, dtype=dtype, device=device)) + + def forward(self, value): + weight = comfy.ops.cast_to_input(self.weight, value) + features = 2.0 * math.pi * value @ weight.T + return torch.cat((features.cos(), features.sin()), dim=-1) + + +class LayerNorm(nn.Module): + def __init__(self, dim, dtype, device): + super().__init__() + self.gamma = nn.Parameter(torch.empty(dim, dtype=dtype, device=device)) + self.register_buffer("beta", torch.empty(dim, dtype=dtype, device=device)) + + def forward(self, x): + return torch.nn.functional.layer_norm( + x, + (x.shape[-1],), + comfy.ops.cast_to_input(self.gamma, x), + comfy.ops.cast_to_input(self.beta, x), + ) + + +class RotaryEmbedding(nn.Module): + def __init__(self, dim, dtype, device): + super().__init__() + self.register_buffer("inv_freq", torch.empty(dim // 2, dtype=dtype, device=device)) + + def forward_from_seq_len(self, length, device, dtype): + positions = torch.arange(length, device=device, dtype=torch.float32) + frequencies = torch.outer(positions, comfy.ops.cast_to_input(self.inv_freq, positions)) + frequencies = frequencies.to(dtype) + cos, sin = frequencies.cos(), frequencies.sin() + return torch.stack((cos, -sin, sin, cos), dim=-1).reshape(1, 1, length, frequencies.shape[-1], 2, 2) + + +def _apply_rope(x, rotation_matrix): + x_dtype = x.dtype + x = x.reshape(*x.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).to(rotation_matrix.dtype) + x = rotation_matrix[..., 0] * x[..., 0] + rotation_matrix[..., 1] * x[..., 1] + return x.movedim(-1, -2).flatten(-2).to(x_dtype) + + +class Attention(nn.Module): + def __init__(self, dim, dim_heads, dtype, device, operations): + super().__init__() + self.num_heads = dim // dim_heads + self.dim_heads = dim_heads + self.to_qkv = operations.Linear(dim, dim * 3, bias=False, dtype=dtype, device=device) + self.to_out = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device) + + def forward(self, x, rotation_matrix): + batch, length, dim = x.shape + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + q = q.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2) + k = k.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2) + v = v.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2) + rotary_dims = rotation_matrix.shape[-3] * 2 + if comfy.model_management.in_training: + q = torch.cat((_apply_rope(q[..., :rotary_dims], rotation_matrix), q[..., rotary_dims:]), dim=-1) + k = torch.cat((_apply_rope(k[..., :rotary_dims], rotation_matrix), k[..., rotary_dims:]), dim=-1) + else: + rotated_q, rotated_k = comfy.quant_ops.ck.apply_rope_split_half(q[..., :rotary_dims], k[..., :rotary_dims], rotation_matrix) + q = torch.cat((rotated_q, q[..., rotary_dims:]), dim=-1) + k = torch.cat((rotated_k, k[..., rotary_dims:]), dim=-1) + attention = optimized_attention_for_device(q.device) + out = attention(q, k, v, self.num_heads, skip_reshape=True) + return self.to_out(out) + + +class GLU(nn.Module): + def __init__(self, dim, inner_dim, dtype, device, operations): + super().__init__() + self.proj = operations.Linear(dim, inner_dim * 2, dtype=dtype, device=device) + + def forward(self, x): + value, gate = self.proj(x).chunk(2, dim=-1) + return value * torch.nn.functional.silu(gate) + + +class FeedForward(nn.Module): + def __init__(self, dim, inner_dim, dtype, device, operations): + super().__init__() + self.ff = nn.Sequential( + GLU(dim, inner_dim, dtype, device, operations), + nn.Identity(), + operations.Linear(inner_dim, dim, dtype=dtype, device=device), + ) + + def forward(self, x): + return self.ff(x) + + +class TransformerBlock(nn.Module): + def __init__(self, dim, dim_heads, inner_dim, dtype, device, operations): + super().__init__() + self.pre_norm = LayerNorm(dim, dtype, device) + self.self_attn = Attention(dim, dim_heads, dtype, device, operations) + self.ff_norm = LayerNorm(dim, dtype, device) + self.ff = FeedForward(dim, inner_dim, dtype, device, operations) + + def forward(self, x, rotation_matrix): + x = x + self.self_attn(self.pre_norm(x), rotation_matrix) + return x + self.ff(self.ff_norm(x)) + + +class ContinuousTransformer(nn.Module): + def __init__(self, dtype, device, operations): + super().__init__() + self.project_in = operations.Linear(2304, 2048, bias=False, dtype=dtype, device=device) + self.project_out = operations.Linear(2048, 128, bias=False, dtype=dtype, device=device) + self.rotary_pos_emb = RotaryEmbedding(32, dtype, device) + self.layers = nn.ModuleList([ + TransformerBlock(2048, 64, 8192, dtype, device, operations) + for _ in range(36) + ]) + + def forward(self, x, timestep_embedding): + x = self.project_in(x) + x = torch.cat((timestep_embedding.unsqueeze(1), x), dim=1) + rotation_matrix = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], x.device, x.dtype) + for layer in self.layers: + x = layer(x, rotation_matrix) + return self.project_out(x[:, 1:]) + + +class DiffusionTransformer(nn.Module): + def __init__(self, dtype, device, operations): + super().__init__() + self.transformer = ContinuousTransformer(dtype, device, operations) + self.timestep_features = FourierFeatures(1, 256, dtype, device) + self.to_timestep_embed = nn.Sequential( + operations.Linear(256, 2048, dtype=dtype, device=device), + nn.SiLU(), + operations.Linear(2048, 2048, dtype=dtype, device=device), + ) + self.preprocess_conv = operations.Conv1d(2304, 2304, 1, bias=False, dtype=dtype, device=device) + self.postprocess_conv = operations.Conv1d(128, 128, 1, bias=False, dtype=dtype, device=device) + + def forward(self, x, timestep, condition): + full = torch.cat((x, torch.zeros_like(x), condition), dim=1) + full = self.preprocess_conv(full) + full + timestep_features = self.timestep_features(timestep[:, None]).to(dtype=x.dtype) + timestep_embedding = self.to_timestep_embed(timestep_features) + out = self.transformer(full.transpose(1, 2), timestep_embedding).transpose(1, 2) + return self.postprocess_conv(out) + out + + +class MiniMaxMusic3DiT(nn.Module): + def __init__(self, dtype=None, device=None, operations=None, **kwargs): + super().__init__() + self.dtype = dtype + self.latent_conditioners = nn.Sequential( + operations.Conv1d(4096, 2048, kernel_size=3, padding=1, dtype=dtype, device=device) + ) + self.diffusion_transformer = DiffusionTransformer(dtype, device, operations) + self.cond_layer_logits = nn.Parameter(torch.empty(8, dtype=dtype, device=device)) + self.cond_layer_scale = nn.Parameter(torch.empty(1, dtype=dtype, device=device)) + + def aligned_condition(self, hidden): + frames = hidden.shape[1] + hidden = hidden.transpose(1, 2).reshape(hidden.shape[0], 8, 4096, frames) + weights = torch.softmax(comfy.ops.cast_to_input(self.cond_layer_logits, hidden), dim=0) + hidden = torch.einsum("blht,l->bht", hidden, weights) + hidden = comfy.ops.cast_to_input(self.cond_layer_scale, hidden) * hidden + condition = self.latent_conditioners(hidden) + return torch.nn.functional.interpolate(condition, size=latent_length(frames), mode="nearest") + + def forward(self, x, timestep, context, conditioning_scale, **kwargs): + condition = self.aligned_condition(context) + condition = condition * conditioning_scale[:, :1, :1] + if condition.shape[-1] < x.shape[-1]: + condition = torch.nn.functional.pad(condition, (0, x.shape[-1] - condition.shape[-1])) + else: + condition = condition[..., :x.shape[-1]] + window = latent_length(MAX_CONDITION_FRAMES) + if x.shape[-1] <= window: + return -self.diffusion_transformer(x, timestep, condition) + + output = torch.zeros_like(x) + count = torch.zeros((1, 1, x.shape[-1]), device=x.device, dtype=x.dtype) + hop = latent_length(CONDITION_HOP_FRAMES) + start = 0 + while start < x.shape[-1]: + end = min(start + window, x.shape[-1]) + output[..., start:end] -= self.diffusion_transformer(x[..., start:end], timestep, condition[..., start:end]) + count[..., start:end] += 1 + if end == x.shape[-1]: + break + start += hop + return output / count diff --git a/comfy/ldm/minimax_music/prompt.py b/comfy/ldm/minimax_music/prompt.py new file mode 100644 index 000000000..5f197ee12 --- /dev/null +++ b/comfy/ldm/minimax_music/prompt.py @@ -0,0 +1,70 @@ +import re + + +SPECIAL_TOKEN_IDS = { + "<|im_start|>": 151644, + "<|im_end|>": 151645, + "<|audio_cfg|>": 151654, + "<|audio_start|>": 151669, + "<|audio_end|>": 151670, + "<|caption_start|>": 151671, + "<|caption_end|>": 151672, + "<|lyrics_start|>": 151673, + "<|lyrics_end|>": 151674, +} +AUDIO_CODE_OFFSET = 151675 + +_SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>") +_LYRIC_TAG_RE = re.compile(r"\s*(\[[^\]]+\])\s*") + + +def _remove_markdown_format(text): + lines = [] + for raw_line in text.splitlines(): + line = re.sub(r"^\s{0,3}#{1,6}\s+", "", raw_line) + line = re.sub(r"^\s*[*+-]\s+", "", line) + while "**" in line: + updated = re.sub(r"\*\*([^*]+)\*\*", r"\1", line) + if updated == line: + break + line = updated + line = re.sub(r"(?<|caption_start|>" + f"{clean_caption(caption)}" + "<|caption_end|><|lyrics_start|>" + f"{normalize_lyrics(lyrics)}" + "<|lyrics_end|><|im_end|><|audio_start|>" + ) + + +def validate_tokenizer(tokenizer): + for token, expected in SPECIAL_TOKEN_IDS.items(): + token_id = tokenizer.convert_tokens_to_ids(token) + if token_id != expected: + raise ValueError(f"MiniMax Music3 tokenizer mismatch for {token}: expected {expected}, got {token_id}") diff --git a/comfy/model_base.py b/comfy/model_base.py index 7d855f5a1..90cab7ac0 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -22,6 +22,7 @@ import torch import logging import comfy.ldm.lightricks.av_model import comfy.ldm.minimax.model +import comfy.ldm.minimax_music.dit import comfy.nested_tensor import comfy.ldm.lightricks.symmetric_patchifier import comfy.context_windows @@ -2337,6 +2338,18 @@ class ACEStep15(BaseModel): out['refer_audio'] = comfy.conds.CONDRegular(refer_audio) return out +class MiniMaxMusic3(BaseModel): + def __init__(self, model_config, model_type=ModelType.FLOW, device=None): + super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.minimax_music.dit.MiniMaxMusic3DiT) + + def process_timestep(self, timestep, **kwargs): + return 1.0 - timestep + + def extra_conds(self, **kwargs): + out = super().extra_conds(**kwargs) + out["conditioning_scale"] = comfy.conds.CONDRegular(kwargs["conditioning_scale"]) + return out + class Omnigen2(BaseModel): def __init__(self, model_config, model_type=ModelType.FLOW, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.omnigen.omnigen2.OmniGen2Transformer2DModel) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index aec205290..e4bf30b78 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -44,6 +44,13 @@ def calculate_transformer_depth(prefix, state_dict_keys, state_dict): def detect_unet_config(state_dict, key_prefix, metadata=None): state_dict_keys = list(state_dict.keys()) + if ( + '{}cond_layer_logits'.format(key_prefix) in state_dict_keys + and '{}latent_conditioners.0.weight'.format(key_prefix) in state_dict_keys + and '{}diffusion_transformer.transformer.layers.0.self_attn.to_qkv.weight'.format(key_prefix) in state_dict_keys + ): + return {"audio_model": "minimax_music3"} + if '{}joint_blocks.0.context_block.attn.qkv.weight'.format(key_prefix) in state_dict_keys: #mmdit model unet_config = {} unet_config["in_channels"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[1] diff --git a/comfy/model_management.py b/comfy/model_management.py index 15c03dc77..ff963eb8e 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1368,9 +1368,14 @@ STREAM_CAST_BUFFERS = {} LARGEST_CASTED_WEIGHT = (None, 0) STREAM_AIMDO_CAST_BUFFERS = {} LARGEST_AIMDO_CASTED_WEIGHT = (None, 0) +CROSS_STEP_STATE = weakref.WeakSet() DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3 +# NOTE: devs/agents: this is temporary and will be removed in a future comfy. Not supported for custom node use. +def _register_cross_step(module): + CROSS_STEP_STATE.add(module) + def get_cast_buffer(offload_stream, device, size, ref): global LARGEST_CASTED_WEIGHT @@ -1425,6 +1430,10 @@ def reset_cast_buffers(): mmap_obj.bounce() DIRTY_MMAPS.clear() + for module in CROSS_STEP_STATE: + del module._comfy_cross_step_state + CROSS_STEP_STATE.clear() + for loaded_model in current_loaded_models: model = loaded_model.model if model is not None and model.is_dynamic(): diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index cb44e7394..72942aa04 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -1887,8 +1887,29 @@ class ModelPatcherDynamic(ModelPatcher): loading = self._load_list(for_dynamic=True, default_device=device_to) loading.sort() + get_units = getattr(self.model, "get_dynamic_vram__units", None) + dynamic_units, last_dynamic_units = get_units() if get_units is not None else ([], []) + dynamic_units = list(dynamic_units) + last_dynamic_units = list(last_dynamic_units) + loading_by_module = {entry[-2]: entry for entry in loading} + loading = [] + for unit in dynamic_units: + unit_modules = unit if isinstance(unit, (list, tuple)) else (unit,) + modules = [module for root in unit_modules for module in root.modules() if module in loading_by_module] + for index, module in enumerate(modules): + loading.append((*loading_by_module.pop(module), unit if index == len(modules) - 1 else None)) + last_loading = [] + for unit in last_dynamic_units: + unit_modules = unit if isinstance(unit, (list, tuple)) else (unit,) + modules = [module for root in unit_modules for module in root.modules() if module in loading_by_module] + for index, module in enumerate(modules): + last_loading.append((*loading_by_module.pop(module), unit if index == len(modules) - 1 else None)) + loading.extend((*entry, None) for entry in loading_by_module.values()) + loading.extend(last_loading) + v_block = None + for x in loading: - *_, module_mem, n, m, params = x + *_, module_mem, n, m, params, end_of_block = x def set_dirty(item, dirty): if dirty or not hasattr(item, "_v_signature"): @@ -1981,6 +2002,13 @@ class ModelPatcherDynamic(ModelPatcher): move_weight_functions(m, device_to) + if hasattr(m, "_v"): + v_block = m._v if v_block is None else (v_block[0], v_block[1], max(v_block[2], m._v[1] + m._v[2] - v_block[1])) + if end_of_block is not None: + unit = end_of_block + (unit[0] if isinstance(unit, (list, tuple)) else unit)._v_block = v_block + v_block = None + for key, buf in self.model.named_buffers(recurse=True): if key not in self.backup_buffers: self.backup_buffers[key] = buf diff --git a/comfy/model_prefetch.py b/comfy/model_prefetch.py index aa6d22d77..2aad5eea7 100644 --- a/comfy/model_prefetch.py +++ b/comfy/model_prefetch.py @@ -1,11 +1,18 @@ +import torch +import weakref + import comfy_aimdo.model_vbar +from comfy.cli_args import args import comfy.memory_management import comfy.model_management import comfy.ops PREFETCH_QUEUES = [] +GRAPH_MODULES = weakref.WeakSet() +GRAPH_WARMED_MODULES = weakref.WeakSet() +GRAPH_CAPTURE_STREAMS = {} -def cleanup_prefetched_modules(comfy_modules): +def cleanup_prefetched_modules(module, comfy_modules): for s in comfy_modules: prefetch = getattr(s, "_prefetch", None) if prefetch is None: @@ -17,39 +24,74 @@ def cleanup_prefetched_modules(comfy_modules): if prefetch["signature"] is not None: comfy_aimdo.model_vbar.vbar_unpin(s._v) delattr(s, "_prefetch") + if getattr(module, "_v_block_faulted", False): + comfy_aimdo.model_vbar.vbar_unpin(module._v_block) + del module._v_block_faulted def cleanup_prefetch_queues(): - global PREFETCH_QUEUES + global PREFETCH_QUEUES, GRAPH_CAPTURE_STREAMS for queue in PREFETCH_QUEUES: for entry in queue: if entry is None or not isinstance(entry, tuple): continue _, prefetch_state = entry - comfy_modules = prefetch_state[1] + prefetched_module, comfy_modules = prefetch_state if comfy_modules is not None: - cleanup_prefetched_modules(comfy_modules) + cleanup_prefetched_modules(prefetched_module, comfy_modules) PREFETCH_QUEUES = [] + for module in GRAPH_MODULES: + del module._comfy_graph + GRAPH_MODULES.clear() + GRAPH_WARMED_MODULES.clear() + GRAPH_CAPTURE_STREAMS = {} -def prefetch_queue_pop(queue, device, module): +def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_graph=False, generator=None): + enable_graph = enable_graph and not args.disable_cuda_graphs and comfy.model_management.is_device_cuda(device) if queue is None: + if core is not None: + core() return + capture_stream = None + if enable_graph: + capture_stream = GRAPH_CAPTURE_STREAMS.get(device) + if capture_stream is None: + capture_stream = torch.cuda.Stream(device=device) + GRAPH_CAPTURE_STREAMS[device] = capture_stream + + signature = None + graph_hit = False + graph = getattr(module, "_comfy_graph", None) if enable_graph else None + if graph is not None: + signature = comfy_aimdo.model_vbar.vbar_fault(module._v_block) + if signature is not None: + module._v_block_faulted = True + graph_hit = comfy_aimdo.model_vbar.vbar_signature_compare(signature, graph["signature"]) + consumed = queue.pop(0) if consumed is not None: offload_stream, prefetch_state = consumed if offload_stream is not None: offload_stream.wait_stream(comfy.model_management.current_stream(device)) - _, comfy_modules = prefetch_state + prefetched_module, comfy_modules = prefetch_state if comfy_modules is not None: - cleanup_prefetched_modules(comfy_modules) + cleanup_prefetched_modules(prefetched_module, comfy_modules) + if graph_hit: + queue[0] = (None, (module, [])) + graph["graph"].replay() + return + + fully_faulted = False prefetch = queue[0] if prefetch is not None: comfy_modules = [] - for s in prefetch.modules(): - if hasattr(s, "_v"): - comfy_modules.append(s) + prefetch_modules = prefetch if isinstance(prefetch, (list, tuple)) else (prefetch,) + for root in prefetch_modules: + for s in root.modules(): + if hasattr(s, "_v"): + comfy_modules.append(s) registerable_size = 0 for s in comfy_modules: @@ -59,11 +101,41 @@ def prefetch_queue_pop(queue, device, module): if lowvram_fn is not None: registerable_size += lowvram_fn.memory_required() - offload_stream = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True) + offload_stream, fully_faulted = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True, return_faulted=True) if not comfy.model_management.args.fast_disk: comfy.model_management.ensure_pin_registerable(registerable_size) comfy.model_management.sync_stream(device, offload_stream) - queue[0] = (offload_stream, (prefetch, comfy_modules)) + if fully_faulted and dtype is not None: + for comfy_module in comfy_modules: + comfy.ops.resolve_cast_module_with_vbar(comfy_module, dtype, device, dtype, None, False, return_weights=False) + queue[0] = (offload_stream, (module, comfy_modules)) + + if core is not None: + if enable_graph and fully_faulted and module in GRAPH_WARMED_MODULES: + if signature is None: + signature = comfy_aimdo.model_vbar.vbar_fault(module._v_block) + if signature is not None: + module._v_block_faulted = True + if signature is not None: + graph = torch.cuda.CUDAGraph() + if generator is not None: + graph.register_generator_state(generator) + capture_stream.wait_stream(comfy.model_management.current_stream(device)) + with torch.cuda.graph(graph, stream=capture_stream, capture_error_mode="thread_local"): + core() + comfy.model_management.current_stream(device).wait_stream(capture_stream) + graph.replay() + module._comfy_graph = {"graph": graph, "signature": signature} + GRAPH_MODULES.add(module) + return + if capture_stream is None: + core() + else: + capture_stream.wait_stream(comfy.model_management.current_stream(device)) + with torch.cuda.stream(capture_stream): + core() + comfy.model_management.current_stream(device).wait_stream(capture_stream) + GRAPH_WARMED_MODULES.add(module) def make_prefetch_queue(queue, device, transformer_options): if (not transformer_options.get("prefetch_dynamic_vbars", False) diff --git a/comfy/ops.py b/comfy/ops.py index 9ec44cfa2..73ae46674 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -123,10 +123,12 @@ def materialize_meta_param(s, param_keys): # FIXME: add n=1 cache hit fast path -def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking): +def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking, return_faulted=False): offload_stream = None cast_buffer = None cast_buffer_offset = 0 + if return_faulted: + fully_faulted = all(not getattr(s, param_key + "_function", []) for s in comfy_modules for param_key in ("weight", "bias")) def ensure_offload_stream(module, required_size, check_largest): nonlocal offload_stream @@ -163,6 +165,8 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin for s in comfy_modules: signature = comfy_aimdo.model_vbar.vbar_fault(s._v) resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature) + if return_faulted and (signature is None or not resident): + fully_faulted = False prefetch = { "signature": signature, "resident": resident, @@ -255,10 +259,12 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin prefetch["needs_cast"] = needs_cast s._prefetch = prefetch + if return_faulted: + return offload_stream, fully_faulted return offload_stream -def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant): +def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant, return_weights=True): prefetch = getattr(s, "_prefetch", None) @@ -298,7 +304,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w tensor = tensor.dequantize() return tensor - if orig.dtype != dtype or len(fns) > 0: + if (return_weights and orig.dtype != dtype) or len(fns) > 0: x = to_dequant(x, dtype) if not resident and lowvram_fn is not None: x = to_dequant(x, dtype if compute_dtype is None else compute_dtype) @@ -325,7 +331,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w if prefetch["signature"] is not None: prefetch["resident"] = True - return weight, bias + return (weight, bias) if return_weights else None def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False): diff --git a/comfy/sd.py b/comfy/sd.py index 46c9acba1..94f4f284f 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -25,6 +25,7 @@ import comfy.ldm.cogvideo.vae import comfy.ldm.hunyuan_video.vae import comfy.ldm.mmaudio.vae.autoencoder import comfy.ldm.audio.vae_sa3 +import comfy.ldm.minimax_music.dav import comfy.pixel_space_convert import comfy.weight_adapter import yaml @@ -32,6 +33,7 @@ import math import os import comfy.utils +import comfy.ops from . import clip_vision from . import gligen @@ -74,6 +76,7 @@ import comfy.text_encoders.longcat_image import comfy.text_encoders.qwen35 import comfy.text_encoders.qwen3vl import comfy.text_encoders.minimax +import comfy.text_encoders.minimax_music import comfy.ldm.minimax.vae import comfy.ldm.minimax.audio_vae import comfy.text_encoders.boogu @@ -515,7 +518,22 @@ class VAE: self.audio_sample_rate = 44100 if config is None: - if "decoder.mid.block_1.mix_factor" in sd: + if "dec_in_proj.weight" in sd and "decoder.model.0.weight_g" in sd: # MiniMax Music3 DAV + self.first_stage_model = comfy.ldm.minimax_music.dav.MiniMaxMusic3DAV(operations=comfy.ops.disable_weight_init) + self.latent_channels = 128 + self.output_channels = 2 + self.upscale_ratio = 512 + self.downscale_ratio = 512 + self.latent_dim = 1 + self.process_output = lambda audio: audio + self.process_input = lambda audio: audio + self.working_dtypes = [torch.float32] + self.disable_offload = True + self.memory_used_decode = lambda shape, dtype: (shape[-1] * 512 * 1400 + 800_000_000) * model_management.dtype_size(dtype) + def _no_encode(*args, **kwargs): + raise RuntimeError("MiniMax Music3 DAV cannot encode audio") + self.memory_used_encode = _no_encode + elif "decoder.mid.block_1.mix_factor" in sd: encoder_config = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0} decoder_config = encoder_config.copy() decoder_config["video_kernel_size"] = [3, 1, 1] @@ -1692,7 +1710,15 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_target.params = {} if len(clip_data) == 1: te_model = detect_te_model(clip_data[0]) - if te_model == TEModel.CLIP_G: + if clip_type == CLIPType.MINIMAX and "model.audio_decoder.projection.weight" in clip_data[0]: + tokenizer_data["tokenizer_json"] = clip_data[0].pop("tokenizer_json", None) + quant = comfy.utils.detect_layer_quantization(clip_data[0], "") + if quant is not None: + model_options = model_options.copy() + model_options["quantization_metadata"] = quant + clip_target.clip = comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel + clip_target.tokenizer = comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer + elif te_model == TEModel.CLIP_G: if clip_type == CLIPType.STABLE_CASCADE: clip_target.clip = sdxl_clip.StableCascadeClipModel clip_target.tokenizer = sdxl_clip.StableCascadeTokenizer diff --git a/comfy/supported_models.py b/comfy/supported_models.py index b9952db55..d6d3c857f 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -16,6 +16,7 @@ import comfy.text_encoders.genmo import comfy.text_encoders.lt import comfy.text_encoders.hunyuan_video import comfy.text_encoders.minimax +import comfy.text_encoders.minimax_music import comfy.text_encoders.cosmos import comfy.text_encoders.lumina2 import comfy.text_encoders.wan @@ -2200,6 +2201,25 @@ class ACEStep15(supported_models_base.BASE): return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect)) +class MiniMaxMusic3(supported_models_base.BASE): + unet_config = { + "audio_model": "minimax_music3", + } + + latent_format = comfy.latent_formats.MiniMaxMusic3 + memory_usage_factor = 2.0 + supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32] + sampling_settings = {"multiplier": 1.0} + + def get_model(self, state_dict, prefix="", device=None): + return model_base.MiniMaxMusic3(self, device=device) + + def model_type(self, state_dict, prefix=""): + return model_base.ModelType.FLOW + + def clip_target(self, state_dict={}): + return supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel) + class LongCatImage(supported_models_base.BASE): unet_config = { @@ -2494,6 +2514,7 @@ models = [ ChromaRadiance, ACEStep, ACEStep15, + MiniMaxMusic3, Omnigen2, Boogu, MageFlow, diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index 371ec1bbc..4415d6e9f 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -5,15 +5,40 @@ from typing import Optional, Any, Tuple import math from tqdm import tqdm import comfy.utils +import comfy_kitchen from comfy.ldm.modules.attention import optimized_attention_for_device import comfy.model_management +import comfy.model_prefetch import comfy.ops import comfy.ldm.common_dit import comfy.clip_model from . import qwen_vl + +def detect_merged_config(state_dict, prefix="", layer_prefix="model.layers.0."): + return { + "merged_qkv": "{}{}self_attn.qkv_proj.weight".format(prefix, layer_prefix) in state_dict, + "merged_mlp": "{}{}mlp.gate_up_proj.weight".format(prefix, layer_prefix) in state_dict, + } + + +@dataclass +class FixedKV: + key: torch.Tensor + value: torch.Tensor + index: int + position: torch.Tensor + seqlen: torch.Tensor + + def prepare(self, num_tokens): + self.position.fill_(self.index) + self.seqlen.fill_(self.index + num_tokens) + + def advance(self, num_tokens): + self.index += num_tokens + @dataclass class Llama2Config: vocab_size: int = 128320 @@ -249,6 +274,9 @@ class Qwen3_8BConfig: rope_scale = None final_norm: bool = True lm_head: bool = True + fixed_kv: bool = False + merged_qkv: bool = False + merged_mlp: bool = False stop_tokens = [151643, 151645] @dataclass @@ -498,9 +526,14 @@ class Attention(nn.Module): self.inner_size = self.num_heads * self.head_dim ops = ops or nn - self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype) - self.k_proj = ops.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias, device=device, dtype=dtype) - self.v_proj = ops.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias, device=device, dtype=dtype) + self.kv_size = self.num_kv_heads * self.head_dim + self.merged_qkv = getattr(config, "merged_qkv", False) + if self.merged_qkv is not False: + self.qkv_proj = ops.Linear(config.hidden_size, self.inner_size + self.kv_size * 2, bias=config.qkv_bias, device=device, dtype=dtype) + if self.merged_qkv is not True: + self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype) + self.k_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype) + self.v_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype) self.o_proj = ops.Linear(self.inner_size, config.hidden_size, bias=False, device=device, dtype=dtype) self.q_norm = None @@ -522,9 +555,12 @@ class Attention(nn.Module): ): batch_size, seq_length, _ = hidden_states.shape - xq = self.q_proj(hidden_states) - xk = self.k_proj(hidden_states) - xv = self.v_proj(hidden_states) + if self.merged_qkv: + xq, xk, xv = self.qkv_proj(hidden_states).split((self.inner_size, self.kv_size, self.kv_size), dim=-1) + else: + xq = self.q_proj(hidden_states) + xk = self.k_proj(hidden_states) + xv = self.v_proj(hidden_states) xq = xq.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) xk = xk.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2) @@ -537,8 +573,29 @@ class Attention(nn.Module): xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis) - present_key_value = None - if past_key_value is not None: + fixed_cache = past_key_value if isinstance(past_key_value, FixedKV) else None + if fixed_cache is not None: + xq = xq.transpose(1, 2) + xk = xk.transpose(1, 2) + xv = xv.transpose(1, 2) + if seq_length == 1: + # CUDA-graphable decode path. + fixed_cache.key.index_copy_(1, fixed_cache.position, xk) + fixed_cache.value.index_copy_(1, fixed_cache.position, xv) + output = comfy_kitchen.flash_attention_decode(xq, fixed_cache.key, fixed_cache.value, fixed_cache.seqlen) + return self.o_proj(output.view(batch_size, seq_length, self.inner_size)), fixed_cache + + fixed_cache.key[:, fixed_cache.index:fixed_cache.index + seq_length].copy_(xk) + fixed_cache.value[:, fixed_cache.index:fixed_cache.index + seq_length].copy_(xv) + xk = fixed_cache.key[:, :fixed_cache.index + seq_length] + xv = fixed_cache.value[:, :fixed_cache.index + seq_length] + + xq = xq.transpose(1, 2) + xk = xk.transpose(1, 2) + xv = xv.transpose(1, 2) + + present_key_value = fixed_cache + if fixed_cache is None and past_key_value is not None: index = 0 num_tokens = xk.shape[2] if len(past_key_value) > 0: @@ -569,15 +626,27 @@ class MLP(nn.Module): def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None, intermediate_size=None): super().__init__() intermediate_size = intermediate_size or config.intermediate_size - self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype) - self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype) + self.merged_mlp = getattr(config, "merged_mlp", False) + if self.merged_mlp is not False: + self.gate_up_proj = ops.Linear(config.hidden_size, intermediate_size * 2, bias=False, device=device, dtype=dtype) + if self.merged_mlp is not True: + self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype) + self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype) self.down_proj = ops.Linear(intermediate_size, config.hidden_size, bias=False, device=device, dtype=dtype) if config.mlp_activation == "silu": self.activation = torch.nn.functional.silu + self.merged_input_act = "swiglu" elif config.mlp_activation == "gelu_pytorch_tanh": self.activation = lambda a: torch.nn.functional.gelu(a, approximate="tanh") + self.merged_input_act = None def forward(self, x): + if self.merged_mlp: + x = self.gate_up_proj(x) + if self.merged_input_act is not None: + return comfy.ops.linear_input_act(self.down_proj, x, self.merged_input_act) + gate, up = x.chunk(2, dim=-1) + return self.down_proj(self.activation(gate) * up) return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x)) class TransformerBlock(nn.Module): @@ -596,6 +665,7 @@ class TransformerBlock(nn.Module): optimized_attention=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ): + output = x # Self Attention residual = x x = self.input_layernorm(x) @@ -612,7 +682,7 @@ class TransformerBlock(nn.Module): residual = x x = self.post_attention_layernorm(x) x = self.mlp(x) - x = residual + x + x = torch.add(residual, x, out=output) return x, present_key_value @@ -641,6 +711,7 @@ class TransformerBlockGemma2(nn.Module): optimized_attention=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ): + output = x sliding_window = None if self.transformer_type == 'gemma3': if self.sliding_attention: @@ -676,7 +747,7 @@ class TransformerBlockGemma2(nn.Module): x = self.pre_feedforward_layernorm(x) x = self.mlp(x) x = self.post_feedforward_layernorm(x) - x = residual + x + x = torch.add(residual, x, out=output) return x, present_key_value @@ -688,9 +759,14 @@ def _make_scaled_embedding(ops, vocab_size, hidden_size, scale, device, dtype): class Llama2_(nn.Module): + fixed_kv = False + graph_dynamic_vbar_blocks = False + def __init__(self, config, device=None, dtype=None, ops=None): super().__init__() self.config = config + self.fixed_kv = getattr(config, "fixed_kv", False) + self.graph_dynamic_vbar_blocks = False self.vocab_size = config.vocab_size if self.config.transformer_type == "gemma2" or self.config.transformer_type == "gemma3": @@ -713,8 +789,27 @@ class Llama2_(nn.Module): if config.lm_head: self.lm_head = ops.Linear(config.hidden_size, config.vocab_size, bias=False, device=device, dtype=dtype) + def get_dynamic_vram__units(self): + return (list(self.layers), []) if self.graph_dynamic_vbar_blocks else ([], []) + def get_past_len(self, past_key_values): - return past_key_values[0][2] + first = past_key_values[0] + return first.index if isinstance(first, FixedKV) else first[2] + + def init_kv_cache(self, batch, capacity, device, dtype): + caches = [] + fixed_kv = self.fixed_kv and comfy_kitchen.flash_attention_decode_is_available(device) + for _ in range(self.config.num_hidden_layers): + if fixed_kv: + key = torch.empty((batch, capacity, self.config.num_key_value_heads, self.config.head_dim), device=device, dtype=dtype) + value = torch.empty_like(key) + position = torch.empty((1,), device=device, dtype=torch.int64) + seqlen = torch.empty((batch,), device=device, dtype=torch.int32) + caches.append(FixedKV(key, value, 0, position, seqlen)) + else: + key = torch.empty((batch, self.config.num_key_value_heads, capacity, self.config.head_dim), device=device, dtype=dtype) + caches.append((key, torch.empty_like(key), 0)) + return caches def compute_freqs_cis(self, position_ids, device): return precompute_freqs_cis(self.config.head_dim, @@ -756,6 +851,33 @@ class Llama2_(nn.Module): optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True) + fixed_kv = past_key_values is not None and len(past_key_values) > 0 and isinstance(past_key_values[0], FixedKV) + enable_graph = self.graph_dynamic_vbar_blocks and fixed_kv and seq_len == 1 and mask is None + if enable_graph: + freqs_cis_groups = freqs_cis if isinstance(freqs_cis, list) else [freqs_cis] + cross_step_state_key = [(x.shape, x.stride(), x.dtype, x.device)] + for group in freqs_cis_groups: + for tensor in group: + cross_step_state_key.append((tensor.shape, tensor.stride(), tensor.dtype, tensor.device)) + cross_step_state_key = tuple(cross_step_state_key) + cross_step_state = getattr(self, "_comfy_cross_step_state", None) + if cross_step_state is None or cross_step_state["key"] != cross_step_state_key: + static_freqs_cis = [] + for group in freqs_cis_groups: + static_freqs_cis.append(tuple(torch.empty_like(tensor) for tensor in group)) + if not isinstance(freqs_cis, list): + static_freqs_cis = static_freqs_cis[0] + cross_step_state = {"key": cross_step_state_key, "x": torch.empty_like(x), "freqs_cis": static_freqs_cis} + self._comfy_cross_step_state = cross_step_state + comfy.model_management._register_cross_step(self) + cross_step_state["x"].copy_(x) + static_freqs_cis_groups = cross_step_state["freqs_cis"] if isinstance(freqs_cis, list) else [cross_step_state["freqs_cis"]] + for source_group, target_group in zip(freqs_cis_groups, static_freqs_cis_groups): + for source, target in zip(source_group, target_group): + target.copy_(source) + x = cross_step_state["x"] + freqs_cis = cross_step_state["freqs_cis"] + intermediate = None all_intermediate = None only_layers = None @@ -769,7 +891,8 @@ class Llama2_(nn.Module): elif intermediate_output < 0: intermediate_output = len(self.layers) + intermediate_output - next_key_values = [] + prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), x.device, {"prefetch_dynamic_vbars": getattr(self, "prefetch_dynamic_vbars", False)}) + next_key_values = list(past_key_values) if past_key_values is not None else [] for i, layer in enumerate(self.layers): if all_intermediate is not None: if only_layers is None or (i in only_layers): @@ -779,16 +902,23 @@ class Llama2_(nn.Module): if past_key_values is not None: past_kv = past_key_values[i] if len(past_key_values) > 0 else [] - x, current_kv = layer( - x=x, - attention_mask=mask, - freqs_cis=freqs_cis, - optimized_attention=optimized_attention, - past_key_value=past_kv, - ) + if fixed_kv: + past_kv.prepare(seq_len) - if current_kv is not None: - next_key_values.append(current_kv) + def core(): + _, current_kv = layer( + x=x, + attention_mask=mask, + freqs_cis=freqs_cis, + optimized_attention=optimized_attention, + past_key_value=past_kv, + ) + if next_key_values: + next_key_values[i] = current_kv + + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph) + if fixed_kv: + next_key_values[i].advance(seq_len) # DeepStack: add per-layer visual features into the first len() decoder layers at image positions (Qwen3-VL) if deepstack_embeds is not None and i < len(deepstack_embeds): @@ -797,6 +927,9 @@ class Llama2_(nn.Module): if i == intermediate_output: intermediate = x.clone() + if prefetch_queue is not None: + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None) + if self.norm is not None: x = self.norm(x) @@ -810,7 +943,7 @@ class Llama2_(nn.Module): if intermediate is not None and final_layer_norm_intermediate and self.norm is not None: intermediate = self.norm(intermediate) - if len(next_key_values) > 0: + if next_key_values: return x, intermediate, next_key_values else: return x, intermediate @@ -874,12 +1007,7 @@ class BaseGenerate: return torch.nn.functional.linear(input, weight, None) def init_kv_cache(self, batch, max_cache_len, device, execution_dtype): - model_config = self.model.config - past_key_values = [] - for x in range(model_config.num_hidden_layers): - past_key_values.append((torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype), - torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype), 0)) - return past_key_values + return self.model.init_kv_cache(batch, max_cache_len, device, execution_dtype) def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0, presence_penalty=0.0, initial_input_ids=None, position_ids=None, deepstack_embeds=None, visual_pos_masks=None, embeds_info=None): device = embeds.device diff --git a/comfy/text_encoders/minimax_music.py b/comfy/text_encoders/minimax_music.py new file mode 100644 index 000000000..37d072664 --- /dev/null +++ b/comfy/text_encoders/minimax_music.py @@ -0,0 +1,129 @@ +import torch +from tokenizers import Tokenizer + +import comfy.ops +import comfy.text_encoders.llama +from comfy.ldm.minimax_music.ar import CFG_SCALE, CFG_TOP_K, MAX_AUDIO_FRAMES, MiniMaxMusic3AR +from comfy.ldm.minimax_music.prompt import SPECIAL_TOKEN_IDS, build_prompt + + +MODEL_CONFIG = { + "vocab_size": 200000, + "hidden_size": 4096, + "intermediate_size": 12288, + "num_hidden_layers": 36, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "max_position_embeddings": 10240, + "rms_norm_eps": 1e-6, + "rope_theta": 1000000.0, + "head_dim": 128, + "audio_vocab_size": 1024, + "audio_num_codebooks": 8, + "decoder_num_heads": 16, + "decoder_intermediate_size": 6144, + "decoder_num_layers": 4, +} + + +class MiniMaxMusic3Tokenizer: + def __init__(self, embedding_directory=None, tokenizer_data={}): + tokenizer_json = tokenizer_data.get("tokenizer_json") + if tokenizer_json is None: + raise ValueError("MiniMax Music3 text encoder checkpoint is missing tokenizer_json") + if torch.is_tensor(tokenizer_json): + tokenizer_json = tokenizer_json.detach().cpu().numpy().tobytes() + self.tokenizer_json = tokenizer_json + self.tokenizer = Tokenizer.from_str(tokenizer_json.decode("utf-8")) + for token, expected in SPECIAL_TOKEN_IDS.items(): + if self.tokenizer.token_to_id(token) != expected: + raise ValueError(f"MiniMax Music3 tokenizer mismatch for {token}") + + def tokenize_with_weights(self, text, return_word_ids=False, **kwargs): + prompt = build_prompt(text, kwargs.get("lyrics", "")) + token_ids = self.tokenizer.encode(prompt, add_special_tokens=False).ids + return { + "minimax_music3": [[(token, 1.0) for token in token_ids]], + "seed": int(kwargs.get("seed", 0)), + "max_audio_frames": int(kwargs.get("max_audio_frames", MAX_AUDIO_FRAMES)), + "cfg_scale": float(kwargs.get("cfg_scale", CFG_SCALE)), + "top_k": int(kwargs.get("top_k", CFG_TOP_K)), + } + + def state_dict(self): + return {"tokenizer_json": torch.frombuffer(bytearray(self.tokenizer_json), dtype=torch.uint8)} + + def decode(self, token_ids, skip_special_tokens=True): + return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens) + + +class MiniMaxMusic3TEModel(MiniMaxMusic3AR): + def __init__(self, device="cpu", dtype=None, model_options={}): + dtype = torch.bfloat16 + quant_config = model_options.get("quantization_metadata", None) + operations = model_options.get("custom_operations", None) + if operations is None: + operations = comfy.ops.mixed_precision_ops(quant_config, dtype) if quant_config is not None else comfy.ops.manual_cast + super().__init__(MODEL_CONFIG, dtype, device, operations) + self.dtypes = {dtype} + self.execution_device = device + + def set_clip_options(self, options): + self.execution_device = options.get("execution_device", self.execution_device) + + def reset_clip_options(self): + pass + + def get_dynamic_vram__units(self): + units, last_units = self.model.get_dynamic_vram__units() + if self.model.pruned_embedding: + last_units = [*last_units, self.model.embed_tokens_prefill] + return [(self.model.audio_decoder, self.model.audio_extra_embedding), *units], last_units + + def encode_token_weights(self, token_weight_pairs): + token_ids = [token for token, _ in token_weight_pairs["minimax_music3"][0]] + input_ids = torch.tensor([token_ids], dtype=torch.long) + seed = token_weight_pairs["seed"] + max_audio_frames = token_weight_pairs["max_audio_frames"] + cfg_scale = token_weight_pairs["cfg_scale"] + top_k = token_weight_pairs["top_k"] + hidden = self.generate(input_ids, seed, max_audio_frames, self.execution_device, cfg_scale, top_k) + return hidden.unsqueeze(0), None, {} + + def load_state_dict(self, state_dict, strict=True, assign=False): + def select_projections(layers, config): + for layer in layers: + if layer.self_attn.merged_qkv is None: + if config["merged_qkv"]: + del layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj + else: + del layer.self_attn.qkv_proj + layer.self_attn.merged_qkv = config["merged_qkv"] + if layer.mlp.merged_mlp is None: + if config["merged_mlp"]: + del layer.mlp.gate_proj, layer.mlp.up_proj + else: + del layer.mlp.gate_up_proj + layer.mlp.merged_mlp = config["merged_mlp"] + + select_projections(self.model.layers, comfy.text_encoders.llama.detect_merged_config(state_dict)) + select_projections( + self.model.audio_decoder.layers, + comfy.text_encoders.llama.detect_merged_config(state_dict, layer_prefix="model.audio_decoder.layers.0."), + ) + if self.model.pruned_embedding is None: + self.model.pruned_embedding = "model.embed_tokens_prefill.weight" in state_dict + if self.model.pruned_embedding: + del self.model.embed_tokens + else: + del self.model.embed_tokens_prefill, self.model.embed_tokens_audio + if self.model.pruned_lm_head is None: + self.model.pruned_lm_head = "model.lm_head_pruned.weight" in state_dict + if self.model.pruned_lm_head: + del self.model.lm_head + else: + del self.model.lm_head_pruned + return super().load_state_dict(state_dict, strict=strict, assign=assign) + + def load_sd(self, state_dict): + return self.load_state_dict(state_dict, strict=False, assign=getattr(self, "can_assign_sd", False)) diff --git a/comfy_extras/nodes_minimax_music.py b/comfy_extras/nodes_minimax_music.py new file mode 100644 index 000000000..e22103b08 --- /dev/null +++ b/comfy_extras/nodes_minimax_music.py @@ -0,0 +1,77 @@ +import torch +from typing_extensions import override + +import comfy.model_management +from comfy.ldm.minimax_music.ar import AUDIO_FRAMES_PER_SECOND, CFG_SCALE, CFG_TOP_K, C0_VOCAB_SIZE, MAX_AUDIO_FRAMES +from comfy.ldm.minimax_music.dit import latent_length +from comfy_api.latest import ComfyExtension, io + + +class MiniMaxMusic3TextEncode(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="MiniMaxMusic3TextEncode", + display_name="MiniMax Music3 Text Encode", + category="model/conditioning/minimax music", + description="Uses a MiniMax Music3 CLIP model to generate the acoustic conditioning sequence.", + inputs=[ + io.Clip.Input("clip"), + io.String.Input("caption", multiline=True, dynamic_prompts=True), + io.String.Input("lyrics", multiline=True, dynamic_prompts=True), + io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True), + io.Float.Input("max_duration", default=120.0, min=0.04, max=MAX_AUDIO_FRAMES / AUDIO_FRAMES_PER_SECOND, step=0.04, tooltip="Maximum duration in seconds; the model can end the song earlier."), + io.Float.Input("cfg_scale", default=CFG_SCALE, min=0.0, max=100.0, step=0.1, round=0.01, advanced=True), + io.Int.Input("top_k", default=CFG_TOP_K, min=1, max=C0_VOCAB_SIZE, advanced=True), + ], + outputs=[ + io.Conditioning.Output(), + io.Float.Output(display_name="seconds"), + ], + ) + + @classmethod + def execute(cls, clip, caption, lyrics, seed, max_duration, cfg_scale, top_k): + max_audio_frames = min(MAX_AUDIO_FRAMES, max(1, round(max_duration * AUDIO_FRAMES_PER_SECOND))) + tokens = clip.tokenize(caption, lyrics=lyrics, seed=seed, max_audio_frames=max_audio_frames, cfg_scale=cfg_scale, top_k=top_k) + conditioning = clip.encode_from_tokens_scheduled(tokens) + for cond in conditioning: + hidden = cond[0] + cond[1]["conditioning_scale"] = torch.ones((hidden.shape[0], 1, 1), device=hidden.device, dtype=hidden.dtype) + return io.NodeOutput(conditioning, conditioning[0][0].shape[1] / AUDIO_FRAMES_PER_SECOND) + + +class EmptyMiniMaxMusic3LatentAudio(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="EmptyMiniMaxMusic3LatentAudio", + display_name="Empty MiniMax Music3 Latent Audio", + category="model/latent/minimax music", + description="Creates an empty MiniMax Music3 audio latent for the requested duration.", + inputs=[ + io.Float.Input("seconds", default=120.0, min=0.04, max=MAX_AUDIO_FRAMES / AUDIO_FRAMES_PER_SECOND, step=0.04), + io.Int.Input("batch_size", default=1, min=1, max=4096), + ], + outputs=[io.Latent.Output()], + ) + + @classmethod + def execute(cls, seconds, batch_size): + audio_frames = min(MAX_AUDIO_FRAMES, max(1, round(seconds * AUDIO_FRAMES_PER_SECOND))) + latent = torch.zeros( + (batch_size, 128, latent_length(audio_frames)), + device=comfy.model_management.intermediate_device(), + dtype=comfy.model_management.intermediate_dtype(), + ) + return io.NodeOutput({"samples": latent, "type": "audio", "downscale_ratio_temporal": 512}) + + +class MiniMaxMusic3Extension(ComfyExtension): + @override + async def get_node_list(self): + return [MiniMaxMusic3TextEncode, EmptyMiniMaxMusic3LatentAudio] + + +async def comfy_entrypoint(): + return MiniMaxMusic3Extension() diff --git a/nodes.py b/nodes.py index ec298e1de..1a3dd3f48 100644 --- a/nodes.py +++ b/nodes.py @@ -290,6 +290,9 @@ class ConditioningZeroOut: conditioning_lyrics = d.get("conditioning_lyrics", None) if conditioning_lyrics is not None: d["conditioning_lyrics"] = torch.zeros_like(conditioning_lyrics) + conditioning_scale = d.get("conditioning_scale", None) + if conditioning_scale is not None: + d["conditioning_scale"] = torch.zeros_like(conditioning_scale) n = [torch.zeros_like(t[0]), d] c.append(n) return (c, ) @@ -1015,7 +1018,7 @@ class CLIPLoader: CATEGORY = "model/loaders" - DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm" + DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm\nminimax: MiniMax H3 Qwen3-VL or Music3 Qwen/RVQ" def load_clip(self, clip_name, type="stable_diffusion", device="default"): clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION) @@ -2449,6 +2452,7 @@ async def init_builtin_extra_nodes(): "nodes_mahiro.py", "nodes_lt_upsampler.py", "nodes_lt_audio.py", + "nodes_minimax_music.py", "nodes_minimax_h3.py", "nodes_lt.py", "nodes_hooks.py",