[Model] MiniMax-H3: accept raw video masks and resolve the latent grid server-side (#7947)

Signed-off-by: chen hongwei <1792043268@qq.com>
Signed-off-by: Wang Zhipeng <wangzhipeng628@gmail.com>
Co-authored-by: Wang Zhipeng <wangzhipeng628@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

docs/serving/videos_api.md

@@ -237,13 +237,22 @@ At least one mask is required. A nontrivial `video_noise_mask` requires
 `source_audio` or a `source_video` with an audio stream. Mask values are in
 `[0, 1]`: `0` preserves the source, `1` regenerates it, and fractional values
 blend the two behaviors. Exact all-one masks are no-ops and do not require a
-source. Source uploads without a mask are rejected. Masks may be a JSON scalar
-or arrays matching the H3 latent/token grid; pixel-resolution masks must be
-resized or pooled by the client before upload.
+source. Source uploads without a mask are rejected. Masks may be a JSON scalar,
+an array matching the H3 latent/token grid, or (video only) a frame-space
+array that the server resizes to the latent grid.
 
 For an aligned output of `F` frames at `W x H`, the video latent grid is
 `[Tv, H/16, W/16]`, where `Tv = 2 + 5 * ((F - 5) / 17)`. The video mask may be
-a scalar, a flat token vector, `[Tv, H/32, W/32]`, or the full latent grid. For
+a scalar, a flat token vector, `[Tv, H/32, W/32]`, the full latent grid, a
+spatial `[h, w]` mask applied to every frame, or a frame-space `[T, h, w]` mask
+with one slice per 24 fps output frame. A shape that matches the token or full
+latent grid is always read as that grid. Frame-space masks are area-resized to
+`[H/16, W/16]`; along time, a mask shorter than `F` is padded with its last
+slice (as the source video is) and a longer one is truncated, then each latent
+takes the maximum over the frames the causal video VAE folds into it, so any
+regenerated frame regenerates its latent. Because the 8 MiB limit below applies
+to the JSON file, clients should area-downsample frame-space masks by 16
+spatially before upload. For
 the model input, timestep, and velocity, a full-grid mask is max-pooled over
 each 2x2 spatial token and fractional values are rounded upward to 1/256
 levels. The final x0 restore uses the original, unquantized mask, so full-grid

tests/model_executor/models/test_minimax_h3_encoder_payload.py

@@ -28,7 +28,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def test_edit_mask_boundary_canonicalizes_supported_request_shapes() -> None:
     video_token = torch.arange(168, dtype=torch.float32).reshape(7, 4, 6) / 168
     video_full = video_token.repeat_interleave(2, dim=1).repeat_interleave(2, dim=2)
-    video_kwargs = {"latent_t": 7, "latent_h": 8, "latent_w": 12}
+    video_kwargs = {"latent_t": 7, "latent_h": 8, "latent_w": 12, "num_frames": 22}
     for value in (video_token.flatten()[None, None], video_token[None], video_full[None]):
         torch.testing.assert_close(_canonical_video_edit_mask(value, **video_kwargs), video_full)
     torch.testing.assert_close(
@@ -46,14 +46,40 @@ def test_edit_mask_boundary_canonicalizes_supported_request_shapes() -> None:
     )
 
 
+def test_edit_mask_resizes_raw_spatial_and_frame_space_masks() -> None:
+    video_kwargs = {"latent_t": 7, "latent_h": 8, "latent_w": 12, "num_frames": 22}
+
+    spatial = torch.full((16, 24), 0.5)
+    spatial_grid = _canonical_video_edit_mask(spatial, **video_kwargs)
+    assert spatial_grid.shape == (7, 8, 12)
+    torch.testing.assert_close(spatial_grid, torch.full((7, 8, 12), 0.5))
+
+    frame_space = torch.zeros(22, 16, 24)
+    frame_space[7] = 1.0  # frame 7 -> token 2
+    frame_grid = _canonical_video_edit_mask(frame_space, **video_kwargs)
+    assert frame_grid.shape == (7, 8, 12)
+    assert bool(torch.all(frame_grid[2] > 0.5).item())
+    assert bool(torch.all(frame_grid[0] == 0.0).item())
+
+
 @pytest.mark.parametrize(
     ("canonicalize", "value", "kwargs", "message"),
     [
-        (_canonical_video_edit_mask, [True], {"latent_t": 7, "latent_h": 8, "latent_w": 12}, "booleans"),
+        (
+            _canonical_video_edit_mask,
+            [True],
+            {"latent_t": 7, "latent_h": 8, "latent_w": 12, "num_frames": 22},
+            "booleans",
+        ),
         (_canonical_audio_edit_mask, float("nan"), {"audio_t": 37}, "finite"),
         (_canonical_audio_edit_mask, 10**400, {"audio_t": 37}, "finite"),
         (_canonical_audio_edit_mask, 1.01, {"audio_t": 37}, r"\[0, 1\]"),
-        (_canonical_video_edit_mask, [0.0, 0.5], {"latent_t": 7, "latent_h": 8, "latent_w": 12}, "shape"),
+        (
+            _canonical_video_edit_mask,
+            [0.0, 0.5],
+            {"latent_t": 7, "latent_h": 8, "latent_w": 12, "num_frames": 22},
+            "shape",
+        ),
     ],
 )
 def test_edit_mask_boundary_rejects_invalid_values(canonicalize, value, kwargs, message) -> None:
@@ -592,3 +618,47 @@ def test_driving_audio_does_not_override_fl2va_task_inference():
     assert prepared.media.task == "fl2va"
     assert prepared.media.audio_mode == "lock_source"
     assert prepared.condition_labels == [("image", 1)]
+
+
+def test_audio_edit_mask_rejects_raw_temporal_lengths() -> None:
+    from vllm_omni.errors import OmniClientError
+
+    for value in (torch.zeros(100), torch.zeros(1, 100), torch.zeros(2, 100), torch.zeros(3, 37)):
+        with pytest.raises(OmniClientError, match="shape"):
+            _canonical_audio_edit_mask(value, audio_t=37)
+
+
+def test_temporal_group_max_pool_rejects_count_mismatch() -> None:
+    from vllm_omni.errors import OmniClientError
+    from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
+        _temporal_group_max_pool,
+    )
+
+    with pytest.raises(OmniClientError, match="does not map"):
+        _temporal_group_max_pool(torch.zeros(34, 1, 1), latent_t=6)
+
+
+def test_video_edit_mask_aligns_frame_space_mask_to_num_frames() -> None:
+    video_kwargs = {"latent_t": 7, "latent_h": 8, "latent_w": 12, "num_frames": 22}
+    mask = torch.zeros(17, 16, 24)
+    mask[0] = 1.0
+    grid = _canonical_video_edit_mask(mask, **video_kwargs)
+    assert grid.shape == (7, 8, 12)
+    assert bool(torch.all(grid[0] > 0.5).item())
+    assert bool(torch.all(grid[1] == 0.0).item())
+
+
+def test_temporal_group_max_pool_causal_boundaries() -> None:
+    from vllm_omni.model_executor.models.minimax_h3.encoder_processing import (
+        _temporal_group_max_pool,
+    )
+
+    def lit(frame: int) -> list[int]:
+        mask = torch.zeros(34, 1, 1)
+        mask[frame] = 1.0
+        out = _temporal_group_max_pool(mask, latent_t=7)
+        return [i for i in range(out.shape[0]) if out[i, 0, 0].item() > 0]
+
+    assert lit(0) == [0]
+    assert lit(4) == [1]
+    assert lit(5) == [2]

