Add MiniMaxH3AddGuide for anchoring image and audio guides at any frame

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 @@ def forward(self, x, t_emb, video_seg, audio_seg):
 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 @@ def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=No
 
         img_pos, img_update = [], []
         audio_pos, audio_update = [], []
-        cursor = text_len
         row = text_len
 
+        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
+            # 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:
-                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
+                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
 
-        target_audio_w = (float(w_grid[0]), float(w_grid[-1]))
         if refs:
             cursor = float(text_len)
             for blk in refs:
@@ -389,7 +409,7 @@ def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=No
         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 @@ def _forward(self, x, timestep, context, transformer_options={}, minimax_payload
         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 @@ def _forward(self, x, timestep, context, transformer_options={}, minimax_payload
         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/model_base.py

@@ -2164,13 +2164,13 @@ def extra_conds(self, **kwargs):
         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 +2184,7 @@ def extra_conds(self, **kwargs):
             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
 

comfy_extras/nodes_minimax_h3.py

@@ -20,6 +20,7 @@
 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 @@ def execute(cls, clip, vae, prompt, width, height, length,
         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 @@ def define_schema(cls):
             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 @@ def execute(cls, clip, vae, audio_vae, prompt, width, height, length, ref_image_
             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 @@ def execute(cls, clip, vae, audio_vae, prompt, width, height, length, ref_image_
         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 @@ async def get_node_list(self):
         return [
             EmptyMiniMaxH3LatentAV,
             MiniMaxH3ImageToVideo,
+            MiniMaxH3AddGuide,
             MiniMaxH3ReferenceToVideo,
             MiniMaxH3SigmaShift,
             ]