Merge branch 'master' into automation/comfyui-frontend-bump

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/lightricks/vae/na_diffusion_decoder.py

@@ -22,6 +22,7 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange
 from torch import nn
+import comfy.model_management
 
 from comfy.ldm.lightricks.model import get_timestep_embedding
 from .causal_video_autoencoder import Encoder, processor
@@ -74,8 +75,12 @@ def default_rope_dim_split(head_dim):
 
 
 def rope_inv_freqs(dim, base=10000.0, device=None):
+    out_device = device
+    if not comfy.model_management.supports_fp64(device):
+        device = torch.device("cpu")
+
     exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim
-    return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(torch.float32)
+    return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(dtype=torch.float32, device=out_device)
 
 
 def _rope_tables(lengths, inv_freqs, device):

comfy/ldm/minimax/model.py

@@ -91,6 +91,18 @@ def _video_t_grid(n, origin):
     return float(origin) + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)])
 
 
+def _ref_t_span(blk):
+    # time-axis span a reference block occupies ahead of the target streams
+    kind = blk["kind"]
+    if kind == "image":
+        return 1.0
+    if kind == "audio":
+        return float(blk["ref_audio_t"])
+    if kind in ("video", "video_audio"):
+        return max(float(blk["ref_audio_t"]), sum(_video_t_spans(blk["latent_t"])))
+    return 0.0
+
+
 def _audio_grid(cursor, t, w_low, w_high):
     # channel-major stereo rows: t advances per latent frame, w pinned to the grid extremes per stereo channel, h stays 0
     g = torch.zeros(t * 2, 3, dtype=torch.float64)
@@ -288,7 +300,7 @@ class FinalLayer(nn.Module):
 class PackedLayout:
     """Static packed-sequence structure for one shape/conditioning signature."""
 
-    def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None, frame_count=None):
+    def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None):
         frame, w_grid = _frame_grid(latent_h, latent_w)
         frame_rows = frame.shape[0]
 
@@ -299,29 +311,37 @@ class PackedLayout:
 
         img_pos, img_update = [], []
         audio_pos, audio_update = [], []
-        cursor = text_len
         row = text_len
 
-        if keyframes:
-            # fl2va: keyframe cond rows right after text, sharing the target spatial grid
-            for kf in keyframes:
-                pixel_index = kf["resolved_frame_index"]
-                if pixel_index == 0:
-                    cond_t = float(text_len)
-                elif frame_count is not None and pixel_index == frame_count - 1:
-                    cond_t = float(text_len) + sum(_video_t_spans(latent_t)) - FRAME_RESCALE
-                else:
-                    raise ValueError("only first/last keyframe anchors are supported")
-                g = torch.empty(frame_rows, 3, dtype=torch.float64)
-                g[:, 0] = cond_t
-                g[:, 1:] = frame
-                segments.append(("cond", frame_rows))
-                pos.append(g)
-                img_pos.append(torch.arange(row, row + frame_rows))
-                img_update.append(torch.zeros(frame_rows, dtype=torch.bool))
-                row += frame_rows
-
         target_audio_w = (float(w_grid[0]), float(w_grid[-1]))
+        # refs pack between text and the targets, so the target timeline starts after their spans
+        cursor = float(text_len)
+        for blk in refs or ():
+            cursor += _ref_t_span(blk)
+
+        if keyframes:
+            # fl2va: keyframe cond rows right after text, sharing the target spatial grid;
+            # anchors count from the target timeline origin, FRAME_RESCALE per pixel frame, 1.0 per audio latent frame
+            for kf in keyframes:
+                cond_t = cursor + FRAME_RESCALE * kf["resolved_frame_index"]
+                video_latent = kf.get("latent")
+                if video_latent is not None:
+                    vt = video_latent.shape[2]
+                    n = vt * frame_rows
+                    segments.append(("cond", n))
+                    pos.append(_video_grid(vt, frame, cond_t))
+                    img_pos.append(torch.arange(row, row + n))
+                    img_update.append(torch.zeros(n, dtype=torch.bool))
+                    row += n
+                audio_latent = kf.get("audio_latent")
+                if audio_latent is not None:
+                    rt = audio_latent.shape[-1]
+                    segments.append(("cond_audio", rt * 2))
+                    pos.append(_audio_grid(cond_t, rt, *target_audio_w))
+                    audio_pos.append(torch.arange(row, row + rt * 2))
+                    audio_update.append(torch.zeros(rt * 2, dtype=torch.bool))
+                    row += rt * 2
+
         if refs:
             cursor = float(text_len)
             for blk in refs:
@@ -389,7 +409,7 @@ class PackedLayout:
         self.audio_update = torch.cat(audio_update)
         self.signature = (text_len, latent_t, latent_h, latent_w, audio_t)
         # contiguous segment table (start, stop, kind)
-        # kinds: text / cond / ref_img / ref_audio / audio / video
+        # kinds: text / cond / cond_audio / ref_img / ref_audio / audio / video
         # the packed sequence is uniform per segment in (modality tag, timestep class),
         # except the text span (tag runs resolved at forward time from the presentation tags)
         seg_abs = []
@@ -529,8 +549,7 @@ class MiniMaxH3Model(nn.Module):
         if layout is None or layout.signature != (text_len, latent_t, lat_h, lat_w, audio_t):
             layout = PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t,
                                   keyframes=payload.get("keyframes"),
-                                  refs=payload.get("refs"),
-                                  frame_count=payload.get("frame_count"))
+                                  refs=payload.get("refs"))
 
         # model_base passes model_sampling.timestep(sigma) = sigma * 1000
         shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
@@ -543,14 +562,14 @@ class MiniMaxH3Model(nn.Module):
         vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP))
         aud_aug = float(payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP))
         has_vis_cond = any(k in ("cond", "ref_img") for _, _, k in layout.segments)
-        has_aud_cond = any(k == "ref_audio" for _, _, k in layout.segments)
+        has_aud_cond = any(k in ("cond_audio", "ref_audio") for _, _, k in layout.segments)
         seg_t = {"text": t_v, "video": t_v, "audio": t_a,
                  "cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug),
-                 "ref_audio": max(t_a, aud_aug)}
+                 "cond_audio": max(t_a, aud_aug), "ref_audio": max(t_a, aud_aug)}
         unique_t = sorted({t_v, t_a} | ({seg_t["cond"]} if has_vis_cond else set())
                           | ({seg_t["ref_audio"]} if has_aud_cond else set()))
         t_row = {t: i for i, t in enumerate(unique_t)}
-        seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "ref_audio": 2}
+        seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "cond_audio": 2, "ref_audio": 2}
 
         text_tags = payload.get("text_token_tags")
         mod_segments = []

comfy/ldm/minimax_music/ar.py

@@ -0,0 +1,343 @@
+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, merged_qkv, dtype, device, operations):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.merged_qkv = merged_qkv
+        if merged_qkv:
+            self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device)
+        else:
+            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, merged_mlp, dtype, device, operations):
+        super().__init__()
+        self.merged_mlp = merged_mlp
+        if merged_mlp:
+            self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device)
+        else:
+            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, merged_qkv, merged_mlp, dtype, device, operations):
+        super().__init__()
+        self.input_layernorm = RVQRMSNorm(hidden_size, dtype, device)
+        self.self_attn = RVQAttention(hidden_size, num_heads, merged_qkv, dtype, device, operations)
+        self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device)
+        self.mlp = RVQMLP(hidden_size, intermediate_size, merged_mlp, 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"])
+        merged_qkv = config.get("decoder_merged_qkv", False)
+        merged_mlp = config.get("decoder_merged_mlp", False)
+        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"]),
+                merged_qkv,
+                merged_mlp,
+                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
+        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
@@ -2164,13 +2165,13 @@ class MiniMaxH3(BaseModel):
         keyframes = kwargs.get("minimax_keyframes", None)
         if keyframes is not None:
             payload["keyframes"] = keyframes
-            payload["frame_count"] = kwargs.get("minimax_frame_count", None)
-            payload["cond_video_latents"] = [kf["latent"] for kf in keyframes]
+            payload["cond_video_latents"] = [kf["latent"] for kf in keyframes if kf.get("latent") is not None]
+            payload["cond_audio_latents"] = [kf["audio_latent"] for kf in keyframes if kf.get("audio_latent") is not None]
         refs = kwargs.get("minimax_refs", None)
         if refs is not None:
             payload["refs"] = refs
-            payload["cond_video_latents"] = [r["latent"] for r in refs if "latent" in r]
-            payload["cond_audio_latents"] = [r["audio_latent"] for r in refs if r.get("audio_latent") is not None]
+            payload["cond_video_latents"] = payload.get("cond_video_latents", []) + [r["latent"] for r in refs if "latent" in r]
+            payload["cond_audio_latents"] = payload.get("cond_audio_latents", []) + [r["audio_latent"] for r in refs if r.get("audio_latent") is not None]
         if kwargs.get("minimax_visual_cond_noise_aug", None) is not None:
             payload["visual_cond_noise_aug"] = kwargs["minimax_visual_cond_noise_aug"]
         if kwargs.get("minimax_audio_cond_noise_aug", None) is not None:
@@ -2184,7 +2185,7 @@ class MiniMaxH3(BaseModel):
             payload["layout"] = comfy.ldm.minimax.model.PackedLayout(
                 cross_attn.shape[1], vs[2], (vs[3] + 1) // 2 * 2, (vs[4] + 1) // 2 * 2,
                 latent_shapes[1][-1], keyframes=payload.get("keyframes"),
-                refs=payload.get("refs"), frame_count=payload.get("frame_count"))
+                refs=payload.get("refs"))
         out['minimax_payload'] = comfy.conds.CONDConstant(payload)
         return out
 
@@ -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]
@@ -830,11 +837,10 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
 
         dit_config["use_adaln_lora"] = True
         dit_config["adaln_lora_dim"] = 256
