Allow using references with MiniMax-H3 Fun Union and fix prefetch race condition (#16020)

Author: kijai

comfy/ldm/minimax/controlnet.py

@@ -42,8 +42,6 @@ def __init__(self, control_in_dim=49, injection_layers=(0, 10, 20, 30, 40), hidd
             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(
@@ -59,8 +57,13 @@ def init_stream(self, h, control_latent, layout, t_emb):
         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))
 
+        # keyframe/reference conditioning rows get a zero control row
+        img_update = layout.img_update.to(h.device)
+        rows = torch.zeros(img_update.shape[0], patch_dim, dtype=torch.float32, device=h.device)
+        rows[img_update] = target_rows
+
         c = h.clone()
-        c[layout.img_pos.to(h.device)] = self.control_proj_in(target_rows).to(h.dtype)
+        c[layout.img_pos.to(h.device)] = self.control_proj_in(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):

comfy_extras/nodes_minimax_h3.py

@@ -414,6 +414,7 @@ def __init__(self, model_patch, vae, control_video, mask, source_video, strength
         self.control_latent = None
         self.control_latent_shape = None
         self.control_stream = None
+        self.pristine_stream = None
         self.active = False
 
     def _fit_frames(self, frames, frame_count, width, height):
@@ -472,26 +473,28 @@ def diffusion_model_wrapper(self, executor, x, timestep, context, transformer_op
         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
+            self.pristine_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"])
+        # stash only: control weight loads here would clobber the base block's freshly staged weights
+        self.pristine_stream = args["img"].clone()
 
     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)
+        if control_index == 0:
+            self.control_latent = self.control_latent.to(out["img"].device)
+            self.control_stream = self.model_patch.model.init_stream(
+                self.pristine_stream, self.control_latent, args["layout"], args["t_emb"])
+            self.pristine_stream = None
         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"])
@@ -510,6 +513,7 @@ def cleanup(self):
         self.control_latent = None
         self.control_latent_shape = None
         self.control_stream = None
+        self.pristine_stream = None
         self.active = False
 
     def models(self):