vllm_omni/model_executor/models/minimax_h3/encoder_processing.py

@@ -294,8 +294,9 @@ def _canonical_video_edit_mask(
     latent_t: int,
     latent_h: int,
     latent_w: int,
+    num_frames: int,
 ) -> torch.Tensor:
-    """Normalize every accepted request shape to one full latent grid."""
+    """Normalize every accepted request shape, including raw frame-space masks, to one full latent grid."""
     mask = _edit_mask(value, name="video_noise_mask")
     token_shape = (latent_t, latent_h // 2, latent_w // 2)
     full_shape = (latent_t, latent_h, latent_w)
@@ -313,12 +314,63 @@ def _canonical_video_edit_mask(
         if not candidate.ndim or candidate.shape[0] != 1:
             break
         candidate = candidate.squeeze(0)
+    if candidate.ndim in (2, 3):
+        return _resize_video_edit_mask(
+            candidate,
+            latent_t=latent_t,
+            latent_h=latent_h,
+            latent_w=latent_w,
+            num_frames=num_frames,
+        )
     raise OmniClientError(
         "MiniMax H3 video_noise_mask shape must be "
-        f"scalar, ({row_count},), {token_shape}, or {full_shape}; got {tuple(mask.shape)}"
+        f"scalar, ({row_count},), {token_shape}, {full_shape}, [H, W], or [T, H, W]; got {tuple(mask.shape)}"
     )
 
 
+def _resize_video_edit_mask(
+    mask: torch.Tensor,
+    *,
+    latent_t: int,
+    latent_h: int,
+    latent_w: int,
+    num_frames: int,
+) -> torch.Tensor:
+    """Resize a raw ``[H, W]`` or frame-space ``[T, H, W]`` mask to ``[latent_t, latent_h, latent_w]``."""
+    if mask.ndim == 2:
+        mask = mask.unsqueeze(0)
+    grid = torch.nn.functional.interpolate(mask.unsqueeze(1), size=(latent_h, latent_w), mode="area").squeeze(1)
+    if grid.shape[0] == 1:
+        return grid.expand(latent_t, latent_h, latent_w).contiguous()
+    # Fit to the output length the way the source video is fitted, so the edit boundary does not drift.
+    if grid.shape[0] < num_frames:
+        grid = torch.cat([grid, grid[-1:].expand(num_frames - grid.shape[0], *grid.shape[1:])], dim=0)
+    return _temporal_group_max_pool(grid[:num_frames], latent_t=latent_t)
+
+
+def _temporal_group_max_pool(mask: torch.Tensor, *, latent_t: int) -> torch.Tensor:
+    """Max-pool frames onto the latents the causal video VAE encodes them into."""
+    clip = 17
+    frame_count = mask.shape[0]
+    num_chunks = -(-frame_count // clip)
+    padded = num_chunks * clip
+    if frame_count < padded:
+        mask = torch.cat([mask, mask[-1:].expand(padded - frame_count, *mask.shape[1:])], dim=0)
+    tokens: list[torch.Tensor] = []
+    for c in range(num_chunks):
+        chunk = mask[c * clip : (c + 1) * clip]
+        # Per 17-frame clip, token 0 covers frame 0 and token k covers frames 4k-3..4k; the VAE drops the last 3.
+        tokens.append(chunk[0:1].amax(dim=0))
+        for k in range(1, 5):
+            tokens.append(chunk[4 * k - 3 : 4 * k + 1].amax(dim=0))
+    result = torch.stack(tokens)[:-3]
+    if result.shape[0] != latent_t:
+        raise OmniClientError(
+            f"MiniMax H3 video_noise_mask temporal depth {frame_count} does not map to {latent_t} latents"
+        )
+    return result
+
+
 def _canonical_audio_edit_mask(value: Any, *, audio_t: int) -> torch.Tensor:
     """Normalize every accepted request shape to a channel-major grid."""
     mask = _edit_mask(value, name="audio_noise_mask")
@@ -476,6 +528,7 @@ def prepare_encoder_inputs(
             latent_t=latent_t,
             latent_h=height // 16,
             latent_w=width // 16,
+            num_frames=num_frames,
         )
         if raw_video_edit_mask is not None
         else None