+        dit_config["num_blocks"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
         if dit_config["model_channels"] == 2048:
-            dit_config["num_blocks"] = 28
             dit_config["num_heads"] = 16
         elif dit_config["model_channels"] == 5120:
-            dit_config["num_blocks"] = 36
             dit_config["num_heads"] = 40
 
         if dit_config["in_channels"] == 16:

comfy/model_management.py

@@ -490,28 +490,36 @@ try:
         except:
             rocm_version = (6, -1)
 
-        def aotriton_supported(gpu_arch):
-            path = torch.__path__[0]
-            path = os.path.join(os.path.join(path, "lib"), "aotriton.images")
-            gfx = set(map(lambda a: a[4:], filter(lambda a: a.startswith("amd-gfx"), os.listdir(path))))
-            if gpu_arch in gfx:
-                return True
-            if "{}x".format(gpu_arch[:-1]) in gfx:
-                return True
-            if "{}xx".format(gpu_arch[:-2]) in gfx:
-                return True
-            return False
+        def aotriton_supported():
+            """Whether pytorch reports flash attention as usable on this gpu.
+
+            can_use_flash_attention() evaluates runtime eligibility for the given
+            parameters; on a ROCm build that includes checking the gpu arch against the
+            kernel images AOTriton was compiled for. Querying it avoids assuming where
+            those images live inside the torch install. The probe tensor is shaped and
+            typed to pass the unrelated SDPA checks, so False means no hardware support
+            rather than a rejected shape.
+            """
+            try:
+                if not torch.backends.cuda.is_flash_attention_available():  # not built with flash attention
+                    return False
+                q = torch.empty((1, 1, 8, 64), dtype=torch.float16, device=get_torch_device())
+                params = torch.backends.cuda.SDPAParams(q, q, q, None, 0.0, False, False)
+                return torch.backends.cuda.can_use_flash_attention(params, False)
+            except (AttributeError, RuntimeError, TypeError) as e:
+                logging.warning("Could not query aotriton support: {}".format(e))
+                return False
 
         logging.info("AMD arch: {}".format(arch))
         logging.info("ROCm version: {}".format(rocm_version))
         if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
-            if aotriton_supported(arch):  # AMD efficient attention implementation depends on aotriton.
+            if aotriton_supported():  # AMD efficient attention implementation depends on aotriton.
                 if torch_version_numeric >= (2, 7):  # works on 2.6 but doesn't actually seem to improve much
                     if any((a in arch) for a in ["gfx90a", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1150", "gfx1151"]):  # TODO: more arches, TODO: gfx950
                         ENABLE_PYTORCH_ATTENTION = True
                 if rocm_version >= (7, 0):
-                   if any((a in arch) for a in ["gfx1200", "gfx1201"]):
-                       ENABLE_PYTORCH_ATTENTION = True
+                    if any((a in arch) for a in ["gfx1200", "gfx1201"]):
+                        ENABLE_PYTORCH_ATTENTION = True
         if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4):
             if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950"]):  # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0
                 SUPPORT_FP8_OPS = True
@@ -1360,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
 
@@ -1417,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) and getattr(module, "_v_block", None) is not None
     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,16 @@ 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.params["projection_config"] = comfy.text_encoders.minimax_music.detect_merged_config(clip_data[0])
+            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,28 @@ 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={}):
+        detect = comfy.text_encoders.minimax_music.detect_merged_config(state_dict, self.text_encoder_key_prefix[0])
+        target = supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel)
+        target.params["projection_config"] = detect
+        return target
+
 
 class LongCatImage(supported_models_base.BASE):
     unet_config = {
@@ -2494,6 +2517,7 @@ models = [
     ChromaRadiance,
     ACEStep,
     ACEStep15,
+    MiniMaxMusic3,
     Omnigen2,
     Boogu,
     MageFlow,

comfy/text_encoders/gemma4.py

@@ -1183,6 +1183,7 @@ def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, po
 
 class Gemma4_Tokenizer():
     tokenizer_json_data = None
+    prime_empty_thought = False
 
     def state_dict(self):
         if self.tokenizer_json_data is not None:
@@ -1333,8 +1334,8 @@ class Gemma4_Tokenizer():
                     num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1]
                     n_audio_tokens = self._audio_token_count(num_samples)
                     media += "<|audio>" + "<|audio|>" * n_audio_tokens + "<audio|>"
-                # Non-thinking mode primes an empty thought channel so the model answers directly.
-                model_open = "" if thinking else "<|channel>thought\n<channel|>"
+                # 12B/31B prime a closed thought block for non-thinking mode, E2B/E4B must not: it cues them into reasoning inline.
+                model_open = "<|channel>thought\n<channel|>" if self.prime_empty_thought and not thinking else ""
                 llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\n<|turn>model\n{model_open}"
 
         text_tokens = super().tokenize_with_weights(llama_text, return_word_ids)
@@ -1418,6 +1419,7 @@ class Gemma4Tokenizer(sd1_clip.SD1Tokenizer):
 class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer):
     """Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram."""
     embedding_size = 3840
+    prime_empty_thought = True
 
     def _extract_audio_features(self, waveform, sample_rate):
         audio = self._resample_16k(waveform, sample_rate)
@@ -1500,7 +1502,7 @@ def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=No
 
 # Variants
 
-def _make_variant(config_cls):
+def _make_variant(config_cls, prime_empty_thought=False):
     audio = config_cls.audio_config is not None
     bases = (Gemma4AudioMixin, Gemma4Base) if audio else (Gemma4Base,)
     class Variant(*bases):
@@ -1510,8 +1512,8 @@ def _make_variant(config_cls):
             if audio:
                 self._init_audio(self.model.config, dtype, device, operations)
     embedding_size = config_cls.hidden_size
-    if embedding_size != Gemma4SDTokenizer.embedding_size:
-        tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size})
+    if embedding_size != Gemma4SDTokenizer.embedding_size or prime_empty_thought:
+        tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size, 'prime_empty_thought': prime_empty_thought})
         class Tokenizer(Gemma4Tokenizer):
             tokenizer_class = tok_cls
         Variant.tokenizer = Tokenizer
@@ -1521,7 +1523,7 @@ def _make_variant(config_cls):
 
 Gemma4_E4B = _make_variant(Gemma4Config)
 Gemma4_E2B = _make_variant(Gemma4_E2B_Config)
-Gemma4_31B = _make_variant(Gemma4_31B_Config)
+Gemma4_31B = _make_variant(Gemma4_31B_Config, prime_empty_thought=True)
 
 
 # Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant).

comfy/text_encoders/llama.py

@@ -5,15 +5,33 @@ 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
 
+
+@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 +267,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 +519,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:
+            self.qkv_proj = ops.Linear(config.hidden_size, self.inner_size + self.kv_size * 2, bias=config.qkv_bias, device=device, dtype=dtype)
+        else:
+            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 +548,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 +566,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 +619,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:
+            self.gate_up_proj = ops.Linear(config.hidden_size, intermediate_size * 2, bias=False, device=device, dtype=dtype)
+        else:
+            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 +658,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 +675,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 +704,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 +740,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 +752,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 +782,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 +844,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 +884,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 +895,24 @@ 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():
+                nonlocal x
+                x, 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 +921,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 +937,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 +1001,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,117 @@
+import torch
+from tokenizers import Tokenizer
+
+import comfy.ops
+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,
+}
+
+
+def detect_merged_config(state_dict, prefix=""):
+    return {
+        "merged_qkv": "{}model.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
+        "merged_mlp": "{}model.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
+        "decoder_merged_qkv": "{}model.audio_decoder.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
+        "decoder_merged_mlp": "{}model.audio_decoder.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
+    }
+
+
+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={}, projection_config=None):
+        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, **(projection_config or {})}, 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):
+        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_api_nodes/apis/bria.py

@@ -57,6 +57,81 @@ class BriaRemoveBackgroundRequest(BaseModel):
     seed: int = Field(...)
 
 
+class BriaGenFillRequest(BaseModel):
+    image: str = Field(...)
+    mask: str = Field(
+        ...,
+        description="Binary mask defining the region to fill: white (255) pixels are generated, "
+        "black (0) pixels are preserved. Must have the same aspect ratio as the image.",
+    )
+    prompt: str = Field(...)
+    negative_prompt: str | None = Field(None)
+    refine_prompt: bool = Field(True)
+    seed: int = Field(...)
+    prompt_content_moderation: bool = Field(False, description="If true, returns 422 on prompt moderation failure.")
+    visual_input_content_moderation: bool = Field(
+        False, description="If true, returns 422 on image or mask moderation failure."
+    )
+    visual_output_content_moderation: bool = Field(
+        False, description="If true, returns 422 on visual output moderation failure."
+    )
+
+
+class BriaEraseRequest(BaseModel):
+    image: str = Field(...)
+    mask: str = Field(
+        ...,
+        description="Binary mask defining the region to erase: white (255) pixels are removed, "
+        "black (0) pixels are preserved. Must have the same aspect ratio as the image.",
+    )
+    mask_type: str = Field("manual", description="'manual' for hand-drawn masks, 'automatic' for segmentation masks.")
+    visual_input_content_moderation: bool = Field(
+        False, description="If true, returns 422 on image or mask moderation failure."
+    )
+    visual_output_content_moderation: bool = Field(
+        False, description="If true, returns 422 on visual output moderation failure."
+    )
+
+
+class BriaExpandRequest(BaseModel):
+    image: str = Field(...)
+    aspect_ratio: str | float | None = Field(
+        None,
+        description="Target ratio: a preset string (1:1, 2:3, 3:2, 3:4, 4:3, 4:5, 5:4, 9:16, 16:9) "
+        "or a float between 0.5 and 3.0. When set, the canvas/placement fields are ignored.",
+    )
+    canvas_size: list[int] | None = Field(None, description="Output canvas [width, height]; area up to 5000x5000.")
+    original_image_size: list[int] | None = Field(
+        None, description="Size [width, height] of the original image inside the canvas."
+    )
+    original_image_location: list[int] | None = Field(
+        None,
+        description="Top-left corner [x, y] of the original image inside the canvas; "
+        "values may fall outside the canvas, cropping the image.",
+    )
+    prompt: str | None = Field(None, description="If omitted, Bria auto-generates a prompt from the image.")
+    negative_prompt: str | None = Field(None)
+    seed: int = Field(...)
+    prompt_content_moderation: bool = Field(False, description="If true, returns 422 on prompt moderation failure.")
+    visual_input_content_moderation: bool = Field(
+        False, description="If true, returns 422 on image moderation failure."
+    )
+    visual_output_content_moderation: bool = Field(
+        False, description="If true, returns 422 on visual output moderation failure."
+    )
+
+
+class BriaIncreaseResolutionRequest(BaseModel):
+    image: str = Field(...)
+    desired_increase: int = Field(..., description="Resolution multiplier, 2 or 4.")
+    visual_input_content_moderation: bool = Field(
+        False, description="If true, returns 422 on image moderation failure."
+    )
+    visual_output_content_moderation: bool = Field(
+        False, description="If true, returns 422 on visual output moderation failure."
+    )
+
+
 class BriaStatusResponse(BaseModel):
     request_id: str = Field(...)
     status_url: str = Field(...)
@@ -72,6 +147,26 @@ class BriaRemoveBackgroundResponse(BaseModel):
     result: BriaRemoveBackgroundResult | None = Field(None)
 
 
+class BriaImageResult(BaseModel):
+    image_url: str = Field(...)
+
+
+class BriaImageResultResponse(BaseModel):
+    status: str = Field(...)
+    result: BriaImageResult | None = Field(None)
+
+
+class BriaExpandResult(BaseModel):
+    image_url: str = Field(...)
+    prompt: str | None = Field(None)
+    seed: int | None = Field(None)
+
+
+class BriaExpandResponse(BaseModel):
+    status: str = Field(...)
+    result: BriaExpandResult | None = Field(None)
+
+
 class BriaImageEditResult(BaseModel):
     structured_prompt: str = Field(...)
     image_url: str = Field(...)

