diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index cf7b9136a..e5b87c7bf 100644 --- a/docs/serving/videos_api.md +++ b/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 diff --git a/tests/model_executor/models/test_minimax_h3_encoder_payload.py b/tests/model_executor/models/test_minimax_h3_encoder_payload.py index c3f64dcd6..3d07a969e 100644 --- a/tests/model_executor/models/test_minimax_h3_encoder_payload.py +++ b/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] diff --git a/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py b/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py index 08c298559..4fe517b50 100644 --- a/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py +++ b/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