Add MiniMaxH3AddGuide for anchoring image and audio guides at any frame (#15439)

This commit is contained in:
Barish Ozbay
2026-08-13 15:55:36 -04:00
committed by GitHub
parent 03fa4e48ba
commit e01fb4c56b
3 changed files with 143 additions and 48 deletions

View File

@@ -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 = []

View File

@@ -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

View File

@@ -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 <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 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
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 @@ class MiniMaxH3Extension(ComfyExtension):
return [
EmptyMiniMaxH3LatentAV,
MiniMaxH3ImageToVideo,
MiniMaxH3AddGuide,
MiniMaxH3ReferenceToVideo,
MiniMaxH3SigmaShift,
]