comfy_api_nodes/apis/minimax.py

@@ -161,12 +161,30 @@ class Hailuo03TaskCreationRequest(BaseModel):
         ..., min_length=1
     )
     resolution: str = Field(...)
-    duration: int = Field(..., ge=5, le=15)
+    duration: int = Field(..., ge=4, le=15)
     ratio: str | None = Field(None)
     seed: int | None = Field(None, ge=0, le=4294967295)
     aigc_watermark: bool | None = Field(None)
 
 
+class Hailuo03ContextIRRequest(BaseModel):
+    model: str = Field(...)
+    content: list[Hailuo03TextContent | Hailuo03ImageContent | Hailuo03VideoContent | Hailuo03AudioContent] = Field(
+        ..., min_length=1
+    )
+    duration: int = Field(..., ge=4, le=15)
+    ratio: str | None = Field(None)
+
+
+class Hailuo03RegenerationRequest(BaseModel):
+    model: str = Field(...)
+    content: list[Hailuo03TextContent | Hailuo03ImageContent | Hailuo03VideoContent | Hailuo03AudioContent] = Field(
+        ..., min_length=1
+    )
+    resolution: str = Field(...)
+    aigc_watermark: bool | None = Field(None)
+
+
 class Hailuo03TaskCreationResponse(BaseModel):
     task_id: str = Field(...)
 
@@ -178,6 +196,7 @@ class Hailuo03TaskError(BaseModel):
 
 class Hailuo03TaskContent(BaseModel):
     url: str | None = Field(None)
+    prompt: str | None = Field(None)
 
 
 class Hailuo03TaskUsage(BaseModel):

comfy_api_nodes/nodes_bria.py

@@ -6,7 +6,13 @@ from typing_extensions import override
 from comfy_api.latest import IO, ComfyExtension, Input
 from comfy_api_nodes.apis.bria import (
     BriaEditImageRequest,
+    BriaEraseRequest,
+    BriaExpandRequest,
+    BriaExpandResponse,
+    BriaGenFillRequest,
     BriaImageEditResponse,
+    BriaImageResultResponse,
+    BriaIncreaseResolutionRequest,
     BriaRemoveBackgroundRequest,
     BriaRemoveBackgroundResponse,
     BriaRemoveVideoBackgroundRequest,
@@ -21,13 +27,30 @@ from comfy_api_nodes.util import (
     convert_mask_to_image,
     download_url_to_image_tensor,
     download_url_to_video_output,
+    downscale_image_tensor_by_max_side,
+    get_image_dimensions,
     poll_op,
     sync_op,
     upload_image_to_comfyapi,
     upload_video_to_comfyapi,
+    validate_string,
     validate_video_duration,
 )
 
+BRIA_MAX_OUTPUT_SIDE = 8192
+BRIA_MIN_RATIO = 0.5
+BRIA_MAX_RATIO = 3.0
+BRIA_MIN_SHORT_SIDE = 224
+
+
+def _upscaled_output_side(height: int, width: int, multiplier: int) -> int:
+    prescale = max(1.0, BRIA_MIN_SHORT_SIDE / min(height, width))
+    return round(max(height, width) * prescale * multiplier)
+
+
+def _smallest_output_side(height: int, width: int, multiplier: int) -> int:
+    return round(max(height, width) / min(height, width) * BRIA_MIN_SHORT_SIDE * multiplier)
+
 
 class BriaImageEditNode(IO.ComfyNode):
 
@@ -243,6 +266,503 @@ class BriaRemoveImageBackground(IO.ComfyNode):
         return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
 
 
+def _mask_to_binary_image(mask: Input.Image, action: str) -> torch.Tensor:
+    binary = (mask > 0.5).float()
+    if not binary.any():
+        raise ValueError(
+            f"The mask is empty, so there is nothing to {action}. Masks are binarized at 50%: "
+            f"areas painted at less than half opacity are ignored."
+        )
+    return convert_mask_to_image(binary)
+
+
+def _validate_mask_aspect_ratio(image: Input.Image, mask: Input.Image) -> None:
+    ih, iw = image.shape[1], image.shape[2]
+    mh, mw = mask.shape[-2], mask.shape[-1]
+    if abs(iw * mh - ih * mw) > 0.01 * ih * mw:
+        raise ValueError(f"Mask must have the same aspect ratio as the image: image is {iw}x{ih}, mask is {mw}x{mh}.")
+
+
+class BriaGenFill(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="BriaGenFill",
+            display_name="Bria Generative Fill",
+            category="partner/image/Bria",
+            description="Generate objects or scenery inside a masked region of an image using Bria.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.Mask.Input(
+                    "mask",
+                    tooltip="White areas are filled with generated content, black areas are preserved. "
+                    "The mask is binarized before sending, so partially painted areas count as white. "
+                    "Must have the same aspect ratio as the image.",
+                ),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Description of what to generate inside the masked region.",
+                ),
+                IO.String.Input("negative_prompt", multiline=True, default=""),
+                IO.Boolean.Input(
+                    "refine_prompt",
+                    default=True,
+                    tooltip="Automatically adjust the prompt for better results; "
+                    "disable to use the prompt exactly as written.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=1,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                ),
+                IO.DynamicCombo.Input(
+                    "moderation",
+                    options=[
+                        IO.DynamicCombo.Option("false", []),
+                        IO.DynamicCombo.Option(
+                            "true",
+                            [
+                                IO.Boolean.Input("prompt_content_moderation", default=False),
+                                IO.Boolean.Input("visual_input_moderation", default=False),
+                                IO.Boolean.Input("visual_output_moderation", default=False),
+                            ],
+                        ),
+                    ],
+                    tooltip="Moderation settings",
+                ),
+            ],
+            outputs=[IO.Image.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":0.0429}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        mask: Input.Image,
+        prompt: str,
+        negative_prompt: str,
+        refine_prompt: bool,
+        seed: int,
+        moderation: InputModerationSettings,
+    ) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1)
+        _validate_mask_aspect_ratio(image, mask)
+        mask_image = _mask_to_binary_image(mask, "fill")
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/bria/v2/image/edit/gen_fill", method="POST"),
+            data=BriaGenFillRequest(
+                image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
+                mask=await upload_image_to_comfyapi(
+                    cls, mask_image, total_pixels=None, wait_label="Uploading mask"
+                ),
+                prompt=prompt,
+                negative_prompt=negative_prompt if negative_prompt else None,
+                refine_prompt=refine_prompt,
+                seed=seed,
+                prompt_content_moderation=moderation.get("prompt_content_moderation", False),
+                visual_input_content_moderation=moderation.get("visual_input_moderation", False),
+                visual_output_content_moderation=moderation.get("visual_output_moderation", False),
+            ),
+            response_model=BriaStatusResponse,
+        )
+        response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
+            status_extractor=lambda r: r.status,
+            response_model=BriaImageResultResponse,
+        )
+        return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
+
+
+class BriaEraser(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="BriaEraser",
+            display_name="Bria Eraser",
+            category="partner/image/Bria",
+            description="Remove objects or areas outlined by a mask from an image using Bria.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.Mask.Input(
+                    "mask",
+                    tooltip="White areas are erased, black areas are preserved. "
+                    "The mask is binarized before sending, so partially painted areas count as white. "
+                    "Must have the same aspect ratio as the image.",
+                ),
+                IO.Combo.Input(
+                    "mask_type",
+                    options=["manual", "automatic"],
+                    tooltip="manual for hand-drawn or brush masks, "
+                    "automatic for masks produced by segmentation models such as SAM.",
+                ),
+                IO.DynamicCombo.Input(
+                    "moderation",
+                    options=[
+                        IO.DynamicCombo.Option("false", []),
+                        IO.DynamicCombo.Option(
+                            "true",
+                            [
+                                IO.Boolean.Input("visual_input_moderation", default=False),
+                                IO.Boolean.Input("visual_output_moderation", default=False),
+                            ],
+                        ),
+                    ],
+                    tooltip="Moderation settings",
+                ),
+            ],
+            outputs=[IO.Image.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":0.0286}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        mask: Input.Image,
+        mask_type: str,
+        moderation: dict,
+    ) -> IO.NodeOutput:
+        _validate_mask_aspect_ratio(image, mask)
+        mask_image = _mask_to_binary_image(mask, "erase")
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/bria/v2/image/edit/erase", method="POST"),
+            data=BriaEraseRequest(
+                image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
+                mask=await upload_image_to_comfyapi(
+                    cls, mask_image, total_pixels=None, wait_label="Uploading mask"
+                ),
+                mask_type=mask_type,
+                visual_input_content_moderation=moderation.get("visual_input_moderation", False),
+                visual_output_content_moderation=moderation.get("visual_output_moderation", False),
+            ),
+            response_model=BriaStatusResponse,
+        )
+        response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
+            status_extractor=lambda r: r.status,
+            response_model=BriaImageResultResponse,
+        )
+        return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
+
+
+class BriaExpandImage(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="BriaExpandImage",
+            display_name="Bria Expand Image",
+            category="partner/image/Bria",
+            description="Expand an image beyond its borders with generated content using Bria.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.DynamicCombo.Input(
+                    "expand_mode",
+                    options=[
+                        *[IO.DynamicCombo.Option(ratio, []) for ratio in
+                          ["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"]],
+                        IO.DynamicCombo.Option(
+                            "custom_ratio",
+                            [
+                                IO.Int.Input(
+                                    "ratio_width",
+                                    default=21,
+                                    min=1,
+                                    max=100,
+                                    tooltip="Width side of the target ratio: 21 and 9 give 21:9.",
+                                ),
+                                IO.Int.Input(
+                                    "ratio_height",
+                                    default=9,
+                                    min=1,
+                                    max=100,
+                                    tooltip="Height side of the target ratio: 21 and 9 give 21:9. "
+                                    f"Bria only accepts width/height between {BRIA_MIN_RATIO} and "
+                                    f"{BRIA_MAX_RATIO}, so anything taller than 1:2 needs the manual mode.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "manual",
+                            [
+                                IO.Int.Input("canvas_width", default=1000, min=64, max=5000),
+                                IO.Int.Input("canvas_height", default=1000, min=64, max=5000),
+                                IO.Int.Input(
+                                    "image_width",
+                                    default=500,
+                                    min=1,
+                                    max=5000,
+                                    tooltip="Width of the original image inside the canvas.",
+                                ),
+                                IO.Int.Input(
+                                    "image_height",
+                                    default=500,
+                                    min=1,
+                                    max=5000,
+                                    tooltip="Height of the original image inside the canvas.",
+                                ),
+                                IO.Int.Input(
+                                    "image_x",
+                                    default=250,
+                                    min=-5000,
+                                    max=5000,
+                                    tooltip="X position of the image's top-left corner inside the canvas; "
+                                    "may fall outside the canvas, cropping the image.",
+                                ),
+                                IO.Int.Input(
+                                    "image_y",
+                                    default=250,
+                                    min=-5000,
+                                    max=5000,
+                                    tooltip="Y position of the image's top-left corner inside the canvas; "
+                                    "may fall outside the canvas, cropping the image.",
+                                ),
+                            ],
+                        ),
+                    ],
+                    tooltip="Target shape of the expanded image: a preset aspect ratio, a custom ratio, "
+                    "or manual placement of the original image on a canvas. "
+                    "Manual is the only mode that can reach a canvas taller than 1:2.",
+                ),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Optional description of the expanded scene; "
+                    "when empty, Bria generates one from the image.",
+                ),
+                IO.String.Input("negative_prompt", multiline=True, default=""),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=1,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                ),
+                IO.DynamicCombo.Input(
+                    "moderation",
+                    options=[
+                        IO.DynamicCombo.Option("false", []),
+                        IO.DynamicCombo.Option(
+                            "true",
+                            [
+                                IO.Boolean.Input("prompt_content_moderation", default=False),
+                                IO.Boolean.Input("visual_input_moderation", default=False),
+                                IO.Boolean.Input("visual_output_moderation", default=False),
+                            ],
+                        ),
+                    ],
+                    tooltip="Moderation settings",
+                ),
+            ],
+            outputs=[
+                IO.Image.Output(),
+                IO.String.Output(display_name="prompt", tooltip="The prompt used for the expansion; "
+                                 "auto-generated by Bria when the prompt input is empty."),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":0.0286}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        expand_mode: dict,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        moderation: InputModerationSettings,
+    ) -> IO.NodeOutput:
+        mode = expand_mode["expand_mode"]
+        aspect_ratio = canvas_size = original_image_size = original_image_location = None
+        if mode == "manual":
+            canvas_size = [expand_mode["canvas_width"], expand_mode["canvas_height"]]
+            original_image_size = [expand_mode["image_width"], expand_mode["image_height"]]
+            original_image_location = [expand_mode["image_x"], expand_mode["image_y"]]
+        elif mode == "custom_ratio":
+            ratio_width, ratio_height = expand_mode["ratio_width"], expand_mode["ratio_height"]
+            aspect_ratio = ratio_width / ratio_height
+            if not BRIA_MIN_RATIO <= aspect_ratio <= BRIA_MAX_RATIO:
+                raise ValueError(
+                    f"Bria accepts a width-to-height ratio between {BRIA_MIN_RATIO} and {BRIA_MAX_RATIO}: "
+                    f"{ratio_width}:{ratio_height} is {aspect_ratio:.4f}. "
+                    f"Use the manual expand mode to reach a canvas of any shape."
+                )
+        else:
+            aspect_ratio = mode
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/bria/v2/image/edit/expand", method="POST"),
+            data=BriaExpandRequest(
+                image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
+                aspect_ratio=aspect_ratio,
+                canvas_size=canvas_size,
+                original_image_size=original_image_size,
+                original_image_location=original_image_location,
+                prompt=prompt if prompt else None,
+                negative_prompt=negative_prompt if negative_prompt else None,
+                seed=seed,
+                prompt_content_moderation=moderation.get("prompt_content_moderation", False),
+                visual_input_content_moderation=moderation.get("visual_input_moderation", False),
+                visual_output_content_moderation=moderation.get("visual_output_moderation", False),
+            ),
+            response_model=BriaStatusResponse,
+        )
+        response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
+            status_extractor=lambda r: r.status,
+            response_model=BriaExpandResponse,
+        )
+        return IO.NodeOutput(
+            await download_url_to_image_tensor(response.result.image_url),
+            response.result.prompt or "",
+        )
+
+
+class BriaIncreaseResolution(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="BriaIncreaseResolution",
+            display_name="Bria Increase Resolution",
+            category="partner/image/Bria",
+            description="Upscale an image by 2x or 4x using Bria, preserving the original content.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.Combo.Input(
+                    "desired_increase",
+                    options=["2", "4"],
+                    tooltip="Resolution multiplier. The output must fit within 8192 pixels on each side.",
+                ),
+                IO.Boolean.Input(
+                    "auto_downscale",
+                    default=False,
+                    tooltip="Automatically lower the multiplier, and downscale the input image if that is "
+                    "still not enough, when the output would exceed the limit.",
+                ),
+                IO.DynamicCombo.Input(
+                    "moderation",
+                    options=[
+                        IO.DynamicCombo.Option("false", []),
+                        IO.DynamicCombo.Option(
+                            "true",
+                            [
+                                IO.Boolean.Input("visual_input_moderation", default=False),
+                                IO.Boolean.Input("visual_output_moderation", default=False),
+                            ],
+                        ),
+                    ],
+                    tooltip="Moderation settings",
+                ),
+            ],
+            outputs=[IO.Image.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":0.0286}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        desired_increase: str,
+        auto_downscale: bool,
+        moderation: dict,
+    ) -> IO.NodeOutput:
+        multiplier = int(desired_increase)
+        height, width = get_image_dimensions(image)
+        if _upscaled_output_side(height, width, multiplier) > BRIA_MAX_OUTPUT_SIDE:
+            candidates = [c for c in (4, 2) if c <= multiplier]
+            if not auto_downscale:
+                predicted = _upscaled_output_side(height, width, multiplier)
+                raise ValueError(
+                    f"Bria can upscale up to a maximum output dimension of {BRIA_MAX_OUTPUT_SIDE} pixels: "
+                    f"input is {width}x{height}, x{multiplier} would be {predicted} pixels on the long side. "
+                    f"Enable auto_downscale, or use a smaller input image or a lower multiplier."
+                )
+            fitted = next(
+                (c for c in candidates if _upscaled_output_side(height, width, c) <= BRIA_MAX_OUTPUT_SIDE), None
+            )
+            if fitted is not None:
+                multiplier = fitted
+            else:
+                shrinkable = next((c for c in sorted(candidates) if _smallest_output_side(height, width, c)
+                                   <= BRIA_MAX_OUTPUT_SIDE), None)
+                if shrinkable is None:
+                    raise ValueError(
+                        f"This image cannot be upscaled by Bria at any multiplier: it is {width}x{height}, and "
+                        f"Bria first enlarges the short side to {BRIA_MIN_SHORT_SIDE} pixels, which pushes the "
+                        f"long side past the {BRIA_MAX_OUTPUT_SIDE} pixel limit. Crop it to a squarer shape first."
+                    )
+                multiplier = shrinkable
+                image = downscale_image_tensor_by_max_side(image, max_side=BRIA_MAX_OUTPUT_SIDE // multiplier)
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/bria/v2/image/edit/increase_resolution", method="POST"),
+            data=BriaIncreaseResolutionRequest(
+                image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
+                desired_increase=multiplier,
+                visual_input_content_moderation=moderation.get("visual_input_moderation", False),
+                visual_output_content_moderation=moderation.get("visual_output_moderation", False),
+            ),
+            response_model=BriaStatusResponse,
+        )
+        response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
+            status_extractor=lambda r: r.status,
+            response_model=BriaImageResultResponse,
+        )
+        return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
+
+
 class BriaRemoveVideoBackground(IO.ComfyNode):
 
     @classmethod
