Implement Minimax Music 3 + Core Support for Cuda Graphs (#15570)

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.")
 

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

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")

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

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

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"(?<!\*)\*([^*\n]+)\*(?!\*)", r"\1", line)
+        lines.append(line.rstrip())
+    text = "\n".join(lines)
+    text = re.sub(r"^\s*[-*_]{3,}\s*$", "", text, flags=re.MULTILINE)
+    return text.replace("• ", "").replace("    ", "")
+
+
+def clean_caption(caption):
+    def replace_special(match):
+        inner = match.group(1).strip()
+        parts = inner.split(None, 1)
+        return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else inner
+
+    text = _SPECIAL_TAG_RE.sub(replace_special, caption)
+    text = _remove_markdown_format(text)
+    return re.sub(r"\n{2,}", "\n", text)
+
+
+def normalize_lyrics(lyrics):
+    parts = _LYRIC_TAG_RE.split(lyrics)
+    text = "\n".join(part.lower() if part.startswith("[") else part for part in parts if part)
+    text = text.replace(" ^ ", "\n")
+    return f"[start]\n{text}"
+
+
+def build_prompt(caption, lyrics):
+    return (
+        "<|im_start|><|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}")

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)

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]

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():

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

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)

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):

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

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,

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

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))

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()

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",