338 lines
13 KiB
Python
338 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
if __package__:
|
|
from .mask_ops import GRID_MODES, H3_PIXEL_CELL, quantize_h3_pixel_mask
|
|
from .media_ops import H3_VIDEO_FPS, floor_h3_frame_count, trim_audio_to_frames
|
|
from .runtime import install_model_wrappers
|
|
else:
|
|
from mask_ops import GRID_MODES, H3_PIXEL_CELL, quantize_h3_pixel_mask
|
|
from media_ops import H3_VIDEO_FPS, floor_h3_frame_count, trim_audio_to_frames
|
|
from runtime import install_model_wrappers
|
|
|
|
|
|
class MiniMaxH3TrimSourceAV:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"frames": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": (
|
|
"Source video frame batch. Trailing frames are dropped to the largest valid "
|
|
"MiniMax H3 17k+5 frame count. Frames are never padded or resized."
|
|
)
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"audio": (
|
|
"AUDIO",
|
|
{
|
|
"tooltip": (
|
|
"Optional synchronized source audio. It is trimmed to the duration of the "
|
|
"returned frame batch at H3's fixed 24 fps."
|
|
)
|
|
},
|
|
)
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "AUDIO", "INT", "STRING")
|
|
RETURN_NAMES = ("trimmed_frames", "trimmed_audio", "h3_length", "trim_info")
|
|
FUNCTION = "trim"
|
|
CATEGORY = "image/video/minimax"
|
|
DESCRIPTION = (
|
|
"Drops trailing source frames to the largest valid MiniMax H3 17k+5 length and, "
|
|
"when connected, trims synchronized audio to the same 24 fps duration."
|
|
)
|
|
|
|
def trim(self, frames, audio=None):
|
|
if frames.ndim != 4 or frames.shape[-1] < 3:
|
|
raise ValueError(
|
|
f"H3 AV Trim expects ComfyUI IMAGE frames [T,H,W,C], got {list(frames.shape)}."
|
|
)
|
|
|
|
source_count = int(frames.shape[0])
|
|
h3_length = floor_h3_frame_count(source_count)
|
|
trimmed_frames = frames[:h3_length]
|
|
trimmed_audio, target_samples = trim_audio_to_frames(audio, h3_length)
|
|
|
|
dropped = source_count - h3_length
|
|
duration = h3_length / H3_VIDEO_FPS
|
|
audio_text = "no audio connected"
|
|
if trimmed_audio is not None:
|
|
actual_samples = int(trimmed_audio["waveform"].shape[-1])
|
|
sample_rate = int(trimmed_audio["sample_rate"])
|
|
audio_text = f"audio {actual_samples}/{target_samples} samples at {sample_rate} Hz"
|
|
if actual_samples < target_samples:
|
|
audio_text += " (source audio was shorter; no silence was added)"
|
|
info = (
|
|
f"H3 source trim: {source_count} -> {h3_length} frames; dropped {dropped}; "
|
|
f"duration {duration:.3f}s at {H3_VIDEO_FPS} fps; {audio_text}."
|
|
)
|
|
return (trimmed_frames, trimmed_audio, h3_length, info)
|
|
|
|
|
|
class MiniMaxH3PerRowMaskPatch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": (
|
|
"MODEL",
|
|
{
|
|
"tooltip": (
|
|
"MiniMax H3 MODEL to patch for per-row mixed-timestep generation masking. "
|
|
"Connect this MODEL output to your guider."
|
|
)
|
|
},
|
|
)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("patched_model",)
|
|
FUNCTION = "patch"
|
|
CATEGORY = "model/patch/minimax"
|
|
DESCRIPTION = (
|
|
"Temporarily enables AbleJones-style mixed-timestep masking for MiniMax H3 without "
|
|
"editing ComfyUI core files. Only the returned MODEL is affected."
|
|
)
|
|
|
|
def patch(self, model):
|
|
return (install_model_wrappers(model),)
|
|
|
|
|
|
class MiniMaxH3SetGenerationMask:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"av_latent": (
|
|
"LATENT",
|
|
{
|
|
"tooltip": (
|
|
"Joint source video+audio latent. Encode the source media, then combine it with "
|
|
"Concat AV Latent before connecting it here."
|
|
)
|
|
},
|
|
),
|
|
"mask": (
|
|
"MASK",
|
|
{
|
|
"tooltip": (
|
|
"Video-space generation mask. It may contain one mask frame or a batch matching "
|
|
"the source video. The sampler resizes and snaps it to H3's patch grid."
|
|
)
|
|
},
|
|
),
|
|
"mask_meaning": (
|
|
["white = generate", "white = preserve"],
|
|
{"tooltip": "Choose whether white pixels are edited or protected."},
|
|
),
|
|
"audio_mode": (
|
|
["generate audio", "preserve source audio"],
|
|
{
|
|
"tooltip": (
|
|
"Generate audio normally, or pin the complete source-audio latent while editing video."
|
|
)
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("masked_av_latent",)
|
|
FUNCTION = "set_mask"
|
|
CATEGORY = "latent/mask/minimax"
|
|
DESCRIPTION = (
|
|
"Attaches a MiniMax H3 video generation mask to a joint source AV latent. "
|
|
"White normally means generate; black means preserve."
|
|
)
|
|
|
|
def set_mask(self, av_latent, mask, mask_meaning, audio_mode):
|
|
import comfy.nested_tensor
|
|
|
|
samples = av_latent.get("samples")
|
|
if samples is None or not getattr(samples, "is_nested", False):
|
|
raise ValueError(
|
|
"H3 Set Generation Mask needs a joint video+audio LATENT. "
|
|
"Use Concat AV Latent before this node."
|
|
)
|
|
streams = samples.unbind()
|
|
if len(streams) < 2:
|
|
raise ValueError("The connected AV latent does not contain both video and audio streams.")
|
|
|
|
video_mask = mask.to(torch.float32).clamp(0.0, 1.0)
|
|
if mask_meaning == "white = preserve":
|
|
video_mask = 1.0 - video_mask
|
|
audio_mask = (
|
|
torch.ones_like(streams[1], dtype=torch.float32)
|
|
if audio_mode == "generate audio"
|
|
else torch.zeros_like(streams[1], dtype=torch.float32)
|
|
)
|
|
|
|
output = av_latent.copy()
|
|
output["noise_mask"] = comfy.nested_tensor.NestedTensor((video_mask, audio_mask))
|
|
return (output,)
|
|
|
|
|
|
class MiniMaxH3MaskGridPreview:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": (
|
|
"IMAGE",
|
|
{
|
|
"tooltip": (
|
|
"The exact resized H3 canvas frames that will be VAE-encoded. Width and height "
|
|
"must be divisible by 32 so the displayed cells align with H3."
|
|
)
|
|
},
|
|
),
|
|
"mask": (
|
|
"MASK",
|
|
{
|
|
"tooltip": (
|
|
"Connect SAM3 Detect, SAM3 Track to Mask, a painted mask, or any standard "
|
|
"ComfyUI MASK. White is treated as the requested generation area."
|
|
)
|
|
},
|
|
),
|
|
"cell_selection": (
|
|
list(GRID_MODES),
|
|
{
|
|
"tooltip": (
|
|
"Runtime exact reproduces H3's latent resize and per-row max decision. Other "
|
|
"modes select cells from direct pixel coverage and can be more or less conservative."
|
|
)
|
|
},
|
|
),
|
|
"cell_adjust": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": -8,
|
|
"max": 8,
|
|
"step": 1,
|
|
"tooltip": (
|
|
"Grow with positive values or shrink with negative values, measured in complete "
|
|
"32x32 H3 cells."
|
|
),
|
|
},
|
|
),
|
|
"overlay_opacity": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.38,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"tooltip": "Opacity of the orange generated-cell overlay.",
|
|
},
|
|
),
|
|
"show_grid": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "Draw cyan boundaries for every 32x32 H3 cell.",
|
|
},
|
|
),
|
|
"show_source_outline": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "Draw the original SAM3 or input-mask boundary in white.",
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK", "IMAGE", "STRING")
|
|
RETURN_NAMES = ("snapped_mask", "grid_preview", "grid_info")
|
|
FUNCTION = "preview"
|
|
CATEGORY = "latent/mask/minimax"
|
|
DESCRIPTION = (
|
|
"Shows exactly which 32x32 MiniMax H3 cells a SAM3 or standard ComfyUI mask will edit, "
|
|
"and outputs the cell-snapped mask for H3 Set Generation Mask."
|
|
)
|
|
|
|
def preview(
|
|
self,
|
|
image,
|
|
mask,
|
|
cell_selection,
|
|
cell_adjust,
|
|
overlay_opacity,
|
|
show_grid,
|
|
show_source_outline,
|
|
):
|
|
import torch.nn.functional as functional
|
|
|
|
if image.ndim != 4 or image.shape[-1] < 3:
|
|
raise ValueError(f"Expected ComfyUI IMAGE [frames,H,W,C], got {list(image.shape)}")
|
|
frames, height, width = image.shape[:3]
|
|
raw, snapped, cells = quantize_h3_pixel_mask(
|
|
mask,
|
|
frames,
|
|
height,
|
|
width,
|
|
mode=cell_selection,
|
|
cell_adjust=cell_adjust,
|
|
)
|
|
|
|
preview = image[..., :3].to(torch.float32).clamp(0.0, 1.0).clone()
|
|
selected = snapped.to(device=preview.device, dtype=preview.dtype).unsqueeze(-1)
|
|
orange = torch.tensor([1.0, 0.18, 0.0], device=preview.device, dtype=preview.dtype)
|
|
opacity = float(overlay_opacity)
|
|
preview = preview * (1.0 - selected * opacity) + orange * selected * opacity
|
|
|
|
if show_grid:
|
|
grid = torch.zeros((height, width), device=preview.device, dtype=torch.bool)
|
|
grid[::H3_PIXEL_CELL, :] = True
|
|
grid[:, ::H3_PIXEL_CELL] = True
|
|
grid[-1, :] = True
|
|
grid[:, -1] = True
|
|
cyan = torch.tensor([0.0, 0.82, 1.0], device=preview.device, dtype=preview.dtype)
|
|
grid_alpha = 0.72
|
|
preview[:, grid] = preview[:, grid] * (1.0 - grid_alpha) + cyan * grid_alpha
|
|
|
|
if show_source_outline:
|
|
source = (raw.to(device=preview.device) >= 0.5).to(preview.dtype)
|
|
dilated = functional.max_pool2d(source, kernel_size=3, stride=1, padding=1)
|
|
eroded = 1.0 - functional.max_pool2d(1.0 - source, kernel_size=3, stride=1, padding=1)
|
|
outline = ((dilated - eroded) > 0).movedim(1, -1)
|
|
preview = torch.where(outline, torch.ones_like(preview), preview)
|
|
|
|
counts = cells.sum(dim=(1, 2, 3)).to(torch.int64).cpu().tolist()
|
|
cells_per_frame = cells.shape[-2] * cells.shape[-1]
|
|
if len(counts) <= 8:
|
|
count_text = ", ".join(str(value) for value in counts)
|
|
else:
|
|
count_text = f"min {min(counts)}, max {max(counts)}, total {sum(counts)}"
|
|
grid_info = (
|
|
f"H3 grid: {cells.shape[-1]}x{cells.shape[-2]} cells at {H3_PIXEL_CELL}x{H3_PIXEL_CELL} "
|
|
f"pixels; selected cells per frame: {count_text} / {cells_per_frame}; "
|
|
f"mode: {cell_selection}; cell adjustment: {cell_adjust:+d}."
|
|
)
|
|
return (snapped, preview, grid_info)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"MiniMaxH3TrimSourceAV": MiniMaxH3TrimSourceAV,
|
|
"MiniMaxH3PerRowMaskPatch": MiniMaxH3PerRowMaskPatch,
|
|
"MiniMaxH3SetGenerationMask": MiniMaxH3SetGenerationMask,
|
|
"MiniMaxH3MaskGridPreview": MiniMaxH3MaskGridPreview,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MiniMaxH3TrimSourceAV": "MiniMax H3 Trim Source AV to 17k+5",
|
|
"MiniMaxH3PerRowMaskPatch": "MiniMax H3 Per-Row Mask Patch (Experimental)",
|
|
"MiniMaxH3SetGenerationMask": "MiniMax H3 Set Generation Mask",
|
|
"MiniMaxH3MaskGridPreview": "MiniMax H3 Mask Grid Preview & Snap",
|
|
}
|