@@ -572,6 +1092,10 @@ class BriaExtension(ComfyExtension):
         return [
             BriaImageEditNode,
             BriaRemoveImageBackground,
+            BriaGenFill,
+            BriaEraser,
+            BriaExpandImage,
+            BriaIncreaseResolution,
             BriaRemoveVideoBackground,
             BriaVideoGreenScreen,
             BriaVideoReplaceBackground,

comfy_api_nodes/nodes_minimax.py

@@ -3,12 +3,14 @@ from typing import Optional
 import torch
 from typing_extensions import override
 
-from comfy_api.latest import IO, ComfyExtension
+from comfy_api.latest import IO, ComfyExtension, Input
 from comfy_api_nodes.apis.minimax import (
     Hailuo03AudioContent,
     Hailuo03AudioContentUrl,
+    Hailuo03ContextIRRequest,
     Hailuo03ImageContent,
     Hailuo03ImageContentUrl,
+    Hailuo03RegenerationRequest,
     Hailuo03TaskCreationRequest,
     Hailuo03TaskCreationResponse,
     Hailuo03TaskQueryResponse,
@@ -456,6 +458,9 @@ HAILUO_03_QUERY_ENDPOINT = "/proxy/minimax/v2/query/video_generation"  # + /{tas
 HAILUO_03_MODELS = {"MiniMax H3": "MiniMax-H3"}
 HAILUO_03_FAILED_STATUSES = ["failed", "cancelled", "expired"]
 
+HAILUO_03_CONTEXT_IR_ENDPOINT = "/proxy/minimax/v2/h3_context_ir"
+HAILUO_03_REGENERATION_ENDPOINT = "/proxy/minimax/v2/video_regeneration"
+
 
 def _hailuo03_model_inputs(include_ratio: bool = True, allow_adaptive: bool = True):
     inputs = [
@@ -487,10 +492,10 @@ def _hailuo03_model_inputs(include_ratio: bool = True, allow_adaptive: bool = Tr
         IO.Int.Input(
             "duration",
             default=5,
-            min=5,
+            min=4,
             max=15,
             step=1,
-            tooltip="Duration of the output video in seconds (5-15).",
+            tooltip="Duration of the output video in seconds (4-15).",
             display_mode=IO.NumberDisplay.slider,
         )
     )
@@ -939,6 +944,592 @@ class MinimaxHailuo03ReferenceNode(IO.ComfyNode):
         )
 
 
+class MinimaxHailuo03ContextIRNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="MinimaxHailuo03ContextIRNode",
+            display_name="MiniMax H3 Context IR (Prompt Enhancer)",
+            category="partner/video/MiniMax",
+            description="Analyze text and media context with MiniMax H3 Context IR and produce an enhanced, "
+            "structured video prompt. Feed the output into the prompt of a MiniMax H3 video node and attach "
+            "the same media there in the same order, because the enhanced prompt refers to the attached "
+            "media by position.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "MiniMax H3",
+                            [
+                                IO.String.Input(
+                                    "prompt",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="Description of the video you intend to generate.",
+                                ),
+                                IO.Int.Input(
+                                    "duration",
+                                    default=5,
+                                    min=4,
+                                    max=15,
+                                    step=1,
+                                    tooltip="Duration of the video you intend to generate, in seconds (4-15).",
+                                    display_mode=IO.NumberDisplay.slider,
+                                ),
+                                IO.Combo.Input(
+                                    "ratio",
+                                    options=["adaptive", "16:9", "4:3", "1:1", "3:4", "9:16", "21:9"],
+                                    default="adaptive",
+                                    tooltip="Aspect ratio of the video you intend to generate. 'adaptive' "
+                                    "requires at least one image, video, or audio input.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_images",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Image.Input("reference_image"),
+                                        names=[
+                                            "image_1",
+                                            "image_2",
+                                            "image_3",
+                                            "image_4",
+                                            "image_5",
+                                            "image_6",
+                                            "image_7",
+                                            "image_8",
+                                            "image_9",
+                                        ],
+                                        min=0,
+                                    ),
+                                    tooltip="Subject or style reference images, referred to in the prompt "
+                                    "as 'Image 1'..'Image 9' in connection order. Up to 9 images.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_videos",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Video.Input("reference_video"),
+                                        names=["video_1", "video_2", "video_3"],
+                                        min=0,
+                                    ),
+                                    tooltip="Motion or scene reference videos, referred to in the prompt "
+                                    "as 'Video 1'..'Video 3' in connection order. Up to 3 videos, "
+                                    "2-15 seconds each, 15 seconds in total.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_audios",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Audio.Input("reference_audio"),
+                                        names=["audio_1", "audio_2", "audio_3"],
+                                        min=0,
+                                    ),
+                                    tooltip="Audio references, referred to in the prompt as "
+                                    "'Audio 1'..'Audio 3' in connection order. Up to 3 clips, "
+                                    "2-15 seconds each, 15 seconds in total. Cannot be used without "
+                                    "a reference image or video.",
+                                ),
+                            ],
+                        )
+                    ],
+                    tooltip="Model to use for prompt enhancement.",
+                ),
+                IO.Image.Input(
+                    "first_frame",
+                    tooltip="First frame of the video you intend to generate. Cannot be combined with "
+                    "reference media.",
+                    optional=True,
+                ),
+                IO.Image.Input(
+                    "last_frame",
+                    tooltip="Last frame of the video you intend to generate. Cannot be combined with "
+                    "reference media.",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                IO.String.Output(),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                depends_on=IO.PriceBadgeDepends(
+                    inputs=["first_frame", "last_frame"],
+                    input_groups=["model.reference_images", "model.reference_videos", "model.reference_audios"],
+                ),
+                expr="""
+                (
+                  $imgsRaw := $lookup(inputGroups, "model.reference_images");
+                  $imgs := $imgsRaw ? $imgsRaw : 0;
+                  $vidsRaw := $lookup(inputGroups, "model.reference_videos");
+                  $vids := $vidsRaw ? $vidsRaw : 0;
+                  $audsRaw := $lookup(inputGroups, "model.reference_audios");
+                  $auds := $audsRaw ? $audsRaw : 0;
+                  $frames := (inputs.first_frame.connected ? 1 : 0) + (inputs.last_frame.connected ? 1 : 0);
+                  ($imgs + $vids + $auds) > 0
+                    ? {"type": "range_usd", "min_usd": 0.06, "max_usd": 0.11, "format": {"approximate": true}}
+                    : $frames > 0
+                      ? {"type": "usd", "usd": 0.05, "format": {"approximate": true}}
+                      : {"type": "usd", "usd": 0.02, "format": {"approximate": true}}
+                )
+                """,
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        first_frame: torch.Tensor | None = None,
+        last_frame: torch.Tensor | None = None,
+    ) -> IO.NodeOutput:
+        validate_string(model["prompt"], strip_whitespace=True, min_length=1)
+
+        reference_images = {k: v for k, v in (model.get("reference_images") or {}).items() if v is not None}
+        reference_videos = {k: v for k, v in (model.get("reference_videos") or {}).items() if v is not None}
+        reference_audios = {k: v for k, v in (model.get("reference_audios") or {}).items() if v is not None}
+        has_frames = first_frame is not None or last_frame is not None
+        has_references = bool(reference_images) or bool(reference_videos) or bool(reference_audios)
+        if has_frames and has_references:
+            raise ValueError(
+                "First/last frame and reference media are mutually exclusive. Use frames for an "
+                "image-to-video prompt, or reference media for a reference-to-video prompt."
+            )
+        if reference_audios and not reference_images and not reference_videos:
+            raise ValueError("Reference audio cannot be used without a reference image or video.")
+        if not has_frames and not has_references and model["ratio"] == "adaptive":
+            raise ValueError(
+                "Ratio 'adaptive' is not supported for text-only requests; select an explicit aspect ratio."
+            )
+
+        for frame in (first_frame, last_frame):
+            if frame is not None:
+                validate_image_aspect_ratio(frame, (2, 5), (5, 2), strict=False)  # 0.4 to 2.5
+                validate_image_dimensions(frame, min_width=256, min_height=256)
+        for image in reference_images.values():
+            validate_image_aspect_ratio(image, (2, 5), (5, 2), strict=False)  # 0.4 to 2.5
+            validate_image_dimensions(image, min_width=256, min_height=256)
+
+        total_video_duration = 0.0
+        for i, video in enumerate(reference_videos.values(), 1):
+            try:
+                fps = float(video.get_frame_rate())
+            except Exception:
+                fps = 0.0
+            if fps and not (23.9 <= fps <= 60.5):
+                raise ValueError(f"Reference video {i} is {fps:.2f} FPS. Supported range is 23.976-60 FPS.")
+            try:
+                dur = video.get_duration()
+            except Exception:
+                continue
+            if dur < 1.8:
+                raise ValueError(f"Reference video {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
+            total_video_duration += dur
+        if total_video_duration > 15.1:
+            raise ValueError(
+                f"Total reference video duration is {total_video_duration:.1f}s. Maximum is 15 seconds."
+            )
+
+        total_audio_duration = 0.0
+        for i, audio in enumerate(reference_audios.values(), 1):
+            dur = int(audio["waveform"].shape[-1]) / int(audio["sample_rate"])
+            if dur < 1.8:
+                raise ValueError(f"Reference audio {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
+            total_audio_duration += dur
+        if total_audio_duration > 15.1:
+            raise ValueError(
+                f"Total reference audio duration is {total_audio_duration:.1f}s. Maximum is 15 seconds."
+            )
+
+        content: list = [Hailuo03TextContent(text=model["prompt"])]
+        if first_frame is not None:
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, first_frame, max_images=1, wait_label="Uploading first frame"
+                            )
+                        )[0],
+                    ),
+                    role="first_frame",
+                )
+            )
+        if last_frame is not None:
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, last_frame, max_images=1, wait_label="Uploading last frame"
+                            )
+                        )[0],
+                    ),
+                    role="last_frame",
+                )
+            )
+        for i, image in enumerate(reference_images.values(), 1):
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, image, max_images=1, wait_label=f"Uploading image {i}"
+                            )
+                        )[0],
+                    ),
+                    role="reference_image",
+                )
+            )
+        for i, video in enumerate(reference_videos.values(), 1):
+            content.append(
+                Hailuo03VideoContent(
+                    video_url=Hailuo03VideoContentUrl(
+                        url=await upload_video_to_comfyapi(cls, video, wait_label=f"Uploading video {i}"),
+                    ),
+                )
+            )
+        for audio in reference_audios.values():
+            content.append(
+                Hailuo03AudioContent(
+                    audio_url=Hailuo03AudioContentUrl(
+                        url=await upload_audio_to_comfyapi(
+                            cls,
+                            audio,
+                            container_format="mp3",
+                            codec_name="libmp3lame",
+                            mime_type="audio/mpeg",
+                        ),
+                    ),
+                )
+            )
+
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=HAILUO_03_CONTEXT_IR_ENDPOINT, method="POST"),
+            response_model=Hailuo03TaskCreationResponse,
+            data=Hailuo03ContextIRRequest(
+                model=HAILUO_03_MODELS[model["model"]],
+                content=content,
+                duration=model["duration"],
+                ratio=None if model["ratio"] == "adaptive" else model["ratio"],
+            ),
+        )
+        task_result = await poll_op(
+            cls,
+            ApiEndpoint(path=f"{HAILUO_03_QUERY_ENDPOINT}/{response.task_id}"),
+            response_model=Hailuo03TaskQueryResponse,
+            status_extractor=lambda r: r.task.status,
+            failed_statuses=HAILUO_03_FAILED_STATUSES,
+            poll_interval=5,
+        )
+        prompt = task_result.task.content.prompt if task_result.task.content else None
+        if not prompt:
+            raise Exception(f"No enhanced prompt in the response: {task_result.model_dump()}")
+        return IO.NodeOutput(prompt)
+
+
+class MinimaxHailuo03RegenerateNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="MinimaxHailuo03RegenerateNode",
+            display_name="MiniMax H3 Regenerate to 2K",
+            category="partner/video/MiniMax",
+            description="Re-render a MiniMax H3 768P output at 2K resolution. Connect the unmodified 768P "
+            "video and the exact prompt used to generate it; if the original generation used first/last "
+            "frames or reference media, attach the same inputs.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "MiniMax H3",
+                            [
+                                IO.String.Input(
+                                    "prompt",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="The exact prompt used to generate the source video.",
+                                ),
+                                IO.Combo.Input(
+                                    "resolution",
+                                    options=["2K"],
+                                    tooltip="Resolution to re-render the source video at.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_images",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Image.Input("reference_image"),
+                                        names=[
+                                            "image_1",
+                                            "image_2",
+                                            "image_3",
+                                            "image_4",
+                                            "image_5",
+                                            "image_6",
+                                            "image_7",
+                                            "image_8",
+                                            "image_9",
+                                        ],
+                                        min=0,
+                                    ),
+                                    tooltip="Reference images from the original generation, in the same "
+                                    "order. Up to 9 images.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_videos",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Video.Input("reference_video"),
+                                        names=["video_1", "video_2", "video_3"],
+                                        min=0,
+                                    ),
+                                    tooltip="Reference videos from the original generation, in the same "
+                                    "order. Up to 3 videos, 2-15 seconds each, 15 seconds in total.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_audios",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Audio.Input("reference_audio"),
+                                        names=["audio_1", "audio_2", "audio_3"],
+                                        min=0,
+                                    ),
+                                    tooltip="Audio references from the original generation, in the same "
+                                    "order. Up to 3 clips, 2-15 seconds each, 15 seconds in total. "
+                                    "Cannot be used without a reference image or video.",
+                                ),
+                            ],
+                        )
+                    ],
+                    tooltip="Model to use for video regeneration.",
+                ),
+                IO.Video.Input(
+                    "video",
+                    tooltip="The MiniMax H3 768P output video to re-render. Connect the unmodified output "
+                    "of a MiniMax H3 video node (24 FPS, 4-15 seconds). 2K outputs cannot be used.",
+                ),
+                IO.Image.Input(
+                    "first_frame",
+                    tooltip="First frame image from the original generation, if one was used.",
+                    optional=True,
+                ),
+                IO.Image.Input(
+                    "last_frame",
+                    tooltip="Last frame image from the original generation, if one was used.",
+                    optional=True,
+                ),
+                IO.Boolean.Input(
+                    "watermark",
+                    default=False,
+                    tooltip="Whether to add an AIGC watermark to the video.",
+                    advanced=True,
+                ),
+            ],
+            outputs=[
+                IO.Video.Output(),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type": "usd", "usd": 0.0715, "format": {"suffix": "/second"}}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        video: Input.Video,
+        watermark: bool,
+        first_frame: torch.Tensor | None = None,
+        last_frame: torch.Tensor | None = None,
+    ) -> IO.NodeOutput:
+        validate_string(model["prompt"], strip_whitespace=True, min_length=1)
+
+        try:
+            fps = float(video.get_frame_rate())
+        except Exception:
+            fps = 0.0
+        if fps and not (23.9 <= fps <= 24.1):
+            raise ValueError(
+                f"The source video is {fps:.2f} FPS. Regeneration accepts unmodified MiniMax H3 768P "
+                "outputs, which are 24 FPS."
+            )
+        try:
+            width, height = video.get_dimensions()
+        except Exception:
+            width = height = 0
+        if width and height and (width % 32 or height % 32 or width * height > 1_032_192):
+            raise ValueError(
+                f"The source video is {width}x{height}. Regeneration accepts MiniMax H3 768P outputs "
+                "(width and height divisible by 32, at most 1,032,192 total pixels); 2K outputs cannot "
+                "be used as a source."
+            )
+        try:
+            frame_count = video.get_frame_count()
+        except Exception:
+            frame_count = 0
+        if frame_count and (frame_count < 107 or frame_count > 362 or (frame_count - 107) % 17):
+            raise ValueError(
+                f"The source video has {frame_count} frames. Regeneration accepts unmodified "
+                "MiniMax H3 outputs, whose length is 107 to 362 frames in steps of 17 "
+                "(4 to 15 seconds at 24 FPS)."
+            )
+
+        reference_images = {k: v for k, v in (model.get("reference_images") or {}).items() if v is not None}
+        reference_videos = {k: v for k, v in (model.get("reference_videos") or {}).items() if v is not None}
+        reference_audios = {k: v for k, v in (model.get("reference_audios") or {}).items() if v is not None}
+        if (first_frame is not None or last_frame is not None) and (
+            reference_images or reference_videos or reference_audios
+        ):
+            raise ValueError(
+                "First/last frame and reference media are mutually exclusive. Use frames for an "
+                "image-to-video prompt, or reference media for a reference-to-video prompt."
+            )
+        if reference_audios and not reference_images and not reference_videos:
+            raise ValueError("Reference audio cannot be used without a reference image or video.")
+
+        for frame in (first_frame, last_frame):
+            if frame is not None:
+                validate_image_aspect_ratio(frame, (2, 5), (5, 2), strict=False)  # 0.4 to 2.5
+                validate_image_dimensions(frame, min_width=256, min_height=256)
+        for image in reference_images.values():
+            validate_image_aspect_ratio(image, (2, 5), (5, 2), strict=False)  # 0.4 to 2.5
+            validate_image_dimensions(image, min_width=256, min_height=256)
+
+        total_video_duration = 0.0
+        for i, ref_video in enumerate(reference_videos.values(), 1):
+            try:
+                ref_fps = float(ref_video.get_frame_rate())
+            except Exception:
+                ref_fps = 0.0
+            if ref_fps and not (23.9 <= ref_fps <= 60.5):
+                raise ValueError(f"Reference video {i} is {ref_fps:.2f} FPS. Supported range is 23.976-60 FPS.")
+            try:
+                dur = ref_video.get_duration()
+            except Exception:
+                continue
+            if dur < 1.8:
+                raise ValueError(f"Reference video {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
+            total_video_duration += dur
+        if total_video_duration > 15.1:
+            raise ValueError(
+                f"Total reference video duration is {total_video_duration:.1f}s. Maximum is 15 seconds."
+            )
+
+        total_audio_duration = 0.0
+        for i, audio in enumerate(reference_audios.values(), 1):
+            dur = int(audio["waveform"].shape[-1]) / int(audio["sample_rate"])
+            if dur < 1.8:
+                raise ValueError(f"Reference audio {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
+            total_audio_duration += dur
+        if total_audio_duration > 15.1:
+            raise ValueError(
+                f"Total reference audio duration is {total_audio_duration:.1f}s. Maximum is 15 seconds."
+            )
+
+        content: list = [
+            Hailuo03VideoContent(
+                video_url=Hailuo03VideoContentUrl(
+                    url=await upload_video_to_comfyapi(cls, video, wait_label="Uploading source video"),
+                ),
+                role="base_video",
+            ),
+            Hailuo03TextContent(text=model["prompt"]),
+        ]
+        if first_frame is not None:
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, first_frame, max_images=1, wait_label="Uploading first frame"
+                            )
+                        )[0],
+                    ),
+                    role="first_frame",
+                )
+            )
+        if last_frame is not None:
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, last_frame, max_images=1, wait_label="Uploading last frame"
+                            )
+                        )[0],
+                    ),
+                    role="last_frame",
+                )
+            )
+        for i, image in enumerate(reference_images.values(), 1):
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, image, max_images=1, wait_label=f"Uploading image {i}"
+                            )
+                        )[0],
+                    ),
+                    role="reference_image",
+                )
+            )
+        for i, ref_video in enumerate(reference_videos.values(), 1):
+            content.append(
+                Hailuo03VideoContent(
+                    video_url=Hailuo03VideoContentUrl(
+                        url=await upload_video_to_comfyapi(cls, ref_video, wait_label=f"Uploading video {i}"),
+                    ),
+                )
+            )
+        for audio in reference_audios.values():
+            content.append(
+                Hailuo03AudioContent(
+                    audio_url=Hailuo03AudioContentUrl(
+                        url=await upload_audio_to_comfyapi(
+                            cls,
+                            audio,
+                            container_format="mp3",
+                            codec_name="libmp3lame",
+                            mime_type="audio/mpeg",
+                        ),
+                    ),
+                )
+            )
+
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=HAILUO_03_REGENERATION_ENDPOINT, method="POST"),
+            response_model=Hailuo03TaskCreationResponse,
+            data=Hailuo03RegenerationRequest(
+                model=HAILUO_03_MODELS[model["model"]],
+                content=content,
+                resolution=model["resolution"],
+                aigc_watermark=watermark,
+            ),
+        )
+        task_result = await poll_op(
+            cls,
+            ApiEndpoint(path=f"{HAILUO_03_QUERY_ENDPOINT}/{response.task_id}"),
+            response_model=Hailuo03TaskQueryResponse,
+            status_extractor=lambda r: r.task.status,
+            failed_statuses=HAILUO_03_FAILED_STATUSES,
+            poll_interval=10,
+        )
+        video_url = task_result.task.content.url if task_result.task.content else None
+        if not video_url:
+            raise Exception(f"No video URL in the response: {task_result.model_dump()}")
+        return IO.NodeOutput(await download_url_to_video_output(video_url))
+
+
 class MinimaxExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
