172 lines
8.8 KiB
Python
172 lines
8.8 KiB
Python
import copy
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
import comfy.conds
|
|
import comfy.patcher_extension
|
|
import nodes
|
|
from comfy.ldm.minimax.model import FRAME_PER_TOKEN
|
|
from comfy_api.latest import ComfyExtension, io, ui
|
|
|
|
|
|
GUIDE_POSITION = "minimax_h3_spatial_guide"
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
|
|
def spatial_guide_layout(native, keyframes):
|
|
_, _, height, width, _ = native.signature
|
|
frame_rows = (height // 2) * (width // 2)
|
|
spacing = 64.0 / math.sqrt(height * width)
|
|
guides = iter(kf for kf in keyframes if kf.get("latent") is not None)
|
|
position_parts, segments = [], []
|
|
img_pos, img_update, audio_pos, audio_update = [], [], [], []
|
|
cursor = 0
|
|
for start, stop, kind in native.segments:
|
|
positions = native.position_ids[start:stop]
|
|
if kind == "cond":
|
|
guide = next(guides)
|
|
if GUIDE_POSITION in guide:
|
|
x, y = guide[GUIDE_POSITION]
|
|
_, _, time, h, w = guide["latent"].shape
|
|
grid = positions.new_empty(time, h // 2, w // 2, 3)
|
|
grid[..., 0] = positions[::frame_rows, 0][:, None, None]
|
|
hh = torch.arange(h // 2, dtype=positions.dtype, device=positions.device)
|
|
ww = torch.arange(w // 2, dtype=positions.dtype, device=positions.device)
|
|
grid[..., 1] = positions[0, 1] + (hh[:, None] + y / 32.0) * spacing
|
|
grid[..., 2] = positions[0, 2] + (ww[None, :] + x / 32.0) * spacing
|
|
positions = grid.reshape(-1, 3)
|
|
count = positions.shape[0]
|
|
position_parts.append(positions)
|
|
segments.append((cursor, cursor + count, kind))
|
|
if kind in ("cond", "ref_img", "video"):
|
|
img_pos.append(torch.arange(cursor, cursor + count))
|
|
img_update.append(torch.full((count,), kind == "video", dtype=torch.bool))
|
|
elif kind in ("cond_audio", "ref_audio", "audio"):
|
|
audio_pos.append(torch.arange(cursor, cursor + count))
|
|
audio_update.append(torch.full((count,), kind == "audio", dtype=torch.bool))
|
|
cursor += count
|
|
|
|
layout = copy.copy(native)
|
|
layout.position_ids = torch.cat(position_parts)
|
|
layout.segments = segments
|
|
layout.seq_len = cursor
|
|
layout.img_pos = torch.cat(img_pos)
|
|
layout.img_update = torch.cat(img_update)
|
|
layout.audio_pos = torch.cat(audio_pos)
|
|
layout.audio_update = torch.cat(audio_update)
|
|
return layout
|
|
|
|
|
|
def spatial_guide_sample(executor, model_wrap, *args, **kwargs):
|
|
for conditions in model_wrap.conds.values():
|
|
for condition in conditions:
|
|
model_conds = condition["model_conds"]
|
|
payload = model_conds["minimax_payload"].cond
|
|
keyframes = payload.get("keyframes", [])
|
|
if any(GUIDE_POSITION in kf for kf in keyframes):
|
|
payload = dict(payload, layout=spatial_guide_layout(payload["layout"], keyframes))
|
|
model_conds["minimax_payload"] = comfy.conds.CONDConstant(payload)
|
|
return executor(model_wrap, *args, **kwargs)
|
|
|
|
|
|
def patch_spatial_guides(model):
|
|
patched = model.clone()
|
|
wrapper_type = comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE
|
|
if not patched.get_wrappers(wrapper_type, GUIDE_POSITION):
|
|
patched.add_wrapper_with_key(wrapper_type, GUIDE_POSITION, spatial_guide_sample)
|
|
return patched
|
|
|
|
|
|
class MiniMaxH3AddSpatialGuide(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="MiniMaxH3AddSpatialGuide",
|
|
display_name="Add Spatial Guide for MiniMax H3",
|
|
category="model/conditioning/minimax",
|
|
description="Place an image or video guide at its original size within the target canvas. Connect both MODEL and positive outputs to the sampler.",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Conditioning.Input("positive"),
|
|
io.Vae.Input("vae"),
|
|
io.Latent.Input("latent"),
|
|
io.Image.Input("image", tooltip="Guide frames at their original size. Video uses 5, 22, 39... frames, matching the native Add Guide node."),
|
|
io.Int.Input("frame_idx", default=0, min=-9999, max=9999),
|
|
io.Int.Input("x", default=0, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION,
|
|
tooltip="Guide's left edge in target pixels. Use half the width difference to center it."),
|
|
io.Int.Input("y", default=0, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION,
|
|
tooltip="Guide's top edge in target pixels. Use half the height difference to center it."),
|
|
],
|
|
outputs=[io.Model.Output(), io.Conditioning.Output(display_name="positive")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, positive, vae, latent, image, frame_idx=0, x=0, y=0) -> 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("MiniMaxH3AddSpatialGuide expects a MiniMax H3 AV latent")
|
|
frame_count = sum(FRAME_PER_TOKEN[k % 5] for k in range(samples.tensors[0].shape[2]))
|
|
guide_frames = image.shape[0]
|
|
guide_frames = 1 if guide_frames < 5 else guide_frames - (guide_frames - 5) % 17
|
|
resolved = frame_idx if frame_idx >= 0 else frame_count + frame_idx
|
|
if resolved < 0 or resolved + guide_frames > frame_count:
|
|
raise ValueError("a {} frame guide at frame_idx {} does not fit in the video's {} frames".format(guide_frames, frame_idx, frame_count))
|
|
|
|
frames = image[:guide_frames, ..., :3]
|
|
height, width = frames.shape[1:3]
|
|
# Complete VAE/DiT patches without scaling or cropping the source.
|
|
if height % 32 or width % 32:
|
|
frames = F.pad(frames.movedim(-1, 1), (0, -width % 32, 0, -height % 32), mode="replicate").movedim(1, -1)
|
|
keyframe = {"resolved_frame_index": resolved, "latent": vae.encode(frames), GUIDE_POSITION: (x, y)}
|
|
positive = [[cond, {**options, "minimax_keyframes": list(options.get("minimax_keyframes", [])) + [keyframe]}]
|
|
for cond, options in positive]
|
|
return io.NodeOutput(patch_spatial_guides(model), positive)
|
|
|
|
|
|
class MiniMaxH3GuidePosition(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="MiniMaxH3GuidePosition",
|
|
display_name="MiniMax H3 Guide Position",
|
|
category="model/conditioning/minimax",
|
|
description="Preview the guide's first frame and drag it within the target canvas. Run this node to refresh input dimensions without sampling.",
|
|
is_output_node=True,
|
|
inputs=[
|
|
io.Int.Input("guide_width", default=640, min=1, max=nodes.MAX_RESOLUTION),
|
|
io.Int.Input("guide_height", default=480, min=1, max=nodes.MAX_RESOLUTION),
|
|
io.Int.Input("canvas_width", default=960, min=1, max=nodes.MAX_RESOLUTION),
|
|
io.Int.Input("canvas_height", default=480, min=1, max=nodes.MAX_RESOLUTION),
|
|
io.Int.Input("x", default=0, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION, socketless=True),
|
|
io.Int.Input("y", default=0, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION, socketless=True),
|
|
io.Image.Input("image", optional=True, tooltip="Connect the same image/video batch as Spatial Guide. The first frame determines the guide preview and its actual size."),
|
|
io.Latent.Input("latent", optional=True, tooltip="Target H3 AV latent. Uses its actual canvas dimensions, including the latent's size alignment."),
|
|
],
|
|
outputs=[io.Int.Output(display_name="x"), io.Int.Output(display_name="y")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, guide_width, guide_height, canvas_width, canvas_height, x, y, image=None, latent=None) -> io.NodeOutput:
|
|
preview = {}
|
|
if image is not None:
|
|
guide_height, guide_width = image.shape[1:3]
|
|
preview["guide_preview"] = ui.PreviewImage(image[:1, ..., :3]).as_dict()["images"]
|
|
if latent is not None:
|
|
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("MiniMaxH3GuidePosition expects a MiniMax H3 AV latent")
|
|
canvas_height, canvas_width = [size * 16 for size in samples.tensors[0].shape[-2:]]
|
|
preview.update({"guide_size": [guide_width, guide_height], "canvas_size": [canvas_width, canvas_height]})
|
|
return io.NodeOutput(x, y, ui=preview)
|
|
|
|
|
|
class MiniMaxH3SpatialGuideExtension(ComfyExtension):
|
|
async def get_node_list(self):
|
|
return [MiniMaxH3AddSpatialGuide, MiniMaxH3GuidePosition]
|
|
|
|
|
|
async def comfy_entrypoint():
|
|
return MiniMaxH3SpatialGuideExtension()
|