[Frontend] MiniMax-H3: Add latent-mask editing to the ComfyUI extension

apps/ComfyUI-vLLM-Omni/README.md

@@ -271,10 +271,10 @@ The [WF-05 template](example_workflows/vLLM-Omni%20MiniMax-H3%20Latent%20Mask%20
 Connect a **Latent Mask Editing** node to **Generate Video → latent_edit** to edit a source clip instead of generating from scratch. It uploads the source media and serializes the video/audio noise masks the MiniMax H3 API accepts:
 
 - `source_video` / `source_audio` — the media to edit.
-- `video_mask` — a ComfyUI mask image; `0` preserves a region, `1` regenerates it, fractional values blend. A 2D mask `[H, W]` is applied to every frame; a 3D mask `[T, H, W]` is treated as a temporal mask (one slice per frame) for continuation or extension. It is resized to the video latent grid.
+- `video_mask` — a ComfyUI mask image; `0` preserves a region, `1` regenerates it, fractional values blend. A 2D mask `[H, W]` is applied to every frame; a 3D mask `[T, H, W]` is treated as a temporal mask (one slice per frame) for continuation or extension. The client area-downsamples it by 16 to keep the upload small, and the server resizes it to the video latent grid.
 - `audio_mask` — a scalar in `[0, 1]`; `0` keeps the source audio, `1` regenerates it, fractional values blend.
 
-A non-trivial mask requires its matching source, and a source without a mask is rejected. This node only forwards inputs to the server; the served model must declare latent-mask editing support (MiniMax-H3).
+The node requires at least one mask; the server then enforces the cross-field rules (a source requires a matching mask, and a non-trivial mask requires its source). This node only forwards inputs to the server; the served model must declare latent-mask editing support (MiniMax-H3).
 
 ### TTS (e.g., Qwen TTS series)
 

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -36,7 +36,6 @@
 
 # A dictionary that contains all nodes you want to export with their names
 NODE_CLASS_MAPPINGS = {
-    "VLLMOmniMiniMaxH3TemporalMask": VLLMOmniMiniMaxH3TemporalMask,
     # === Generation ===
     "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
     "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
@@ -46,6 +45,7 @@
     "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
     "VLLMOmniVideoReferences": VLLMOmniVideoReferences,
     "VLLMOmniLatentMaskEditing": VLLMOmniLatentMaskEditing,
+    "VLLMOmniMiniMaxH3TemporalMask": VLLMOmniMiniMaxH3TemporalMask,
     # === Params ===
     "VLLMOmniARSampling": VLLMOmniARSampling,
     "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
@@ -59,7 +59,6 @@
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "VLLMOmniMiniMaxH3TemporalMask": "MiniMax-H3 Temporal Mask",
     # === Generation ===
     "VLLMOmniGenerateImage": "Generate Image",
     "VLLMOmniGenerateVideo": "Generate Video",
@@ -69,6 +68,7 @@
     "VLLMOmniVoiceClone": "TTS Voice Cloning",
     "VLLMOmniVideoReferences": "Video References",
     "VLLMOmniLatentMaskEditing": "Latent Mask Editing",
+    "VLLMOmniMiniMaxH3TemporalMask": "MiniMax-H3 Temporal Mask",
     # === Params ===
     "VLLMOmniARSampling": "AR Sampling Params",
     "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -8,7 +8,7 @@
 from comfy_api.input import AudioInput, VideoInput
 
 from .utils.api_client import VLLMOmniClient
-from .utils.latent_mask import _align_frame_count, _video_latent_t
+from .utils.latent_mask import _align_frame_count
 from .utils.logger import get_logger
 from .utils.models import lookup_model_spec
 from .utils.types import (
@@ -1133,13 +1133,13 @@ def build(self, images, source_fps, duration, mode, preserve_fraction=0.5):
         else:
             raise ValueError(f"Unknown temporal mask mode: {mode}")
         available = min(frames, math.floor(boundary * 24 + 1e-8))
+        # Snap the preserved prefix to whole VAE clips (5 + 17n frames) so every
+        # latent it maps to is fully preserved.
         prefix = 0 if available < 5 else 5 + 17 * ((available - 5) // 17)
-        preserved = _video_latent_t(prefix) if prefix else 0
-        total = _video_latent_t(frames)
-        mask = torch.ones(total, 1, 1)
-        mask[:preserved] = 0
+        # One slice per output frame; the server pools frames to the latent grid.
+        mask = torch.ones(frames, 1, 1)
+        mask[:prefix] = 0
         indices = torch.arange(frames, device=images.device)
         source_indices = (indices * (source_fps / 24)).floor().long().clamp(max=images.shape[0] - 1)
         preview_images = images.index_select(0, source_indices)
-        preview_mask = mask.index_select(0, torch.arange(frames) * total // frames)
-        return mask, 24.0, preview_images, preview_mask
+        return mask, 24.0, preview_images, mask

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -30,7 +30,7 @@
     video_to_base64,
     video_to_bytes,
 )
-from .latent_mask import scalar_mask_to_json, video_mask_to_grid_json
+from .latent_mask import scalar_mask_to_json, video_mask_to_json
 from .logger import get_logger, pretty_printer
 from .models import lookup_model_spec
 from .types import (
@@ -388,6 +388,15 @@ async def generate_video(
                 content_type="image/png",
             )
 
+        if len(keyframe_images) == 2:
+            for image_filename, image in keyframe_images:
+                form.add_field(
+                    "input_references",
+                    image_tensor_to_png_bytes(image, image_filename),
+                    filename=image_filename,
+                    content_type="image/png",
+                )
+
         # === latent-mask editing (MiniMax H3) ===
         if latent_edit is not None:
             source_video = latent_edit.get("source_video")
@@ -398,29 +407,24 @@ async def generate_video(
             if video_mask is None and audio_mask is None:
                 raise ValueError("Latent-mask editing requires at least one mask.")
 
-            video_mask_trivial = video_mask is None or bool((video_mask == 1.0).all().item())
-            audio_mask_trivial = audio_mask is None or audio_mask == 1.0
-            if not video_mask_trivial and source_video is None:
-                raise ValueError("A non-trivial video mask requires a source video.")
-            if not audio_mask_trivial and source_audio is None and source_video is None:
-                raise ValueError("A non-trivial audio mask requires a source audio or a source video with audio.")
-
             if source_video is not None:
                 form.add_field(
                     "source_video",
-                    video_to_bytes(source_video, "source.mp4"),
+                    video_to_bytes(source_video),
                     filename="source.mp4",
                     content_type="video/mp4",
                 )
             if source_audio is not None:
+                # The filename extension selects the audio codec (mp3), so it is
+                # meaningful even though aiohttp uses the explicit filename below.
                 form.add_field(
                     "source_audio",
                     audio_to_bytes(source_audio, "source_audio.mp3"),
                     filename="source_audio.mp3",
                     content_type="audio/mpeg",
                 )
             if video_mask is not None:
-                mask_json = video_mask_to_grid_json(video_mask, width=width, height=height, num_frames=num_frames)
+                mask_json = video_mask_to_json(video_mask)
                 form.add_field(
                     "video_noise_mask",
                     mask_json.encode("utf-8"),
@@ -435,15 +439,6 @@ async def generate_video(
                     content_type="application/json",
                 )
 
-        if len(keyframe_images) == 2:
-            for image_filename, image in keyframe_images:
-                form.add_field(
-                    "input_references",
-                    image_tensor_to_png_bytes(image, image_filename),
-                    filename=image_filename,
-                    content_type="image/png",
-                )
-
         # === model specific params. Either use a specialized builder, or add flattened fields as-is ===
         if model_params is not None:
             model_params = dict(model_params)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/latent_mask.py

@@ -1,51 +1,57 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+"""Client-side serialization of MiniMax-H3 latent-edit masks.
+
+A video mask is sent in frame space — a 2D spatial ``[H, W]`` (applied to every
+frame) or a 3D ``[T, H, W]`` (one slice per frame) — and the server resolves it
+to the latent grid. An audio mask is a scalar applied to all time steps.
+"""
+
 import json
+import math
 
 import torch
 import torch.nn.functional as F
 
+# The server rejects mask JSON above 8 MiB, so the spatial axes are
+# area-downsampled by the VAE spatial stride before upload. This depends only on
+# the stride, not on the H3 shape lattice, which the server still owns.
+_VAE_SPATIAL_STRIDE = 16
+_MASK_DECIMALS = 4
+
 
 def scalar_mask_to_json(value: float) -> str:
     if not 0.0 <= value <= 1.0:
         raise ValueError(f"mask value must be in [0, 1], got {value}")
     return str(value)
 
 
+def video_mask_to_json(mask: torch.Tensor) -> str:
+    """Serialize a frame-space video mask to JSON for the ``video_noise_mask`` field.
+
+    ``mask`` is a 2D spatial mask or a 3D frame-space mask. Its spatial axes are
+    area-downsampled by the VAE stride; the server resizes the result to the
+    latent grid, so the client neither floors the canvas nor maps frames to
+    latents.
+    """
+    if mask.ndim not in (2, 3):
+        raise ValueError(f"expected a 2D or 3D mask tensor, got {mask.ndim}D")
+    frames = mask.unsqueeze(0) if mask.ndim == 2 else mask
+    size = tuple(max(1, math.ceil(dim / _VAE_SPATIAL_STRIDE)) for dim in frames.shape[1:])
+    frames = F.interpolate(frames.unsqueeze(1).float(), size=size, mode="area").squeeze(1)
+    if mask.ndim == 2:
+        frames = frames.squeeze(0)
+    # Round in float64 so tolist() emits short decimals rather than float32 noise.
+    return json.dumps(frames.double().round(decimals=_MASK_DECIMALS).tolist(), separators=(",", ":"))
+
+
+# The temporal-mask node snaps its preserve boundary to the H3 frame lattice;
+# this mirrors the server shape planner.
 def _align_frame_count(frame_count: int) -> int:
     if frame_count <= 0:
         return 1
     current = int(frame_count)
     while current % 17 != 5:
         current += 1
     return current
-
-
-def _video_latent_t(frame_count: int) -> int:
-    if frame_count <= 5:
-        return 2
-    return ((int(frame_count) - 5) // 17) * 5 + 2
-
-
-def video_mask_to_grid(mask: torch.Tensor, *, width: int, height: int, num_frames: int) -> torch.Tensor:
-    tv = _video_latent_t(_align_frame_count(num_frames))
-    height = int(height) // 32 * 32
-    width = int(width) // 32 * 32
-    gh, gw = height // 16, width // 16
-
-    if mask.ndim == 2:
-        mask = mask.unsqueeze(0)
-    elif mask.ndim != 3:
-        raise ValueError(f"expected a 2D or 3D mask tensor, got {mask.ndim}D")
-
-    grid = F.interpolate(mask.unsqueeze(1).float(), size=(gh, gw), mode="area").squeeze(1)
-    if grid.shape[0] == 1:
-        grid = grid.expand(tv, gh, gw)
-    elif grid.shape[0] != tv:
-        grid = F.interpolate(grid.unsqueeze(0).unsqueeze(0), size=(tv, gh, gw), mode="nearest").squeeze(0).squeeze(0)
-    return grid
-
-
-def video_mask_to_grid_json(mask: torch.Tensor, *, width: int, height: int, num_frames: int) -> str:
-    return json.dumps(video_mask_to_grid(mask, width=width, height=height, num_frames=num_frames).tolist())

apps/ComfyUI-vLLM-Omni/docs/wf05-h3-latent-editing.md

@@ -12,6 +12,6 @@ Edit a source video with MiniMax-H3 using spatial or temporal masks. The workflo
 Open **vLLM-Omni MiniMax-H3 Latent Mask Editing.json** from the example workflows and set the Generate Video URL/model for the case you want to run.
 
 1. Upload your source video in Load Video.
-2. Adjust the case's mask. For removal and inpainting, set Mask Canvas to the video dimensions, Mask Rectangle to the region size, and Combine Masks x/y to its position. For continuation and extension, set MiniMax-H3 Temporal Mask's mode and duration; continuation also uses `preserve_fraction`. Match the duration in Generate Video, and use a duration longer than the source for extension. Mask values are `0` to preserve and `1` to regenerate.
+2. Adjust the case's mask. For removal and inpainting, set Mask Canvas to the video dimensions, Mask Rectangle to the region size, and Combine Masks x/y to its position, or replace that chain with Load Image (as Mask) for an arbitrary-shape mask. For continuation and extension, set MiniMax-H3 Temporal Mask's mode and duration; continuation also uses `preserve_fraction`. Match the duration in Generate Video, and use a duration longer than the source for extension. Mask values are `0` to preserve and `1` to regenerate.
 3. Adjust the prompt to describe the desired content. Set the output dimensions and duration in Generate Video. If editing audio, update the sound description and audio mask (`0` preserves, `1` regenerates the whole clip).
 4. Select the case's Mask Preview Save Video node and choose **Execute to selected output nodes** to preview the mask as a red overlay without inference, or select its Result Save Video node to generate the edited video. Running the entire workflow generates all four cases.