@@ -950,6 +1541,8 @@ class MinimaxExtension(ComfyExtension):
             MinimaxHailuo03TextToVideoNode,
             MinimaxHailuo03FirstLastFrameNode,
             MinimaxHailuo03ReferenceNode,
+            MinimaxHailuo03ContextIRNode,
+            MinimaxHailuo03RegenerateNode,
         ]
 
 

comfy_extras/nodes_minimax_h3.py

@@ -20,6 +20,7 @@ import comfy.model_sampling
 import comfy.nested_tensor
 import comfy.utils
 import node_helpers
+from comfy.ldm.minimax.model import FRAME_PER_TOKEN, FRAME_RESCALE
 from comfy_api.latest import ComfyExtension, io
 
 CANVAS_MULTIPLE = 32
@@ -67,6 +68,16 @@ def _resize(image, width, height, crop):
     return samples.movedim(1, -1)
 
 
+def _encode_ref_audio(audio_vae, audio):
+    waveform = audio["waveform"]  # [B, C, L]
+    sr = audio["sample_rate"]
+    vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
+    if sr != vae_sr:
+        waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
+    z = audio_vae.encode(waveform[:1].movedim(1, -1))  # [1, 32, 2, T]
+    return z, z.shape[-1]
+
+
 def _empty_av_latent(width, height, length, batch_size=1):
     frame_count, latent_t, audio_t = temporal_shape(length)
     video = torch.zeros([batch_size, 24, latent_t, height // 16, width // 16],
@@ -144,13 +155,87 @@ class MiniMaxH3ImageToVideo(io.ComfyNode):
         if keyframes:
             for kf in keyframes:
                 kf["latent"] = vae.encode(kf.pop("image"))
-            cond = node_helpers.conditioning_set_values(cond, {
-                "minimax_keyframes": keyframes,
-                "minimax_frame_count": frame_count,
-            })
+            cond = node_helpers.conditioning_set_values(cond, {"minimax_keyframes": keyframes})
         return io.NodeOutput(cond, latent)
 
 
+class MiniMaxH3AddGuide(io.ComfyNode):
+    """Anchor image and/or audio guides at an arbitrary pixel frame of the target video."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxH3AddGuide",
+            display_name="Add Guide for MiniMax H3",
+            category="model/conditioning/minimax",
+            description="Anchor an image, a short clip, audio, or a clip with its soundtrack at any frame of a MiniMax H3 video. Chain several nodes to anchor several frames.",
+            inputs=[
+                io.Conditioning.Input("positive"),
+                io.Vae.Input("vae", optional=True, tooltip="Video VAE, needed when an image is connected."),
+                io.Vae.Input("audio_vae", optional=True, tooltip="Audio VAE, needed when an audio is connected."),
+                io.Latent.Input("latent"),
+                io.Image.Input("image", optional=True, tooltip="Image or video frames to anchor. Multi-frame batches are anchored as a clip and cropped down to the model's valid clip lengths: 5, 22, 39... (17k + 5) frames. Batches shorter than 5 frames use only the first image."),
+                io.Audio.Input("audio", optional=True,
+                               tooltip="Soundtrack to anchor starting at the same frame index, cropped to the video's remaining duration."),
+                io.Int.Input("frame_idx", default=0, min=-9999, max=9999,
+                             tooltip="Frame index to anchor the image or the clip's first frame at. Negative values are counted from the end of the video."),
+            ],
+            outputs=[io.Conditioning.Output(display_name="positive")],
+        )
+
+    @classmethod
+    def execute(cls, positive, latent, frame_idx, vae=None, audio_vae=None, image=None, audio=None) -> io.NodeOutput:
+        samples = latent["samples"]
+        if not samples.is_nested or len(samples.tensors) != 2 or samples.tensors[0].ndim != 5 or samples.tensors[0].shape[1] != 24:
+            raise ValueError("MiniMaxH3AddGuide expects a MiniMax H3 AV latent")
+        if image is None and audio is None:
+            raise ValueError("MiniMaxH3AddGuide needs an image or an audio to anchor")
+        video = samples.tensors[0]
+        height = video.shape[3] * 16
+        width = video.shape[4] * 16
+        frame_count = sum(FRAME_PER_TOKEN[k % 5] for k in range(video.shape[2]))
+
+        guide_frames = 1
+        if image is not None:
+            if vae is None:
+                raise ValueError("anchoring guide frames needs the vae input")
+            guide_frames = image.shape[0]
+            if guide_frames < 5:
+                guide_frames = 1
+            else:
+                while guide_frames % 17 != 5:
+                    guide_frames -= 1
+
+        resolved_frame_index = frame_idx if frame_idx >= 0 else frame_count + frame_idx
+        if resolved_frame_index < 0 or resolved_frame_index + guide_frames > frame_count:
+            if guide_frames == 1:
+                raise ValueError("frame_idx {} is outside the video's {} frames".format(frame_idx, frame_count))
+            raise ValueError("a {} frame guide clip at frame_idx {} does not fit in the video's {} frames".format(
+                guide_frames, frame_idx, frame_count))
+
+        keyframe = {"resolved_frame_index": resolved_frame_index}
+        if image is not None:
+            frames = _resize(image[:guide_frames], width, height, "center")
+            keyframe["latent"] = vae.encode(frames)
+
+        if audio is not None:
+            if audio_vae is None:
+                raise ValueError("anchoring guide audio needs the audio_vae input")
+            audio_latent, audio_rt = _encode_ref_audio(audio_vae, audio)
+            # the streams share one time axis: FRAME_RESCALE per pixel frame, 1.0 per audio latent frame
+            max_rt = math.floor(samples.tensors[1].shape[-1] - FRAME_RESCALE * resolved_frame_index)
+            if max_rt < 1:
+                raise ValueError("frame_idx {} is past the end of the video's audio track".format(frame_idx))
+            if audio_rt > max_rt:
+                audio_latent = audio_latent[..., :max_rt].clone()
+            keyframe["audio_latent"] = audio_latent
+
+        keyframes = list(positive[0][1].get("minimax_keyframes", []))
+        keyframes.append(keyframe)
+        positive = node_helpers.conditioning_set_values(positive, {"minimax_keyframes": keyframes})
+        return io.NodeOutput(positive)
+
+
 class MiniMaxH3ReferenceToVideo(io.ComfyNode):
     """ref2va: prompt + reference images / videos / audio -> conditioning + AV latent.
 
@@ -197,16 +282,6 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
             outputs=[io.Conditioning.Output(display_name="positive"), io.Latent.Output()],
         )
 
-    @staticmethod
-    def _encode_ref_audio(audio_vae, audio):
-        waveform = audio["waveform"]  # [B, C, L]
-        sr = audio["sample_rate"]
-        vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
-        if sr != vae_sr:
-            waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
-        z = audio_vae.encode(waveform[:1].movedim(1, -1))  # [1, 32, 2, T]
-        return z, z.shape[-1]
-
     @classmethod
     def execute(cls, clip, vae, audio_vae, prompt, width, height, length, ref_image_size="match",
                 ref_images=None, ref_videos=None, ref_video_audios=None, ref_audios=None) -> io.NodeOutput:
@@ -254,7 +329,7 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
             z = vae.encode(frames)
             audio_latent, ref_audio_t = (None, 0)
             if soundtrack is not None:
-                audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, soundtrack)
+                audio_latent, ref_audio_t = _encode_ref_audio(audio_vae, soundtrack)
                 # the soundtrack gets its own <Audio j> label, emitted before <Video k>
                 ref_items.append({"type": "audio"})
             # Qwen sees the video at 2 fps with timestamps
