Add MiniMaxH3AddGuide for anchoring image and audio guides at any frame (#15439)
This commit is contained in:
parent
03fa4e48ba
commit
e01fb4c56b
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
]
|
||||
|
|
|
|||
Loading…
Reference in New Issue