From d3eaf6adb388c27245374ac26e585fef4936c485 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:44:42 -0700 Subject: [PATCH] Minimax h3 controlnet as a model patch instead of a controlnet. (#15975) --- comfy/ldm/minimax/controlnet.py | 82 ++++++++++++ comfy/ldm/minimax/model.py | 2 +- comfy_extras/nodes_minimax_h3.py | 203 ++++++++++++++++++++++++++++++ comfy_extras/nodes_model_patch.py | 45 +++++++ 4 files changed, 331 insertions(+), 1 deletion(-) create mode 100644 comfy/ldm/minimax/controlnet.py diff --git a/comfy/ldm/minimax/controlnet.py b/comfy/ldm/minimax/controlnet.py new file mode 100644 index 000000000..86df2e7ee --- /dev/null +++ b/comfy/ldm/minimax/controlnet.py @@ -0,0 +1,82 @@ +"""MiniMax H3 Fun ControlNet-Union model patch.""" + +import torch +import torch.nn as nn + +import comfy.ldm.common_dit +from .model import DiTBlock, patchify_video + + +class ControlDiTBlock(DiTBlock): + def __init__(self, hidden, heads, head_dim, ffn, t_dim, eps, qk_eps, first_block=False, + apply_silu=True, adaln_dtype=None, dtype=None, device=None, operations=None): + super().__init__(hidden, heads, head_dim, ffn, t_dim, eps, qk_eps, apply_silu=apply_silu, + adaln_dtype=adaln_dtype, dtype=dtype, device=device, operations=operations) + if first_block: + self.before_proj = operations.Linear(hidden, hidden, bias=True, dtype=dtype, device=device) + self.after_proj = operations.Linear(hidden, hidden, bias=True, dtype=dtype, device=device) + + +class MiniMaxH3FunControl(torch.nn.Module): + def __init__(self, control_in_dim=49, injection_layers=(0, 10, 20, 30, 40), hidden_size=5376, + num_attention_heads=56, attention_head_dim=128, ffn_hidden_size=14336, + time_embed_dim=2688, patch_size=(1, 2, 2), norm_eps=1e-5, qk_norm_eps=1e-5, + use_adaln_curves=False, dtype=None, device=None, operations=None): + super().__init__() + self.dtype = dtype + self.patch_size = tuple(patch_size) + self.injection_layers = tuple(injection_layers) + if not self.injection_layers or self.injection_layers[0] != 0: + raise ValueError("MiniMax H3 Fun control injection layers must start at layer 0") + if self.injection_layers != tuple(sorted(set(self.injection_layers))): + raise ValueError("MiniMax H3 Fun control injection layers must be unique and increasing") + self.control_in_dim = control_in_dim + patch_dim = control_in_dim * self.patch_size[0] * self.patch_size[1] * self.patch_size[2] + self.control_proj_in = operations.Linear(patch_dim, hidden_size, bias=True, dtype=torch.float32, device=device) + self.control_blocks = nn.ModuleList([ + ControlDiTBlock(hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size, + time_embed_dim, norm_eps, qk_norm_eps, first_block=(i == 0), + apply_silu=not use_adaln_curves, + adaln_dtype=torch.float32 if use_adaln_curves else dtype, + dtype=dtype, device=device, operations=operations) + for i in range(len(self.injection_layers))]) + + def init_stream(self, h, control_latent, layout, t_emb): + if any(kind not in ("text", "audio", "video") for _, _, kind in layout.segments): + raise ValueError("MiniMax H3 Fun ControlNet does not support keyframe or reference conditioning") + adaln_in = self.control_blocks[0].adaln_proj.linear.in_features + if t_emb.shape[-1] != adaln_in: + raise RuntimeError( + "MiniMax H3 controlnet adaln width {} does not match the base model's timestep embedding width {}: " + "the controlnet and base checkpoint use different adaln forms (curve basis vs full), " + "convert the controlnet to match the base model.".format(adaln_in, t_emb.shape[-1])) + + patch_dim = self.control_in_dim * self.patch_size[0] * self.patch_size[1] * self.patch_size[2] + control_latent = comfy.ldm.common_dit.pad_to_patch_size(control_latent.to(torch.float32), self.patch_size) + target_rows = patchify_video(control_latent, self.patch_size) + if target_rows.shape[1] < patch_dim: + target_rows = torch.nn.functional.pad(target_rows, (0, patch_dim - target_rows.shape[1])) + elif target_rows.shape[1] > patch_dim: + raise ValueError("MiniMax H3 control input has {} columns but the model patch expects {}".format(target_rows.shape[1], patch_dim)) + + c = h.clone() + c[layout.img_pos.to(h.device)] = self.control_proj_in(target_rows).to(h.dtype) + return self.control_blocks[0].before_proj(c).add_(h) + + def step(self, index, c, t_emb, mod_segments, rope_freqs, transformer_options): + block = self.control_blocks[index] + c = DiTBlock.forward(block, c, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options) + return c, block.after_proj(c) + + +def is_minimax_h3_fun_state_dict(state_dict): + required = ( + "control_proj_in.weight", + "control_blocks.0.adaln_proj.linear.weight", + "control_blocks.0.after_proj.weight", + "control_blocks.0.before_proj.weight", + "control_blocks.0.attn.qkv_proj.weight", + "control_blocks.0.attn.q_norm.weight", + "control_blocks.0.mlp.fc1.weight", + ) + return all(key in state_dict for key in required) diff --git a/comfy/ldm/minimax/model.py b/comfy/ldm/minimax/model.py index 35f387b2d..b83863c66 100644 --- a/comfy/ldm/minimax/model.py +++ b/comfy/ldm/minimax/model.py @@ -731,7 +731,7 @@ class MiniMaxH3Model(nn.Module): transformer_options=args["transformer_options"])} h = blocks_replace[("double_block", i)]( {"img": h, "t_emb": t_emb, "mod_segments": mod_segments, "rope_freqs": rope_freqs, - "transformer_options": transformer_options}, + "layout": layout, "transformer_options": transformer_options}, {"original_block": block_wrap})["img"] else: h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options) diff --git a/comfy_extras/nodes_minimax_h3.py b/comfy_extras/nodes_minimax_h3.py index 0a08f185f..e7544db41 100644 --- a/comfy_extras/nodes_minimax_h3.py +++ b/comfy_extras/nodes_minimax_h3.py @@ -12,12 +12,14 @@ audio stream's shifted schedule internally). import math import torch +import torch.nn.functional as F import torchaudio import nodes import comfy.model_management import comfy.model_sampling import comfy.nested_tensor +import comfy.patcher_extension import comfy.utils import node_helpers from comfy.ldm.minimax.model import FRAME_PER_TOKEN, FRAME_RESCALE @@ -399,6 +401,206 @@ class MiniMaxH3SigmaShift(io.ComfyNode): return io.NodeOutput(m) +class MiniMaxH3FunControlPatch: + def __init__(self, model_patch, vae, control_video, mask, source_video, strength, sigma_start, sigma_end): + self.model_patch = model_patch + self.vae = vae + self.control_video = control_video + self.mask = mask + self.source_video = source_video + self.strength = strength + self.sigma_start = sigma_start + self.sigma_end = sigma_end + self.control_latent = None + self.control_latent_shape = None + self.control_stream = None + self.active = False + + def _fit_frames(self, frames, frame_count, width, height): + indices = torch.arange(frame_count, device=frames.device).clamp(max=frames.shape[0] - 1) + return comfy.utils.common_upscale(frames[indices], width, height, "bilinear", "center") + + def _encode(self, frames, target_shape): + latent = self.vae.encode(frames.movedim(1, -1)).to(torch.float32) + if tuple(latent.shape) != target_shape: + raise ValueError("MiniMax H3 Fun VAE output shape {} does not match the target {}".format(tuple(latent.shape), target_shape)) + return latent + + def prepare_control_latent(self, target_shape): + target_shape = tuple(target_shape) + if self.control_latent is not None and self.control_latent_shape == target_shape: + return + + latent_frames, latent_height, latent_width = target_shape[2:] + frame_count = max((latent_frames - 2) // 5, 0) * 17 + 5 + spatial_compression = self.vae.spacial_compression_encode() + width = latent_width * spatial_compression + height = latent_height * spatial_compression + loaded_models = comfy.model_management.loaded_models(only_currently_used=True) + + try: + hint = None + if self.control_video is not None: + frames = self._fit_frames(self.control_video, frame_count, width, height) + hint = self._encode(frames, target_shape) + + if self.mask is not None: + mask = (self.mask.reshape(-1, 1, self.mask.shape[-2], self.mask.shape[-1]) > 0.5).to(torch.float32) + indices = torch.arange(frame_count, device=mask.device).clamp(max=mask.shape[0] - 1) + mask = comfy.utils.common_upscale(mask[indices], width, height, "bilinear", "center") + visibility = 1.0 - (mask > 0.5).to(torch.float32) + if self.source_video is None: + source = torch.zeros(frame_count, 3, height, width, dtype=visibility.dtype, device=visibility.device) + else: + source = self._fit_frames(self.source_video, frame_count, width, height) + masked_latent = self._encode(source * visibility.to(source.device), target_shape) + if hint is None: + hint = torch.zeros_like(masked_latent) + visibility_latent = F.interpolate( + visibility.squeeze(1)[None, None], size=(latent_frames, latent_height, latent_width), + mode="trilinear", align_corners=False) + hint = torch.cat([hint, visibility_latent.to(hint.device), masked_latent.to(hint.device)], dim=1) + finally: + comfy.model_management.load_models_gpu(loaded_models) + + self.control_latent = hint + self.control_latent_shape = target_shape + + def diffusion_model_wrapper(self, executor, x, timestep, context, transformer_options={}, **kwargs): + sigmas = transformer_options.get("sigmas") + sigma = float(sigmas[0]) if sigmas is not None else float(timestep.flatten()[0]) / 1000.0 + self.active = self.sigma_end <= sigma <= self.sigma_start + self.control_stream = None + if self.active: + payload = kwargs.get("minimax_payload") or {} + if payload.get("keyframes") or payload.get("refs"): + raise ValueError("MiniMax H3 Fun ControlNet does not support keyframe or reference conditioning") + self.prepare_control_latent(x[0].shape) + try: + return executor(x, timestep, context, transformer_options, **kwargs) + finally: + self.control_stream = None + + def before_block(self, block_index, args): + if not self.active or block_index != self.model_patch.model.injection_layers[0]: + return + self.control_latent = self.control_latent.to(args["img"].device) + self.control_stream = self.model_patch.model.init_stream( + args["img"], self.control_latent, args["layout"], args["t_emb"]) + + def after_block(self, block_index, args, out): + if not self.active: + return out + control_index = self.model_patch.model.injection_layers.index(block_index) + self.control_stream, skip = self.model_patch.model.step( + control_index, self.control_stream, args["t_emb"], args["mod_segments"], args["rope_freqs"], + transformer_options=args["transformer_options"]) + skip[args["layout"].audio_pos.to(skip.device)] = 0 + out["img"].add_(skip, alpha=self.strength) + return out + + def to(self, device_or_dtype): + if isinstance(device_or_dtype, torch.device): + if self.control_latent is not None: + self.control_latent = self.control_latent.to(device_or_dtype) + self.control_stream = None + return self + + def cleanup(self): + self.control_latent = None + self.control_latent_shape = None + self.control_stream = None + self.active = False + + def models(self): + return [self.model_patch] + + def register(self, model): + model.add_wrapper(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, self.diffusion_model_wrapper) + for block_index in self.model_patch.model.injection_layers: + blocks_replace = model.model_options.get("transformer_options", {}).get("patches_replace", {}).get("dit", {}) + previous = blocks_replace.get(("double_block", block_index)) + model.set_model_patch_replace( + MiniMaxH3FunControlBlockPatch(self, block_index, previous), "dit", "double_block", block_index) + + +class MiniMaxH3FunControlBlockPatch: + def __init__(self, control_patch, block_index, previous): + self.control_patch = control_patch + self.block_index = block_index + self.previous = previous + + def __call__(self, args, extra_args): + self.control_patch.before_block(self.block_index, args) + if self.previous is None: + out = extra_args["original_block"](args) + else: + out = self.previous(args, extra_args) + return self.control_patch.after_block(self.block_index, args, out) + + def to(self, device_or_dtype): + self.control_patch.to(device_or_dtype) + if hasattr(self.previous, "to"): + self.previous = self.previous.to(device_or_dtype) + return self + + def cleanup(self): + self.control_patch.cleanup() + if hasattr(self.previous, "cleanup"): + self.previous.cleanup() + + def models(self): + models = self.control_patch.models() + if hasattr(self.previous, "models"): + models += self.previous.models() + return models + + +class MiniMaxH3FunControlNetApply(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="MiniMaxH3FunControlNetApply", + description="Apply a MiniMax H3 Fun ControlNet to a text-to-video model as a model patch.", + display_name="Apply MiniMax H3 Fun ControlNet", + search_aliases=["minimax controlnet", "h3 controlnet", "video inpaint controlnet"], + category="model/patch/minimax", + inputs=[ + io.Model.Input("model"), + io.ModelPatch.Input("model_patch"), + io.Vae.Input("vae"), + io.Float.Input("strength", default=1.0, min=0.0, max=10.0, step=0.01), + io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001, advanced=True), + io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.001, advanced=True), + io.Image.Input("control_video", optional=True), + io.Mask.Input("mask", optional=True, tooltip="1 marks the regions to regenerate."), + io.Image.Input("source_video", optional=True, tooltip="Video behind the mask; only read when a mask is given."), + ], + outputs=[io.Model.Output()], + ) + + @classmethod + def execute(cls, model, model_patch, vae, strength, start_percent, end_percent, + control_video=None, mask=None, source_video=None) -> io.NodeOutput: + if strength == 0 or (control_video is None and mask is None): + return io.NodeOutput(model) + + model_patched = model.clone() + model_sampling = model.get_model_object("model_sampling") + patch = MiniMaxH3FunControlPatch( + model_patch, + vae, + control_video[..., :3].movedim(-1, 1) if control_video is not None else None, + mask, + source_video[..., :3].movedim(-1, 1) if mask is not None and source_video is not None else None, + strength, + float(model_sampling.percent_to_sigma(start_percent)), + float(model_sampling.percent_to_sigma(end_percent)), + ) + patch.register(model_patched) + return io.NodeOutput(model_patched) + + class MiniMaxH3Extension(ComfyExtension): async def get_node_list(self): return [ @@ -407,6 +609,7 @@ class MiniMaxH3Extension(ComfyExtension): MiniMaxH3AddGuide, MiniMaxH3ReferenceToVideo, MiniMaxH3SigmaShift, + MiniMaxH3FunControlNetApply, ] diff --git a/comfy_extras/nodes_model_patch.py b/comfy_extras/nodes_model_patch.py index d81112932..5ce3bb25a 100644 --- a/comfy_extras/nodes_model_patch.py +++ b/comfy_extras/nodes_model_patch.py @@ -1,3 +1,5 @@ +import json + import torch from torch import nn import folder_paths @@ -9,6 +11,7 @@ import comfy.latent_formats import comfy.ldm.lumina.controlnet import comfy.ldm.supir.supir_modules import comfy.ldm.anima.lllite +import comfy.ldm.minimax.controlnet import comfy.ldm.wan.uni3c import comfy.ldm.lightricks.duration_head from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel @@ -266,6 +269,48 @@ class ModelPatchLoader: if torch.count_nonzero(ref_weight) == 0: config['broken'] = True model = comfy.ldm.lumina.controlnet.ZImage_Control(device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast, **config) + elif comfy.ldm.minimax.controlnet.is_minimax_h3_fun_state_dict(sd): + load_device = comfy.model_management.get_torch_device() + quant = comfy.utils.detect_layer_quantization(sd, "") + if quant is not None: + dtype = torch.bfloat16 + operations = comfy.ops.mixed_precision_ops(quant, dtype) + else: + dtype = comfy.model_management.unet_dtype( + model_params=-1, + supported_dtypes=[torch.bfloat16, torch.float32], + weight_dtype=comfy.utils.weight_dtype(sd), + ) + manual_cast_dtype = comfy.model_management.unet_manual_cast( + dtype, load_device, supported_dtypes=[torch.bfloat16, torch.float32]) + operations = comfy.ops.pick_operations(dtype, manual_cast_dtype) + + num_blocks = 0 + while "control_blocks.{}.after_proj.weight".format(num_blocks) in sd: + num_blocks += 1 + injection_layers = tuple(range(0, num_blocks * 10, 10)) + if metadata is not None and "control_blocks_places" in metadata: + injection_layers = tuple(json.loads(metadata["control_blocks_places"])) + if len(injection_layers) != num_blocks: + raise ValueError("MiniMax H3 Fun control_blocks_places metadata does not match the checkpoint") + qkv = sd["control_blocks.0.attn.qkv_proj.weight"] + head_dim = sd["control_blocks.0.attn.q_norm.weight"].shape[0] + use_adaln_curves = metadata is not None and metadata.get("minimax_h3_fun_controlnet") == "adaln_basis" + time_embed_dim = 8 if use_adaln_curves else 2688 + model = comfy.ldm.minimax.controlnet.MiniMaxH3FunControl( + control_in_dim=49, + injection_layers=injection_layers, + hidden_size=sd["control_proj_in.weight"].shape[0], + num_attention_heads=qkv.shape[0] // (3 * head_dim), + attention_head_dim=head_dim, + ffn_hidden_size=sd["control_blocks.0.mlp.fc1.weight"].shape[0] // 2, + time_embed_dim=time_embed_dim, + use_adaln_curves=use_adaln_curves, + operations=operations, + device=comfy.model_management.unet_offload_device(), + dtype=dtype, + ) + model.requires_grad_(False) elif 'controlnet_patch_embedding.weight' in sd: # Uni3C controlnet for Wan attn_key_replace = {".self_attn.to_q.": ".self_attn.q.", ".self_attn.to_k.": ".self_attn.k.",