@@ -269,7 +344,7 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
         for audio in (ref_audios or {}).values():
             if audio is None:
                 continue
-            audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, audio)
+            audio_latent, ref_audio_t = _encode_ref_audio(audio_vae, audio)
             ref_items.append({"type": "audio"})
             ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t, "audio_latent": audio_latent})
 
@@ -329,6 +404,7 @@ class MiniMaxH3Extension(ComfyExtension):
         return [
             EmptyMiniMaxH3LatentAV,
             MiniMaxH3ImageToVideo,
+            MiniMaxH3AddGuide,
             MiniMaxH3ReferenceToVideo,
             MiniMaxH3SigmaShift,
             ]

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

comfyui_version.py

@@ -1,3 +1,3 @@
 # This file is automatically generated by the build process when version is
 # updated in pyproject.toml.
-__version__ = "0.32.0"
+__version__ = "0.33.0"

main.py

@@ -248,7 +248,7 @@ import hook_breaker_ac10a0
 import comfy.memory_management
 import comfy.model_patcher
 
-if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia() and not comfy.model_management.is_wsl()):
+if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia()):
     if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)):
         logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows")
     else:

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

openapi.yaml

@@ -35,6 +35,10 @@ components:
                     description: Timestamp when the asset was last accessed
                     format: date-time
                     type: string
