Minimax h3 controlnet as a model patch instead of a controlnet.

comfy/ldm/minimax/controlnet.py

@@ -0,0 +1,117 @@
+"""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):
+    common = (
+        "control_proj_in.weight",
+        "control_blocks.0.adaln_proj.linear.weight",
+        "control_blocks.0.after_proj.weight",
+        "control_blocks.0.before_proj.weight",
+    )
+    native = (
+        "control_blocks.0.attn.qkv_proj.weight",
+        "control_blocks.0.attn.q_norm.weight",
+        "control_blocks.0.mlp.fc1.weight",
+    )
+    diffusers = (
+        "control_blocks.0.attn.to_q.weight",
+        "control_blocks.0.attn.to_k.weight",
+        "control_blocks.0.attn.to_v.weight",
+        "control_blocks.0.attn.norm_q.weight",
+        "control_blocks.0.ff.net.0.proj.weight",
+    )
+    return all(key in state_dict for key in common) and (all(key in state_dict for key in native) or all(key in state_dict for key in diffusers))
+
+
+def convert_minimax_h3_fun_state_dict(state_dict):
+    if "control_blocks.0.attn.to_q.weight" not in state_dict:
+        return state_dict
+
+    converted = {}
+    for key, value in state_dict.items():
+        if key.endswith(".attn.to_q.weight"):
+            base = key[:-len("to_q.weight")]
+            converted[base + "qkv_proj.weight"] = torch.cat([
+                state_dict[base + "to_q.weight"],
+                state_dict[base + "to_k.weight"],
+                state_dict[base + "to_v.weight"],
+            ], dim=0)
+        elif key.endswith(".attn.to_k.weight") or key.endswith(".attn.to_v.weight"):
+            continue
+        elif key.endswith(".ff.net.0.proj.weight"):
+            half = value.shape[0] // 2
+            converted[key.replace(".ff.net.0.proj.", ".mlp.fc1.")] = torch.cat([value[half:], value[:half]], dim=0)
+        else:
+            converted[key.replace(".attn.norm_q.", ".attn.q_norm.")
+                         .replace(".attn.norm_k.", ".attn.k_norm.")
+                         .replace(".attn.to_out.0.", ".attn.out_proj.")
+                         .replace(".ff.net.2.", ".mlp.fc2.")] = value
+    return converted

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,15 +12,18 @@
 import math
 
 import torch
+import torch.nn.functional as F
 import torchaudio
 
 import nodes
+import comfy.ldm.minimax.controlnet
 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
+from comfy.ldm.minimax.model import FRAME_PER_TOKEN, FRAME_RESCALE, MiniMaxH3Model
 from comfy_api.latest import ComfyExtension, io
 
 CANVAS_MULTIPLE = 32
@@ -399,6 +402,229 @@ 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)
+        if not isinstance(model_patch.model, comfy.ldm.minimax.controlnet.MiniMaxH3FunControl):
+            raise ValueError("this node needs a MiniMax H3 Fun ControlNet model patch")
+
+        diffusion_model = model.get_model_object("diffusion_model")
+        if not isinstance(diffusion_model, MiniMaxH3Model):
+            raise ValueError("MiniMax H3 Fun ControlNet only works with MiniMax H3 models")
+        control_model = model_patch.model
+        if control_model.control_in_dim < diffusion_model.latents_dim:
+            raise ValueError("MiniMax H3 Fun ControlNet input width does not match the base model")
+        if mask is not None and control_model.control_in_dim < diffusion_model.latents_dim * 2 + 1:
+            raise ValueError("this MiniMax H3 Fun model patch does not support inpainting")
+        if control_model.control_proj_in.out_features != diffusion_model.hidden_size:
+            raise ValueError("MiniMax H3 Fun ControlNet hidden width does not match the base model")
+        if control_model.patch_size != diffusion_model.patch_size:
+            raise ValueError("MiniMax H3 Fun ControlNet patch size does not match the base model")
+        if control_model.control_blocks[0].adaln_proj.linear.in_features != diffusion_model.blocks[0].adaln_proj.linear.in_features:
+            raise ValueError("MiniMax H3 Fun ControlNet adaln form does not match the base model")
+        control_attn = control_model.control_blocks[0].attn
+        base_attn = diffusion_model.blocks[0].attn
+        if control_attn.heads != base_attn.heads or control_attn.head_dim != base_attn.head_dim:
+            raise ValueError("MiniMax H3 Fun ControlNet attention shape does not match the base model")
+        if control_model.injection_layers[-1] >= len(diffusion_model.blocks):
+            raise ValueError("MiniMax H3 Fun ControlNet has more injection layers than the base 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 +633,7 @@ async def get_node_list(self):
             MiniMaxH3AddGuide,
             MiniMaxH3ReferenceToVideo,
             MiniMaxH3SigmaShift,
+            MiniMaxH3FunControlNetApply,
             ]