From e01fb4c56b7a88149d469b99cbbfe3223d715054 Mon Sep 17 00:00:00 2001 From: Barish Ozbay <17261091+drozbay@users.noreply.github.com> Date: Thu, 13 Aug 2026 15:55:36 -0400 Subject: [PATCH] Add MiniMaxH3AddGuide for anchoring image and audio guides at any frame (#15439) --- comfy/ldm/minimax/model.py | 73 +++++++++++++-------- comfy/model_base.py | 10 +-- comfy_extras/nodes_minimax_h3.py | 108 ++++++++++++++++++++++++++----- 3 files changed, 143 insertions(+), 48 deletions(-) diff --git a/comfy/ldm/minimax/model.py b/comfy/ldm/minimax/model.py index f745db884..b6feb8860 100644 --- a/comfy/ldm/minimax/model.py +++ b/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 @@ class FinalLayer(nn.Module): 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 @@ class PackedLayout: img_pos, img_update = [], [] audio_pos, audio_update = [], [] - cursor = text_len row = text_len - if keyframes: - # fl2va: keyframe cond rows right after text, sharing the target spatial grid - 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 - 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; + # anchors count from the target timeline origin, FRAME_RESCALE per pixel frame, 1.0 per audio latent frame + for kf in keyframes: + 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 + if refs: cursor = float(text_len) for blk in refs: @@ -389,7 +409,7 @@ class PackedLayout: 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 @@ class MiniMaxH3Model(nn.Module): 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 @@ class MiniMaxH3Model(nn.Module): 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 = [] diff --git a/comfy/model_base.py b/comfy/model_base.py index 90cab7ac0..6705eb6c3 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -2165,13 +2165,13 @@ class MiniMaxH3(BaseModel): 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: @@ -2185,7 +2185,7 @@ class MiniMaxH3(BaseModel): 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 diff --git a/comfy_extras/nodes_minimax_h3.py b/comfy_extras/nodes_minimax_h3.py index 0b1840e85..0a08f185f 100644 --- a/comfy_extras/nodes_minimax_h3.py +++ b/comfy_extras/nodes_minimax_h3.py @@ -20,6 +20,7 @@ import comfy.model_sampling 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 @@ class MiniMaxH3ImageToVideo(io.ComfyNode): 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 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode): 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 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode): 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