| """Public general masked-target nodes for MiniMax H3 video editing.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from .masking_ops import ( |
| GRID_MODES, |
| H3_PIXEL_CELL, |
| H3_VIDEO_FPS, |
| MASK_CONVERSION_MODES, |
| floor_h3_frame_count, |
| normalize_comfy_mask, |
| quantize_h3_pixel_mask, |
| reduce_h3_mask_to_video_latent, |
| resize_mask_to_video_latent, |
| temporal_audio_mask, |
| trim_audio_to_frames, |
| ) |
| from .masking_support import require_h3_mask_support |
|
|
|
|
| MASK_MEANINGS = ("white = generate", "white = preserve") |
| AUDIO_MASK_MODES = ( |
| "preserve source audio", |
| "generate all audio", |
| "follow video mask", |
| "custom audio mask", |
| ) |
|
|
|
|
| def _nested_parts(value, label): |
| if isinstance(value, (tuple, list)): |
| return list(value) |
| if bool(getattr(value, "is_nested", False)) and hasattr(value, "unbind"): |
| return list(value.unbind()) |
| raise ValueError( |
| "%s must be a nested H3 video/audio value, not a single tensor." % |
| label |
| ) |
|
|
|
|
| def _target_streams(latent): |
| if not isinstance(latent, dict) or latent.get("samples") is None: |
| raise ValueError("H3 Masked Target needs a LATENT containing samples.") |
| streams = _nested_parts(latent["samples"], "H3 target latent") |
| if len(streams) < 2: |
| raise ValueError( |
| "H3 Masked Target needs both video and audio streams. Encode the " |
| "source media and combine it as a joint AV latent first." |
| ) |
| video, audio = streams[:2] |
| if not isinstance(video, torch.Tensor) or video.ndim != 5: |
| raise ValueError( |
| "H3 target video latent must be [B,C,T,H,W], got %s." % |
| (list(getattr(video, "shape", ())),) |
| ) |
| if not isinstance(audio, torch.Tensor) or audio.ndim != 4: |
| raise ValueError( |
| "H3 target audio latent must be [B,C,channels,T], got %s." % |
| (list(getattr(audio, "shape", ())),) |
| ) |
| if int(video.shape[0]) != 1 or int(audio.shape[0]) != 1: |
| raise ValueError("H3 masked targets currently support batch size 1.") |
| return video, audio |
|
|
|
|
| def _existing_video_mask(value, video): |
| if value.ndim == 5: |
| work = value[:, :1].to(device=video.device, dtype=torch.float32) |
| if int(work.shape[0]) != 1: |
| raise ValueError("Existing H3 video mask has an unsupported batch.") |
| if tuple(work.shape[2:]) != tuple(video.shape[2:]): |
| work = F.interpolate( |
| work, |
| size=tuple(int(v) for v in video.shape[2:]), |
| mode="trilinear", |
| align_corners=False, |
| ) |
| return work.clamp(0.0, 1.0).contiguous() |
| if value.ndim == 4 and int(value.shape[0]) == 1: |
| return _existing_video_mask(value.unsqueeze(1), video) |
| return resize_mask_to_video_latent(value, video) |
|
|
|
|
| def _existing_audio_mask(value, audio): |
| work = value.to(device=audio.device, dtype=torch.float32) |
| if work.ndim == 1: |
| work = work[None, None, None] |
| elif work.ndim == 2: |
| work = work[None, None] |
| elif work.ndim == 3: |
| work = work.unsqueeze(0) |
| if work.ndim != 4 or int(work.shape[0]) != 1: |
| raise ValueError( |
| "Existing H3 audio mask must broadcast to [1,C,channels,T]; got " |
| "%s." % (list(work.shape),) |
| ) |
| work = work.amax(dim=1, keepdim=True) |
| if int(work.shape[2]) != int(audio.shape[2]): |
| work = work.amax(dim=2, keepdim=True).expand( |
| 1, 1, int(audio.shape[2]), int(work.shape[3])) |
| if int(work.shape[-1]) != int(audio.shape[-1]): |
| lead = int(work.shape[2]) |
| work = F.interpolate( |
| work.reshape(lead, 1, int(work.shape[-1])), |
| size=int(audio.shape[-1]), |
| mode="linear", |
| align_corners=False, |
| ).reshape(1, 1, lead, int(audio.shape[-1])) |
| return work.clamp(0.0, 1.0).contiguous() |
|
|
|
|
| def _compose_existing_masks(latent, video_mask, audio_mask, video, audio): |
| existing = latent.get("noise_mask") |
| if existing is None: |
| return video_mask, audio_mask, False |
| parts = _nested_parts(existing, "Existing H3 noise mask") |
| if len(parts) < 2: |
| raise ValueError( |
| "Existing H3 noise mask has no audio component and cannot be " |
| "composed safely." |
| ) |
| old_video = _existing_video_mask(parts[0], video) |
| old_audio = _existing_audio_mask(parts[1], audio) |
| |
| |
| return ( |
| torch.minimum(video_mask, old_video), |
| torch.minimum(audio_mask, old_audio), |
| True, |
| ) |
|
|
|
|
| class MiniMaxH3ContexTrimSourceAV: |
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "frames": ("IMAGE", { |
| "tooltip": "Source video frames. Trailing frames are " |
| "dropped to the largest valid H3 17k+5 length; " |
| "frames are never padded or resized.", |
| }), |
| }, |
| "optional": { |
| "audio": ("AUDIO", { |
| "tooltip": "Optional synchronized source audio, trimmed " |
| "to the returned frame duration at 24 fps.", |
| }), |
| }, |
| } |
|
|
| RETURN_TYPES = ("IMAGE", "AUDIO", "INT", "STRING") |
| RETURN_NAMES = ("trimmed_frames", "trimmed_audio", "h3_length", |
| "trim_info") |
| OUTPUT_TOOLTIPS = ( |
| "Source frames trimmed to the largest valid H3 17k+5 length.", |
| "Synchronized audio trimmed to the returned frame duration, or an " |
| "empty output when no audio was connected.", |
| "The valid H3 frame count represented by both returned streams.", |
| "Human-readable summary of the frame and audio trimming performed.", |
| ) |
| FUNCTION = "trim" |
| CATEGORY = "conditioning/minimax/context_loop/masking" |
| DESCRIPTION = ("Trim source video and optional audio to the largest H3 " |
| "17k+5 run without inventing frames or silence.") |
|
|
| def trim(self, frames, audio=None): |
| if not isinstance(frames, torch.Tensor) or frames.ndim != 4: |
| raise ValueError( |
| "H3 source AV trim expects IMAGE [T,H,W,C], got %s." % |
| (list(getattr(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 / float(H3_VIDEO_FPS) |
| audio_text = "no audio connected" |
| if trimmed_audio is not None: |
| actual = int(trimmed_audio["waveform"].shape[-1]) |
| sample_rate = int(trimmed_audio["sample_rate"]) |
| audio_text = "audio %d/%d samples at %d Hz" % ( |
| actual, target_samples, sample_rate) |
| if actual < target_samples: |
| audio_text += " (source shorter; no silence added)" |
| info = ( |
| "H3 source trim: %d -> %d frames; dropped %d; %.3fs at 24 fps; %s." |
| % (source_count, h3_length, dropped, duration, audio_text) |
| ) |
| return trimmed_frames, trimmed_audio, h3_length, info |
|
|
|
|
| class MiniMaxH3ContexMaskedTarget: |
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "target_latent": ("LATENT", { |
| "tooltip": "Joint source video/audio latent used as the " |
| "actual H3 target. Known regions are preserved " |
| "from this latent; do not add the same source " |
| "again merely as a video reference.", |
| }), |
| "mask": ("MASK", { |
| "tooltip": "Static or per-frame video-space mask. It is " |
| "mapped to H3's causal video-latent frames and " |
| "2x2 latent-token cells. Exact mode accepts " |
| "one static mask or one mask per source frame.", |
| }), |
| "mask_meaning": (list(MASK_MEANINGS), { |
| "default": MASK_MEANINGS[0], |
| "tooltip": "Choose whether white pixels regenerate or " |
| "remain protected.", |
| }), |
| "audio_mode": (list(AUDIO_MASK_MODES), { |
| "default": AUDIO_MASK_MODES[0], |
| "tooltip": "Preserve or generate all source audio, map " |
| "the video mask onto time, or use the optional " |
| "custom audio mask.", |
| }), |
| }, |
| "optional": { |
| "audio_mask": ("MASK", { |
| "tooltip": "Used only by custom audio mask. Spatial " |
| "dimensions are reduced; white time regions " |
| "generate and black regions remain protected.", |
| }), |
| "mask_conversion": (list(MASK_CONVERSION_MODES), { |
| "default": MASK_CONVERSION_MODES[0], |
| "tooltip": "H3 exact maps pixel masks with the video " |
| "VAE's causal 1/4/4/4/4 frame groups and " |
| "2x2 latent-token max coverage. Legacy " |
| "trilinear retains the older interpolated " |
| "conversion for workflow comparison. Changing " |
| "this changes generation; update the Plan's " |
| "generation fingerprint before resuming.", |
| }), |
| }, |
| } |
|
|
| RETURN_TYPES = ("LATENT", "STRING") |
| RETURN_NAMES = ("masked_target", "mask_info") |
| OUTPUT_TOOLTIPS = ( |
| "Copy of target_latent carrying the composed H3 video and audio " |
| "denoise mask.", |
| "Summary of the latent dimensions and generated percentages for " |
| "each stream.", |
| ) |
| FUNCTION = "apply" |
| CATEGORY = "conditioning/minimax/context_loop/masking" |
| DESCRIPTION = ("Attach causal, token-aligned H3 video/audio denoise masks " |
| "to a real source AV target for inpainting and temporal " |
| "edits. Ordinary AV extension needs no user mask.") |
|
|
| def apply( |
| self, |
| target_latent, |
| mask, |
| mask_meaning, |
| audio_mode, |
| audio_mask=None, |
| mask_conversion=MASK_CONVERSION_MODES[0], |
| ): |
| require_h3_mask_support("general masked-target generation") |
| video, audio = _target_streams(target_latent) |
|
|
| generation_mask = normalize_comfy_mask(mask) |
| if mask_meaning == "white = preserve": |
| generation_mask = 1.0 - generation_mask |
| elif mask_meaning != "white = generate": |
| raise ValueError("Unknown H3 mask meaning %r." % mask_meaning) |
| if mask_conversion == MASK_CONVERSION_MODES[0]: |
| video_mask = reduce_h3_mask_to_video_latent( |
| generation_mask, video) |
| conversion_summary = "H3 exact causal/token max" |
| elif mask_conversion == MASK_CONVERSION_MODES[1]: |
| video_mask = resize_mask_to_video_latent(generation_mask, video) |
| conversion_summary = "legacy trilinear" |
| else: |
| raise ValueError( |
| "Unknown H3 mask conversion mode %r." % mask_conversion) |
|
|
| if audio_mode == "preserve source audio": |
| out_audio_mask = torch.zeros( |
| (1, 1, int(audio.shape[2]), int(audio.shape[3])), |
| device=audio.device, |
| dtype=torch.float32, |
| ) |
| elif audio_mode == "generate all audio": |
| out_audio_mask = torch.ones( |
| (1, 1, int(audio.shape[2]), int(audio.shape[3])), |
| device=audio.device, |
| dtype=torch.float32, |
| ) |
| elif audio_mode == "follow video mask": |
| out_audio_mask = temporal_audio_mask(generation_mask, audio) |
| elif audio_mode == "custom audio mask": |
| if audio_mask is None: |
| raise ValueError( |
| "custom audio mask mode requires the optional audio_mask " |
| "input. White generates; black preserves." |
| ) |
| out_audio_mask = temporal_audio_mask(audio_mask, audio) |
| else: |
| raise ValueError("Unknown H3 audio mask mode %r." % audio_mode) |
|
|
| video_mask, out_audio_mask, composed = _compose_existing_masks( |
| target_latent, |
| video_mask, |
| out_audio_mask, |
| video, |
| audio, |
| ) |
|
|
| import comfy.nested_tensor |
|
|
| output = target_latent.copy() |
| output["noise_mask"] = comfy.nested_tensor.NestedTensor( |
| (video_mask, out_audio_mask)) |
| video_percent = 100.0 * float(video_mask.mean().item()) |
| audio_percent = 100.0 * float(out_audio_mask.mean().item()) |
| info = ( |
| "H3 masked target: latent video %dx%dx%d, generate %.2f%%; " |
| "audio %d steps, generate %.2f%%; %s; mask %s; existing mask %s." |
| % ( |
| int(video.shape[4]), |
| int(video.shape[3]), |
| int(video.shape[2]), |
| video_percent, |
| int(audio.shape[-1]), |
| audio_percent, |
| audio_mode, |
| conversion_summary, |
| "intersected" if composed else "none", |
| ) |
| ) |
| return output, info |
|
|
|
|
| class MiniMaxH3ContexMaskGridPreview: |
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "image": ("IMAGE", { |
| "tooltip": "Exact resized H3 canvas used by the video " |
| "VAE. Width and height must be divisible by 32.", |
| }), |
| "mask": ("MASK", { |
| "tooltip": "Static or tracked ComfyUI mask interpreted " |
| "using mask_meaning.", |
| }), |
| "cell_selection": (list(GRID_MODES), { |
| "default": GRID_MODES[0], |
| "tooltip": "Runtime exact matches H3 latent reduction. " |
| "Coverage modes provide deliberate conservative " |
| "or strict alternatives.", |
| }), |
| "cell_adjust": ("INT", { |
| "default": 0, "min": -8, "max": 8, "step": 1, |
| "tooltip": "Grow or shrink by complete 32x32 H3 cells.", |
| }), |
| "preview_frame": ("INT", { |
| "default": 0, "min": 0, "max": 10000, "step": 1, |
| "tooltip": "Frame displayed by grid_preview. The snapped " |
| "MASK output still contains the complete batch.", |
| }), |
| "overlay_opacity": ("FLOAT", { |
| "default": 0.38, "min": 0.0, "max": 1.0, "step": 0.01, |
| "tooltip": "Opacity of the orange snapped-generation " |
| "overlay in grid_preview.", |
| }), |
| "show_grid": ("BOOLEAN", { |
| "default": True, |
| "tooltip": "Draw the effective 32x32 H3 source-pixel " |
| "cell grid over the preview.", |
| }), |
| "show_source_outline": ("BOOLEAN", { |
| "default": True, |
| "tooltip": "Outline the original unsnapped mask so its " |
| "difference from the effective mask is visible.", |
| }), |
| }, |
| "optional": { |
| "mask_meaning": (list(MASK_MEANINGS), { |
| "default": MASK_MEANINGS[0], |
| "tooltip": "Use the same white-pixel convention here " |
| "and on Apply Target Mask. The overlay always " |
| "shows the effective generation region.", |
| }), |
| }, |
| } |
|
|
| RETURN_TYPES = ("MASK", "IMAGE", "STRING") |
| RETURN_NAMES = ("snapped_mask", "grid_preview", "grid_info") |
| OUTPUT_TOOLTIPS = ( |
| "Complete mask batch snapped to H3's effective causal 32x32 " |
| "source-pixel cells, using the selected white-pixel convention.", |
| "Selected source frame with the snapped mask and optional grid or " |
| "source outline overlaid.", |
| "Summary of the selected frame, grid dimensions, and snapped-mask " |
| "coverage.", |
| ) |
| FUNCTION = "preview" |
| CATEGORY = "conditioning/minimax/context_loop/masking" |
| DESCRIPTION = ("Preview and snap a video mask to H3's effective 32x32 " |
| "source-pixel generation cells.") |
|
|
| def preview( |
| self, |
| image, |
| mask, |
| cell_selection, |
| cell_adjust, |
| preview_frame, |
| overlay_opacity, |
| show_grid, |
| show_source_outline, |
| mask_meaning=MASK_MEANINGS[0], |
| ): |
| if not isinstance(image, torch.Tensor) or image.ndim != 4: |
| raise ValueError( |
| "H3 mask preview expects IMAGE [frames,H,W,C], got %s." % |
| (list(getattr(image, "shape", ())),) |
| ) |
| frames, height, width = (int(v) for v in image.shape[:3]) |
| generation_mask = normalize_comfy_mask(mask) |
| if mask_meaning == "white = preserve": |
| generation_mask = 1.0 - generation_mask |
| elif mask_meaning != "white = generate": |
| raise ValueError("Unknown H3 mask meaning %r." % mask_meaning) |
| raw, snapped_generation, cells = quantize_h3_pixel_mask( |
| generation_mask, |
| frames, |
| height, |
| width, |
| mode=cell_selection, |
| cell_adjust=int(cell_adjust), |
| ) |
| selected_frame = min(max(0, int(preview_frame)), frames - 1) |
| preview = image[selected_frame:selected_frame + 1, ..., :3].to( |
| torch.float32).clamp(0.0, 1.0).clone() |
| selected = snapped_generation[selected_frame].to( |
| device=preview.device, dtype=preview.dtype)[None, ..., None] |
| 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 bool(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) |
| preview[0, grid] = preview[0, grid] * 0.28 + cyan * 0.72 |
|
|
| if bool(show_source_outline): |
| source = (raw[selected_frame:selected_frame + 1].to( |
| device=preview.device) >= 0.5).to(preview.dtype) |
| dilated = F.max_pool2d( |
| source, kernel_size=3, stride=1, padding=1) |
| eroded = 1.0 - F.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 > 0).sum( |
| dim=(1, 2, 3)).to(torch.int64).cpu().tolist() |
| cells_per_frame = int(cells.shape[-2]) * int(cells.shape[-1]) |
| if len(counts) <= 8: |
| count_text = ", ".join(str(value) for value in counts) |
| else: |
| count_text = "min %d, max %d, total %d" % ( |
| min(counts), max(counts), sum(counts)) |
| info = ( |
| "H3 grid: %dx%d cells at 32x32 pixels; selected/frame %s / %d; " |
| "mode %s; %s; adjust %+d; preview frame %d/%d." |
| % ( |
| int(cells.shape[-1]), |
| int(cells.shape[-2]), |
| count_text, |
| cells_per_frame, |
| cell_selection, |
| mask_meaning, |
| int(cell_adjust), |
| selected_frame, |
| frames - 1, |
| ) |
| ) |
| snapped = ( |
| 1.0 - snapped_generation |
| if mask_meaning == "white = preserve" |
| else snapped_generation |
| ) |
| return snapped.contiguous(), preview, info |
|
|
|
|
| NODE_CLASS_MAPPINGS = { |
| "MiniMaxH3ContexTrimSourceAV": MiniMaxH3ContexTrimSourceAV, |
| "MiniMaxH3ContexMaskedTarget": MiniMaxH3ContexMaskedTarget, |
| "MiniMaxH3ContexMaskGridPreview": MiniMaxH3ContexMaskGridPreview, |
| } |
|
|
| NODE_DISPLAY_NAME_MAPPINGS = { |
| "MiniMaxH3ContexTrimSourceAV": "MiniMax H3 Masking 路 Trim Source AV", |
| "MiniMaxH3ContexMaskedTarget": "MiniMax H3 Masking 路 Apply Target Mask", |
| "MiniMaxH3ContexMaskGridPreview": "MiniMax H3 Masking 路 Grid Preview", |
| } |
|
|
