Minimax h3 controlnet as a model patch instead of a controlnet. (#15975)

Author: comfyanonymous

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)

comfy/ldm/minimax/model.py

@@ -731,7 +731,7 @@ def block_wrap(args):
                                          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)

comfy_extras/nodes_minimax_h3.py

@@ -12,12 +12,14 @@
 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 ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingAV, comfy.model_sa
         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 @@ async def get_node_list(self):
             MiniMaxH3AddGuide,
             MiniMaxH3ReferenceToVideo,
             MiniMaxH3SigmaShift,
+            MiniMaxH3FunControlNetApply,
             ]
 
 

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.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 @@ def load_model_patch(self, name):
                     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.",