| """Pure tensor/media helpers for general MiniMax H3 masked targets. |
| |
| The 32-pixel grid mirrors H3's 16x video VAE scale followed by the DiT's 2x2 |
| latent patching. A mask value of one means generate and zero means preserve. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| H3_FRAME_OFFSET = 5 |
| H3_FRAME_STRIDE = 17 |
| H3_VIDEO_FPS = 24 |
| H3_VAE_SCALE = 16 |
| H3_PIXEL_CELL = 32 |
| H3_FRAMES_PER_VIDEO_LATENT = (1, 4, 4, 4, 4) |
| MASK_CONVERSION_MODES = ( |
| "H3 exact (causal/token max)", |
| "legacy trilinear", |
| ) |
| GRID_MODES = ( |
| "runtime exact (latent max)", |
| "any pixel coverage", |
| "50% pixel coverage", |
| "full pixel coverage", |
| ) |
|
|
|
|
| def floor_h3_frame_count(frame_count: int) -> int: |
| """Return the largest valid 17k+5 frame count not above ``frame_count``.""" |
| count = int(frame_count) |
| if count < H3_FRAME_OFFSET: |
| raise ValueError( |
| "MiniMax H3 needs at least %d video frames; received %d." % |
| (H3_FRAME_OFFSET, count) |
| ) |
| return H3_FRAME_OFFSET + ( |
| (count - H3_FRAME_OFFSET) // H3_FRAME_STRIDE |
| ) * H3_FRAME_STRIDE |
|
|
|
|
| def trim_audio_to_frames( |
| audio: Any, |
| frame_count: int, |
| fps: int = H3_VIDEO_FPS, |
| ): |
| """Trim ComfyUI AUDIO to the exact video duration without padding.""" |
| if audio is None: |
| return None, None |
| if not isinstance(audio, dict): |
| raise ValueError( |
| "H3 source AV trim needs a ComfyUI AUDIO value when connected." |
| ) |
| waveform = audio.get("waveform") |
| sample_rate = audio.get("sample_rate") |
| if waveform is None or not hasattr(waveform, "shape") or waveform.ndim < 2: |
| raise ValueError("Connected AUDIO has no valid waveform tensor.") |
| try: |
| resolved_rate = int(sample_rate) |
| except (TypeError, ValueError) as exc: |
| raise ValueError("Connected AUDIO has no valid sample_rate.") from exc |
| if resolved_rate < 1: |
| raise ValueError("Connected AUDIO sample_rate must be positive.") |
|
|
| target_samples = round(int(frame_count) * resolved_rate / int(fps)) |
| trimmed_samples = min(int(waveform.shape[-1]), target_samples) |
| output = audio.copy() |
| output["waveform"] = waveform[..., :trimmed_samples] |
| return output, target_samples |
|
|
|
|
| def normalize_comfy_mask(mask: torch.Tensor) -> torch.Tensor: |
| """Normalize common Comfy MASK spellings to float ``[frames,H,W]``.""" |
| if not isinstance(mask, torch.Tensor): |
| raise ValueError("H3 masking expects a ComfyUI MASK tensor.") |
| if mask.ndim == 2: |
| mask = mask.unsqueeze(0) |
| elif mask.ndim == 4 and int(mask.shape[1]) == 1: |
| mask = mask[:, 0] |
| elif mask.ndim == 4 and int(mask.shape[-1]) == 1: |
| mask = mask[..., 0] |
| if mask.ndim != 3: |
| raise ValueError( |
| "Expected a ComfyUI MASK [frames,H,W], got %s." % |
| (list(mask.shape),) |
| ) |
| if not int(mask.shape[0]) or not int(mask.shape[1]) or not int(mask.shape[2]): |
| raise ValueError("H3 masking received an empty mask tensor.") |
| return mask.to(torch.float32).clamp(0.0, 1.0) |
|
|
|
|
| def resize_mask_to_canvas( |
| mask: torch.Tensor, |
| frames: int, |
| height: int, |
| width: int, |
| ) -> torch.Tensor: |
| """Resize a Comfy mask to ``[frames,1,height,width]``.""" |
| if frames < 1 or height < 1 or width < 1: |
| raise ValueError("Mask canvas dimensions must be positive.") |
| work = normalize_comfy_mask(mask)[None, None] |
| work = F.interpolate( |
| work, |
| size=(int(frames), int(height), int(width)), |
| mode="trilinear", |
| align_corners=False, |
| ) |
| return work[0].movedim(0, 1).contiguous() |
|
|
|
|
| def resize_mask_to_video_latent( |
| mask: torch.Tensor, |
| video: torch.Tensor, |
| ) -> torch.Tensor: |
| """Legacy trilinear conversion to the target H3 video latent dimensions.""" |
| if video.ndim != 5: |
| raise ValueError( |
| "H3 target video latent must be [B,C,T,H,W], got %s." % |
| (list(video.shape),) |
| ) |
| if int(video.shape[0]) != 1: |
| raise ValueError("H3 masked targets currently support batch size 1.") |
| work = normalize_comfy_mask(mask).to(device=video.device)[None, None] |
| work = F.interpolate( |
| work, |
| size=(int(video.shape[2]), int(video.shape[3]), int(video.shape[4])), |
| mode="trilinear", |
| align_corners=False, |
| ) |
| return work.to(torch.float32).contiguous() |
|
|
|
|
| def h3_pixel_frames_for_video_latents(latent_steps: int) -> int: |
| """Return the pixel-frame span represented by H3 video latent steps.""" |
| steps = int(latent_steps) |
| if steps < 1: |
| raise ValueError("H3 video latent length must be positive.") |
| pattern = H3_FRAMES_PER_VIDEO_LATENT |
| complete, trailing = divmod(steps, len(pattern)) |
| return complete * sum(pattern) + sum(pattern[:trailing]) |
|
|
|
|
| def h3_video_latents_for_pixel_frames(pixel_frames: int) -> int: |
| """Return H3 video steps for an exact 1-frame or 17k+5 pixel run.""" |
| frames = int(pixel_frames) |
| if frames == 1: |
| return 1 |
| if frames < H3_FRAME_OFFSET or ( |
| frames - H3_FRAME_OFFSET) % H3_FRAME_STRIDE: |
| raise ValueError( |
| "H3 exact mask conversion needs a 1-frame image or a 17k+5 " |
| "video run; received %d frames." % frames) |
| return 2 + 5 * ((frames - H3_FRAME_OFFSET) // H3_FRAME_STRIDE) |
|
|
|
|
| def h3_video_latent_frame_groups( |
| latent_steps: int, |
| ) -> tuple[tuple[int, int], ...]: |
| """Return the exact causal pixel-frame range feeding each H3 latent step.""" |
| groups = [] |
| start = 0 |
| pattern = H3_FRAMES_PER_VIDEO_LATENT |
| for index in range(int(latent_steps)): |
| end = start + pattern[index % len(pattern)] |
| groups.append((start, end)) |
| start = end |
| if not groups: |
| raise ValueError("H3 video latent length must be positive.") |
| return tuple(groups) |
|
|
|
|
| def _h3_token_snap_max(mask: torch.Tensor) -> torch.Tensor: |
| """Unify each 2x2 H3 latent-pixel token using conservative max coverage.""" |
| if mask.ndim != 3: |
| raise ValueError( |
| "H3 token snapping expects [frames,H,W], got %s." % |
| (list(mask.shape),) |
| ) |
| height, width = int(mask.shape[-2]), int(mask.shape[-1]) |
| work = F.pad( |
| mask[:, None], |
| (0, -width % 2, 0, -height % 2), |
| mode="replicate", |
| ) |
| work = F.max_pool2d(work, kernel_size=2, stride=2) |
| work = work.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1) |
| return work[:, 0, :height, :width].contiguous() |
|
|
|
|
| def reduce_h3_mask_to_video_latent( |
| mask: torch.Tensor, |
| video: torch.Tensor, |
| ) -> torch.Tensor: |
| """Map a pixel mask to H3's causal VAE groups and 2x2 token grid. |
| |
| A tracked mask must contain exactly the pixel-frame span represented by |
| the target latent. A single mask is broadcast over that span. Spatial, |
| temporal, and token reduction all use max coverage so a moving or thin |
| selected region cannot disappear between causal frames or token rows. |
| """ |
| if video.ndim != 5: |
| raise ValueError( |
| "H3 target video latent must be [B,C,T,H,W], got %s." % |
| (list(video.shape),) |
| ) |
| if int(video.shape[0]) != 1: |
| raise ValueError("H3 masked targets currently support batch size 1.") |
|
|
| latent_steps = int(video.shape[2]) |
| expected_frames = h3_pixel_frames_for_video_latents(latent_steps) |
| source = normalize_comfy_mask(mask).to(device=video.device) |
| source_frames = int(source.shape[0]) |
| static = source_frames == 1 |
| if not static and source_frames != expected_frames: |
| raise ValueError( |
| "H3 exact mask conversion needs one static mask or exactly %d " |
| "tracked masks for this %d-step target; received %d. In a loop, " |
| "connect Loop Mask Slice. Choose legacy trilinear only to retain " |
| "the older interpolated behavior." % |
| (expected_frames, latent_steps, source_frames) |
| ) |
|
|
| reduced = F.adaptive_max_pool2d( |
| source[:, None], |
| (int(video.shape[3]), int(video.shape[4])), |
| )[:, 0] |
| reduced = _h3_token_snap_max(reduced) |
| if static: |
| return reduced[None, None].expand( |
| 1, 1, latent_steps, |
| int(reduced.shape[-2]), int(reduced.shape[-1]), |
| ).to(torch.float32).contiguous() |
| groups = h3_video_latent_frame_groups(latent_steps) |
| reduced = torch.stack( |
| [reduced[start:end].amax(dim=0) for start, end in groups], dim=0) |
| return reduced[None, None].to(torch.float32).contiguous() |
|
|
|
|
| def temporal_audio_mask( |
| mask: torch.Tensor, |
| audio: torch.Tensor, |
| ) -> torch.Tensor: |
| """Reduce a video/custom mask to an H3 audio-latent timeline.""" |
| if audio.ndim != 4: |
| raise ValueError( |
| "H3 target audio latent must be [B,C,channels,T], got %s." % |
| (list(audio.shape),) |
| ) |
| if int(audio.shape[0]) != 1: |
| raise ValueError("H3 masked targets currently support batch size 1.") |
| source = normalize_comfy_mask(mask).to(device=audio.device) |
| timeline = source.amax(dim=(-2, -1))[None, None] |
| timeline = F.interpolate( |
| timeline, |
| size=int(audio.shape[-1]), |
| mode="linear", |
| align_corners=False, |
| ) |
| return timeline[:, :, None].expand( |
| int(audio.shape[0]), 1, int(audio.shape[2]), int(audio.shape[3]) |
| ).to(torch.float32).contiguous() |
|
|
|
|
| def _grow_or_shrink_cells(cells: torch.Tensor, amount: int) -> torch.Tensor: |
| if amount == 0: |
| return cells |
| radius = abs(int(amount)) |
| kernel = radius * 2 + 1 |
| if amount > 0: |
| return F.max_pool2d( |
| cells, kernel_size=kernel, stride=1, padding=radius) |
| inverse = 1.0 - cells |
| inverse = F.pad( |
| inverse, |
| (radius, radius, radius, radius), |
| mode="constant", |
| value=1.0, |
| ) |
| return 1.0 - F.max_pool2d(inverse, kernel_size=kernel, stride=1) |
|
|
|
|
| def quantize_h3_pixel_mask( |
| mask: torch.Tensor, |
| frames: int, |
| height: int, |
| width: int, |
| mode: str = GRID_MODES[0], |
| cell_adjust: int = 0, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Return resized raw mask, snapped pixel mask, and 32px cell map.""" |
| if mode not in GRID_MODES: |
| raise ValueError("Unknown H3 mask grid mode: %s" % mode) |
| if height % H3_PIXEL_CELL or width % H3_PIXEL_CELL: |
| raise ValueError( |
| "Accurate H3 grid preview requires width and height divisible by " |
| "%d; received %dx%d." % (H3_PIXEL_CELL, width, height) |
| ) |
|
|
| raw = resize_mask_to_canvas(mask, frames, height, width) |
| grid_h = height // H3_PIXEL_CELL |
| grid_w = width // H3_PIXEL_CELL |
|
|
| if mode == "runtime exact (latent max)": |
| latent = F.adaptive_max_pool2d( |
| raw, |
| (height // H3_VAE_SCALE, width // H3_VAE_SCALE), |
| )[:, 0] |
| latent = _h3_token_snap_max(latent) |
| latent_steps = h3_video_latents_for_pixel_frames(frames) |
| groups = h3_video_latent_frame_groups(latent_steps) |
| latent = torch.cat([ |
| latent[start:end].amax(dim=0, keepdim=True).expand( |
| end - start, int(latent.shape[1]), int(latent.shape[2])) |
| for start, end in groups |
| ], dim=0) |
| cells = latent[:, None, ::2, ::2] |
| else: |
| pixels = raw.reshape( |
| frames, 1, grid_h, H3_PIXEL_CELL, grid_w, H3_PIXEL_CELL) |
| if mode == "any pixel coverage": |
| cells = (pixels.amax(dim=(-3, -1)) >= 0.5).to(raw.dtype) |
| elif mode == "50% pixel coverage": |
| cells = (pixels.mean(dim=(-3, -1)) >= 0.5).to(raw.dtype) |
| else: |
| cells = (pixels.amin(dim=(-3, -1)) >= 0.5).to(raw.dtype) |
|
|
| cells = _grow_or_shrink_cells(cells, int(cell_adjust)) |
| snapped = cells.repeat_interleave( |
| H3_PIXEL_CELL, dim=-2).repeat_interleave(H3_PIXEL_CELL, dim=-1) |
| return raw, snapped[:, 0].contiguous(), cells |
|
|