+                loader_path:
+                    description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
+                    nullable: true
+                    type: string
                 metadata:
                     additionalProperties: true
                     description: System-managed metadata from download sources (HuggingFace, CivitAI, etc.) - read-only, not user-modifiable
@@ -165,6 +169,10 @@ components:
                     format: uuid
                     nullable: true
                     type: string
+                loader_path:
+                    description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
+                    nullable: true
+                    type: string
                 mime_type:
                     description: Updated MIME type of the asset
                     type: string
@@ -188,6 +196,31 @@ components:
                 - id
                 - updated_at
             type: object
+        ChurnkeyAuthResponse:
+            description: |
+                Credentials the Churnkey embed requires to launch the cancel flow.
+                `auth_hash` is hex-encoded HMAC-SHA256 of `customer_id` signed with the
+                server's CHURNKEY_HMAC_SECRET; it is bound to that single customer ID
+                and must not be reused for other customers.
+            properties:
+                auth_hash:
+                    description: Hex-encoded HMAC-SHA256(customer_id, CHURNKEY_HMAC_SECRET)
+                    type: string
+                customer_id:
+                    description: Stripe customer ID for the workspace
+                    type: string
+                mode:
+                    description: Churnkey environment matching the configured app
+                    enum:
+                        - live
+                        - test
+                        - sandbox
+                    type: string
+            required:
+                - customer_id
+                - auth_hash
+                - mode
+            type: object
         CreateWorkflowRequest:
             description: Request body for creating a new saved workflow.
             properties:
@@ -511,6 +544,25 @@ components:
             required:
                 - history
             type: object
+        JobAssetsResponse:
+            description: Paginated list of the assets produced by a single job.
+            properties:
+                assets:
+                    description: The job's output assets for the requested page (empty when the job produced none)
+                    items:
+                        $ref: '#/components/schemas/JobOutputAsset'
+                    type: array
+                job_id:
+                    description: ID of the job these assets belong to
+                    format: uuid
+                    type: string
+                pagination:
+                    $ref: '#/components/schemas/PaginationInfo'
+            required:
+                - job_id
+                - assets
+                - pagination
+            type: object
         JobCancelResponse:
             description: Response for POST /api/jobs/{job_id}/cancel. Returned on both fresh cancels and idempotent no-ops.
             properties:
@@ -565,6 +617,9 @@ components:
                     additionalProperties: true
                     description: Primary preview output (only for terminal states)
                     type: object
+                previewable_outputs_count:
+                    description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
+                    type: integer
                 status:
                     description: User-friendly job status
                     enum:
@@ -597,6 +652,13 @@ components:
                 workflow_id:
                     description: UUID identifying the workflow graph definition
                     type: string
+                workflow_version_id:
+                    description: |
+                        UUID of the cloud workflow version this job is pinned to, if the
+                        submission carried one (see PromptRequest's workflow_version_id).
+                        Absent for jobs submitted without that association, including
+                        every job submitted through the public API v2 today.
+                    type: string
                 workspace_id:
                     description: |
                         ID of the workspace that owns this job. A successful (200)
@@ -645,6 +707,9 @@ components:
                     additionalProperties: true
                     description: Primary preview output (only present for terminal states)
                     type: object
+                previewable_outputs_count:
+                    description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
+                    type: integer
                 status:
                     description: User-friendly job status
                     enum:
@@ -662,6 +727,56 @@ components:
                 - status
                 - create_time
             type: object
+        JobOutputAsset:
+            description: |
+                An asset produced by a job, enriched with the per-output node context
+                (`node_id`, `output_key`, `output_index`) correlated from the job's
+                execution outputs by content hash. The node-context fields are null
+                when the asset cannot be matched to an output entry.
+            properties:
+                created_at:
+                    description: Timestamp when the asset was created
+                    format: date-time
+                    type: string
+                hash:
+                    description: Blake3 hash of the asset content.
+                    pattern: ^blake3:[a-f0-9]{64}$
+                    type: string
+                id:
+                    description: Unique identifier for the asset
+                    format: uuid
+                    type: string
+                mime_type:
+                    description: MIME type of the asset
+                    type: string
+                name:
+                    description: Name of the asset file
+                    type: string
+                node_id:
+                    description: ID of the workflow node that produced this asset, if known
+                    nullable: true
+                    type: string
+                output_index:
+                    description: Zero-based index of this asset within the node's output slot, if known
+                    nullable: true
+                    type: integer
+                output_key:
+                    description: Output slot key under the producing node (e.g. "images"), if known
+                    nullable: true
+                    type: string
+                preview_url:
+                    description: Relative URL for asset preview/thumbnail
+                    format: uri-reference
+                    type: string
+                size:
+                    description: Size of the asset in bytes
+                    format: int64
+                    type: integer
+            required:
+                - id
+                - name
+                - created_at
+            type: object
         JobStatusResponse:
             description: Job status information
             properties:
@@ -1522,7 +1637,11 @@ paths:
             operationId: listAssets
             parameters:
                 - deprecated: true
-                  description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
+                  description: |
+                    Deprecated alias for `tags_all`, kept permanently for existing
+                    callers. Filter assets that have ALL of these tags. Combining it
+                    with `tags_all`, or exceeding 100 tags (counted after removing
+                    empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: include_tags
@@ -1532,7 +1651,11 @@ paths:
                     type: array
                   style: form
                 - deprecated: true
-                  description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
+                  description: |
+                    Deprecated alias for `tags_none`, kept permanently for existing
+                    callers. Exclude assets that have ANY of these tags. Combining it
+                    with `tags_none`, or exceeding 100 tags (counted after removing
+                    empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: exclude_tags
@@ -1541,7 +1664,14 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Filter assets that have ALL of these tags
+                - description: |
+                    Filter assets that have ALL of these tags. Tag values are opaque
+                    byte-strings compared exactly and case-sensitively; unknown tags
+                    are not an error — they simply match nothing. Replaces the
+                    deprecated `include_tags`. Sending both spellings, listing the
+                    same tag here and in `tags_none`, or exceeding 100 tags per list
+                    (counted after removing empty values and duplicates) returns 400
+                    `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: tags_all
@@ -1550,7 +1680,14 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Filter assets that have AT LEAST ONE of these tags
+                - description: |
+                    Filter assets that have AT LEAST ONE of these tags. Combines with
+                    `tags_all`/`tags_none` by intersection (`tags_none` always wins;
+                    overlap with `tags_none` is allowed and leaves a dead term).
+                    Supplying a positive tag filter (`tags_any`, `tags_all`, or
+                    `include_tags`) replaces the default category filter that is
+                    otherwise applied. Lists over 100 tags (counted after removing
+                    empty values and duplicates) return 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: tags_any
@@ -1559,7 +1696,11 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Exclude assets that have ANY of these tags
+                - description: |
+                    Exclude assets that have ANY of these tags. Replaces the
+                    deprecated `exclude_tags`. Sending both spellings, or exceeding
+                    100 tags per list (counted after removing empty values and
+                    duplicates), returns 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: tags_none
@@ -2342,7 +2483,11 @@ paths:
             operationId: getAssetTagHistogram
             parameters:
                 - deprecated: true
-                  description: 'Deprecated alias of tags_all: filter assets that have ALL of these tags'
+                  description: |
+                    Deprecated alias for `tags_all`, kept permanently for existing
+                    callers. Filter assets that have ALL of these tags. The same
+                    combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: include_tags
@@ -2352,7 +2497,11 @@ paths:
                     type: array
                   style: form
                 - deprecated: true
-                  description: 'Deprecated alias of tags_none: exclude assets that have ANY of these tags'
+                  description: |
+                    Deprecated alias for `tags_none`, kept permanently for existing
+                    callers. Exclude assets that have ANY of these tags. The same
+                    combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: exclude_tags
@@ -2361,7 +2510,10 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Filter assets that have ALL of these tags
+                - description: |
+                    Filter assets that have ALL of these tags. Replaces the deprecated
+                    `include_tags`. The same combination and list-size rules as on
+                    `/api/assets` apply (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: tags_all
@@ -2370,7 +2522,11 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Filter assets that have AT LEAST ONE of these tags
+                - description: |
+                    Filter assets that have AT LEAST ONE of these tags. Combines with
+                    `tags_all`/`tags_none` by intersection (`tags_none` always wins).
+                    The same combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: tags_any
@@ -2379,7 +2535,10 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Exclude assets that have ANY of these tags
+                - description: |
+                    Exclude assets that have ANY of these tags. Replaces the deprecated
+                    `exclude_tags`. The same combination and list-size rules as on
+                    `/api/assets` apply (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: tags_none
@@ -2440,6 +2599,49 @@ paths:
             summary: Get tag histogram for filtered assets
             tags:
                 - file
+    /api/billing/churnkey/auth:
+        get:
+            description: |
+                Returns the Stripe customer identifier and a server-signed
+                HMAC-SHA256 of the customer ID, used to launch the Churnkey-hosted
+                cancellation flow embed.
+            operationId: getChurnkeyAuth
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ChurnkeyAuthResponse'
+                    description: Success
+                "401":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Unauthorized
+                "404":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Workspace has no Stripe customer (never subscribed)
+                "500":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Internal server error
+                "503":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Churnkey is not configured on the server
+            security:
+                - BearerAuth: []
+            summary: Get Churnkey HMAC auth credentials
+            tags:
+                - billing
     /api/embeddings:
         get:
             description: Returns the list of text-encoder embeddings available on disk.
@@ -2460,9 +2662,10 @@ paths:
                 Returns a list of model folders available in the system.
                 This is an experimental endpoint that replaces the legacy /models endpoint.
                 Each folder's name is the identifier to pass to /api/experiment/models/{folder}.
-                Once the model_type migration is active the names are model_type folder_names
-                (e.g. `ultralytics_bbox`); a folder with no folder_name mapping is returned by
-                its directory path.
+                The folder vocabulary is resolved per request from the caller's identity: where the
+                model_type migration is active for that caller the names are model_type folder_names
+                (e.g. `ultralytics_bbox`), and a folder with no folder_name mapping is returned by its
+                directory path. An authenticated response can therefore differ from an anonymous one.
             operationId: getModelFolders
             responses:
                 "200":
@@ -2479,7 +2682,10 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/ErrorResponse'
                     description: Internal server error
-            security: []
+            security:
+                - ApiKeyAuth: []
+                - BearerAuth: []
+                - {}
             summary: Get available model folders
             tags:
                 - file
@@ -2488,6 +2694,10 @@ paths:
             description: |
                 Returns a list of models available in the specified folder.
                 This is an experimental endpoint that provides enhanced model information.
+                Accepted folder identifiers are those returned by /api/experiment/models for the same
+                caller. That vocabulary is request-scoped, so list folders and fetch a folder's models
+                with the same credentials — a name obtained anonymously may not resolve when
+                authenticated, and vice versa.
             operationId: getModelsInFolder
             parameters:
                 - description: The folder name to list models from
@@ -2518,7 +2728,10 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/ErrorResponse'
                     description: Internal server error
-            security: []
+            security:
+                - ApiKeyAuth: []
+                - BearerAuth: []
+                - {}
             summary: Get models in a specific folder
             tags:
                 - file
@@ -3155,6 +3368,74 @@ paths:
             summary: Get full job details
             tags:
                 - workflow
+    /api/jobs/{job_id}/assets:
+        get:
+            description: |
+                Retrieve a paginated list of the assets produced by a specific job,
+                enriched with the per-output node context (`node_id`, `output_key`,
+                `output_index`) correlated from the job's execution outputs by content
+                hash. Unlike `GET /api/assets?job_ids={id}`, this endpoint is scoped to a
+                single job and carries node-level placement, making it suited to job
+                output views rather than the general asset browser. Returns an empty
+                `assets` array for jobs that produced no assets.
+            operationId: getJobAssets
+            parameters:
+                - description: Job identifier (UUID)
+                  in: path
+                  name: job_id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+                - description: Maximum number of assets to return (1-500)
+                  in: query
+                  name: limit
+                  schema:
+                    default: 20
+                    maximum: 500
+                    minimum: 1
+                    type: integer
+                - description: Number of assets to skip for pagination
+                  in: query
+                  name: offset
+                  schema:
+                    default: 0
+                    minimum: 0
+                    type: integer
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/JobAssetsResponse'
+                    description: Success - Job assets returned
+                "400":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Invalid request parameters
+                "401":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Unauthorized - Authentication required
+                "404":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Job not found or does not belong to the user
+                "500":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Internal server error
+            summary: List a job's output assets
+            tags:
+                - workflow
     /api/jobs/{job_id}/cancel:
         post:
             description: |
@@ -3360,6 +3641,12 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/PromptErrorResponse'
                     description: Payment required - Insufficient credits
+                "403":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/PromptErrorResponse'
+                    description: Workspace governance policy blocks one or more partner providers (error.type PARTNER_NODE_DISABLED; error.class_types lists the offending nodes, error.providers the disabled providers)
                 "413":
                     content:
                         application/json:
@@ -3371,7 +3658,7 @@ paths:
                         application/json:
                             schema:
                                 $ref: '#/components/schemas/PromptErrorResponse'
-                    description: Payment required - User has not paid
+                    description: 'Retryable backpressure. Two distinct causes, disambiguated by the body''s `error.type`, NOT by parsing `error.message`: `PAYMENT_REQUIRED` / `FREE_TIER_UNAVAILABLE` / `FREE_TIER_EXHAUSTED` / `PARTNER_NODE_PAYMENT_REQUIRED` (a billing gate - retrying without paying never succeeds), or `QUEUE_LIMIT` (this workspace''s bounded job queue is full - retrying after some queued jobs complete will succeed).'
                 "500":
                     content:
                         application/json:
@@ -5210,6 +5497,8 @@ tags:
       name: user
     - description: Background task management
       name: task
+    - description: Workspace billing and subscription management
+      name: billing
     - description: Workflow storage and version management
       name: workflows
     - description: Job queue state and control

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ComfyUI"
-version = "0.32.0"
+version = "0.33.0"
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.10"

requirements.txt

@@ -1,5 +1,5 @@
 comfyui-frontend-package==1.49.6
-comfyui-workflow-templates==0.11.40
+comfyui-workflow-templates==0.11.41
 comfyui-embedded-docs==0.5.9
 torch
 torchsde
@@ -22,7 +22,7 @@ alembic
 SQLAlchemy>=2.0.0
 filelock
 av>=16.0.0
-comfy-kitchen==0.2.30
+comfy-kitchen==0.2.31
 comfy-aimdo==0.4.13
 requests
 simpleeval>=1.0.0

tests-unit/comfy_test/gemma4_template_test.py

@@ -0,0 +1,61 @@
+"""Gemma4 chat template regression tests."""
+
+import pytest
+import torch
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+import comfy.text_encoders.gemma4 as gemma4  # noqa: E402
+
+PROMPT = "describe a cute anime girl with fennec ears"
+THOUGHT_BLOCK = "<|channel>thought\n<channel|>"
+
+# E2B/E4B and 12B/31B ship different canonical chat templates: only the latter prime a
+# closed thought block when thinking is off.
+NO_PRIMING = [gemma4.Gemma4_E2B, gemma4.Gemma4_E4B]
+PRIMING = [gemma4.Gemma4_31B, gemma4.Gemma4_12B]
+
+
+class _CaptureTemplate:
+    """Stands in for SDTokenizer.tokenize_with_weights so the built template is checked without model files."""
+    llama_text = ""
+
+    def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
+        self.llama_text = text
+        return {}
+
+
+def build_template(variant, **kwargs):
+    prime = variant.tokenizer.tokenizer_class.prime_empty_thought
+    probe = type("Probe", (gemma4.Gemma4_Tokenizer, _CaptureTemplate), {"prime_empty_thought": prime})()
+    probe.tokenize_with_weights(PROMPT, **kwargs)
+    return probe.llama_text
+
+
+@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING)
+def test_thinking_enabled_only_asks_via_the_system_turn(variant):
+    template = build_template(variant, skip_template=False, thinking=True)
+    assert template == f"<|turn>system\n<|think|>\n<turn|>\n<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n"
+
+
+@pytest.mark.parametrize("variant", NO_PRIMING)
+def test_thinking_disabled_does_not_prime_a_thought_channel(variant):
+    template = build_template(variant, skip_template=False, thinking=False)
+    assert template == f"<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n"
+    assert "channel" not in template
+    assert "<|think|>" not in template
+
+
+@pytest.mark.parametrize("variant", PRIMING)
+def test_thinking_disabled_primes_a_thought_channel(variant):
+    template = build_template(variant, skip_template=False, thinking=False)
+    assert template == f"<|turn>user\n{PROMPT}<turn|>\n<|turn>model\n{THOUGHT_BLOCK}"
+
+
+@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING)
+@pytest.mark.parametrize("thinking", [False, True])
+def test_skip_template_passes_text_through_unchanged(variant, thinking):
+    assert build_template(variant, skip_template=True, thinking=thinking) == PROMPT