From 5a93ec1b4e990f42973954d63809f932e699426a Mon Sep 17 00:00:00 2001 From: xiaoyu-xyz <45517337+xiaoyu-xyz@users.noreply.github.com> Date: Sun, 20 Sep 2026 11:54:44 +0800 Subject: [PATCH] [Model][Frontend] MiniMax-H3: Add latent-mask editing (#7465) Signed-off-by: xiaoyu <1259084489@qq.com> --- docs/serving/videos_api.md | 47 +++ .../minimax_h3/test_minimax_h3_contract.py | 1 + .../test_minimax_h3_latent_edit_pipeline.py | 258 +++++++++++++++ .../minimax_h3/test_minimax_h3_latent_mask.py | 81 +++++ .../test_minimax_h3_step_execution.py | 119 ++++++- .../openai_api/test_video_server.py | 270 ++++++++++++++- .../models/test_minimax_h3_encoder_payload.py | 125 ++++++- vllm_omni/diffusion/model_metadata.py | 6 + .../models/minimax_h3/batched_packing.py | 28 +- .../models/minimax_h3/denoise_loop.py | 138 ++++++-- .../models/minimax_h3/latent_mask.py | 183 +++++++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 220 ++++++++++++- vllm_omni/entrypoints/openai/api_server.py | 9 +- vllm_omni/entrypoints/openai/serving_video.py | 66 +++- .../openai/video/generation/helpers.py | 273 +++++++++++++-- .../models/minimax_h3/conditioning.py | 311 +++++++++++++++++- .../models/minimax_h3/encoder_processing.py | 224 ++++++++++++- .../models/minimax_h3/reference_video.py | 49 ++- 18 files changed, 2286 insertions(+), 122 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_latent_edit_pipeline.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_latent_mask.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/latent_mask.py diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index f48f2ca99..da0fa1fd3 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -73,6 +73,10 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4 | `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` | | `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` | | `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs | +| `source_video` | file | null | MiniMax H3 latent-edit source video (`.mp4` or `.mov`, up to 512 MiB) | +| `source_audio` | file | null | Optional MiniMax H3 latent-edit source audio (`.wav` or `.mp3`, up to 512 MiB) | +| `video_noise_mask` | file | null | UTF-8 JSON MiniMax H3 video mask; `0` preserves and `1` regenerates a token | +| `audio_noise_mask` | file | null | UTF-8 JSON MiniMax H3 audio mask; `0` preserves and `1` regenerates a token | | `width` | integer | model default | Output video width | | `height` | integer | model default | Output video height | | `num_frames` | integer | 1 | Number of generated frames | @@ -214,6 +218,49 @@ upstream vLLM's media URL allowlist protection. Deployments should therefore treat remote media URLs as untrusted and choose this setting as part of their URL access policy. +### MiniMax H3 Latent-Mask Editing + +MiniMax H3 accepts request-scoped source media and video/audio noise masks. +At least one mask is required. A nontrivial `video_noise_mask` requires +`source_video`, while a nontrivial `audio_noise_mask` requires either +`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. + +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 +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 +masks retain cell-level preservation inside a model token. The audio length is +`Ta = round(F * 40 / 24)`, and its mask may be a scalar, `[Ta]`, `[2, Ta]`, or +a flat `2 * Ta` vector. If source audio is short, its missing latent tail is +forced to mask value `1` so H3 generates that portion. Each mask must be sent +as a UTF-8 JSON file part and is limited to 8 MiB. A file may contain either a +JSON scalar or an array. + +```bash +curl -s http://localhost:8091/v1/videos/sync \ + -F "prompt=Partially restyle the complete clip and soundtrack" \ + -F 'extra_params={"task":"t2va","duration":4.0,"aspect_ratio":"16:9"}' \ + -F "source_video=@source.mp4;type=video/mp4" \ + -F "source_audio=@source.wav;type=audio/wav" \ + -F "video_noise_mask=@video-mask.json;type=application/json" \ + -F "audio_noise_mask=@audio-mask.json;type=application/json" \ + -o edited.mp4 +``` + +For example, each mask file in the request above may contain the JSON scalar +`0.5`. + +The server streams source files to temporary request-scoped storage and removes +them after synchronous or asynchronous generation finishes. Models that do not +declare latent-mask editing support reject these fields. + ### Speech-to-Video For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 49197d315..d6c70dfaf 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -1150,6 +1150,7 @@ def test_minimax_h3_advertises_the_official_ref2va_image_limit(): from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9 + assert get_diffusion_model_metadata("MiniMaxH3Pipeline").supports_latent_mask_editing def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch): diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_latent_edit_pipeline.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_latent_edit_pipeline.py new file mode 100644 index 000000000..a57a7fe09 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_latent_edit_pipeline.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pipeline-level MiniMax H3 latent initialization tests.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_prepare_edit_video_clones_the_last_frame_with_lossless_rgb(monkeypatch, tmp_path): + from vllm_omni.model_executor.models.minimax_h3 import reference_video + + source = tmp_path / "source.webm" + commands: list[list[str]] = [] + monkeypatch.setattr( + reference_video, + "_probe_video", + lambda _path: {"frame_count": 12, "audio_codecs": ("opus",)}, + ) + monkeypatch.setattr( + reference_video.subprocess, + "run", + lambda command, **_kwargs: commands.append(command), + ) + + prepared = reference_video.prepare_edit_video( + source, + target_width=448, + target_height=256, + target_frame_count=73, + workdir=str(tmp_path / "prepared"), + ) + + assert prepared["input_has_audio"] is True + assert prepared["frame_count"] == 73 + command = commands[0] + assert command[command.index("-vf") + 1].endswith("tpad=stop_mode=clone:stop=-1") + assert command[command.index("-c:v") + 1] == "libx264rgb" + assert command[command.index("-pix_fmt") + 1] == "rgb24" + + +@pytest.mark.parametrize("source_audio", [None, "override.wav"]) +def test_encoder_preparation_routes_edit_sources_without_qwen_inputs( + monkeypatch, + tmp_path, + source_audio, +): + from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing + + latent_t, audio_t = 7, 37 + video_clean = torch.full((42, 96), 3.0) + audio_source = torch.zeros(4, 32) + audio_source[:, 0] = torch.tensor([1.0, 2.0, 11.0, 12.0]) + waveform = torch.zeros(3_200) + frames = torch.zeros(22, 64, 96, 3, dtype=torch.uint8).numpy() + calls: dict[str, Any] = {} + + def prepare_video(value, **kwargs): + calls["prepare"] = (value, kwargs) + return { + "original_path": "source.mov", + "prepared_path": str(tmp_path / "prepared.mp4"), + "input_has_audio": True, + } + + def load_video_audio(value, **kwargs): + calls["video_audio"] = (value, kwargs) + return waveform, 32_000 + + def load_audio_file(value): + calls["audio_file"] = value + return waveform, 32_000 + + class VideoVAE: + def encode_video(self, value): + calls["video_encode_shape"] = value.shape + return video_clean.clone(), (latent_t, 4, 6) + + class AudioVAE: + def encode_waveform(self, value, sample_rate): + calls["audio_encode"] = (value.shape, sample_rate) + return audio_source.clone(), 2 + + monkeypatch.setattr(processing, "resolve_minimax_h3_shape", lambda *_args: (64, 96, 22, latent_t, audio_t)) + monkeypatch.setattr(processing, "prepare_edit_video", prepare_video) + monkeypatch.setattr(processing, "load_video_frames", lambda _path: frames) + monkeypatch.setattr(processing, "load_video_audio", load_video_audio) + monkeypatch.setattr(processing, "load_audio_file", load_audio_file) + + prepared = processing.prepare_encoder_inputs( + { + "prompt": "edit the clip", + "multi_modal_data": { + "source_video": "source.mov", + "source_audio": source_audio, + "video_noise_mask": torch.zeros(1, 1, latent_t, 4, 6), + "audio_noise_mask": torch.tensor([0.0, 0.5, *([0.0] * (audio_t - 2))]).reshape(1, 1, -1), + }, + }, + SimpleNamespace(extra_args={"task": "t2va"}), + ) + media = type(prepared.media).from_mm_tensors( + prepared.media.to_mm_tensors(), + prepared.media.to_metadata(), + ) + conditioning = processing.encode_media( + media, + video_vae=VideoVAE(), + audio_vae=AudioVAE(), + emit_conditioning=True, + ) + + assert prepared.images == [] + assert prepared.qwen_videos == [] + assert prepared.condition_labels == [] + assert prepared.media.images == () + assert prepared.media.videos == () + assert calls["prepare"][0] == "source.mov" + assert calls["prepare"][1]["target_width"] == 96 + assert calls["prepare"][1]["target_height"] == 64 + assert calls["prepare"][1]["target_frame_count"] == 22 + if source_audio is None: + assert calls["video_audio"] == ( + "source.mov", + {"duration_seconds": pytest.approx(22 / 24)}, + ) + assert "audio_file" not in calls + else: + assert calls["audio_file"] == "override.wav" + assert "video_audio" not in calls + + assert conditioning is not None + torch.testing.assert_close(conditioning.video_edit_clean_rows, video_clean) + assert conditioning.video_edit_mask.shape == (latent_t, 4, 6) + torch.testing.assert_close(conditioning.video_edit_mask, torch.zeros(latent_t, 4, 6)) + assert conditioning.audio_edit_source_t == 2 + assert conditioning.audio_edit_mask.shape == (2, audio_t) + torch.testing.assert_close( + conditioning.audio_edit_clean_rows[:, 0], + torch.tensor([1.0, 2.0, *([0.0] * (audio_t - 2)), 11.0, 12.0, *([0.0] * (audio_t - 2))]), + ) + torch.testing.assert_close( + conditioning.audio_edit_mask, + torch.tensor([0.0, 0.5, *([0.0] * (audio_t - 2))]).repeat(2, 1), + ) + + from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MiniMaxH3EncoderConditioning, + MiniMaxH3TextConditioning, + ) + + combined = MiniMaxH3EncoderConditioning.from_components( + MiniMaxH3TextConditioning( + torch.zeros(1, 5120, dtype=torch.bfloat16), + torch.ones(1, dtype=torch.int64), + ), + conditioning, + ) + round_trip = MiniMaxH3EncoderConditioning.from_omni_payload(combined.to_omni_payload()) + torch.testing.assert_close(round_trip.video_edit_clean_rows, video_clean) + torch.testing.assert_close(round_trip.video_edit_mask, torch.zeros(latent_t, 4, 6)) + torch.testing.assert_close(round_trip.audio_edit_clean_rows, conditioning.audio_edit_clean_rows) + torch.testing.assert_close( + round_trip.audio_edit_mask, + torch.tensor([0.0, 0.5, *([0.0] * (audio_t - 2))]).repeat(2, 1), + ) + assert round_trip.audio_edit_source_t == 2 + + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline._fasth3 = None + pipeline._active_turbo_spec = lambda _sampling: None + pipeline._has_active_native_lora = lambda _sampling: False + pipeline._resolve_task = lambda task, **_kwargs: task + pipeline._resolve_sigma_positions = lambda _task, _sampling: (None, 3) + pipeline._quality_policy = SimpleNamespace(resolve=lambda **_kwargs: SimpleNamespace(cache_dit=None)) + pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda _plan: None) + context = pipeline._prepare_encoder_conditioning_inputs( + round_trip, + SimpleNamespace(extra_args={}, quality=None, seed=17, num_outputs_per_prompt=1), + ) + torch.testing.assert_close(context["video_edit_clean_rows"], video_clean) + torch.testing.assert_close(context["video_edit_mask_rows"], torch.zeros(42)) + assert context["video_edit_restore_mask_rows"].shape == (42, 96) + torch.testing.assert_close( + context["audio_edit_mask_rows"], + torch.tensor([0.0, 0.5, *([1.0] * (audio_t - 2))]).repeat(2), + ) + torch.testing.assert_close(context["audio_edit_mask_rows"], context["audio_edit_restore_mask_rows"]) + + +def _build_inputs(*, video_mask: torch.Tensor, audio_mask: torch.Tensor): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + latent_t, latent_h, latent_w, audio_t = 2, 4, 6, 3 + video_clean = torch.linspace(-1.0, 1.0, latent_t * (latent_h // 2) * (latent_w // 2) * 96).reshape(-1, 96) + audio_clean = torch.linspace(-0.5, 0.5, audio_t * 2 * 32).reshape(-1, 32) + inputs = pipeline._build_denoise_inputs( + task="t2va", + text_embeddings=torch.zeros(2, 4), + text_tags=torch.ones(2, dtype=torch.long), + seed=17, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + num_frames=22, + num_steps=3, + video_shift=12.0, + audio_shift=3.0, + base_schedule=None, + visual_condition=None, + visual_condition_shape=None, + audio_condition=None, + ref_audio_t=None, + video_edit_clean_rows=video_clean, + video_edit_mask_rows=video_mask, + video_edit_restore_mask_rows=video_mask, + audio_edit_clean_rows=audio_clean, + audio_edit_mask_rows=audio_mask, + audio_edit_restore_mask_rows=audio_mask, + ) + return pipeline, inputs, video_clean, audio_clean + + +def test_build_denoise_inputs_uses_h3_source_anchors(): + pipeline, inputs, video_clean, audio_clean = _build_inputs( + video_mask=torch.zeros(12), + audio_mask=torch.zeros(6), + ) + video_edit = inputs["video_edit"] + audio_edit = inputs["audio_edit"] + assert video_edit is not None and audio_edit is not None + initial_video, _ = pipeline._initial_noise( + seed=17, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + ) + torch.testing.assert_close(video_edit.clean_rows, video_clean) + torch.testing.assert_close(video_edit.anchor_rows, 0.999 * video_clean + 0.001 * initial_video) + torch.testing.assert_close(audio_edit.clean_rows, audio_clean) + torch.testing.assert_close(audio_edit.anchor_rows, audio_clean) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_latent_mask.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_latent_mask.py new file mode 100644 index 000000000..91665454f --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_latent_mask.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Focused CPU tests for MiniMax H3 latent-mask editing.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_video_full_grid_quantization_and_cell_level_restore(): + from vllm_omni.diffusion.models.minimax_h3.latent_mask import minimax_h3_video_edit_masks + + shape = dict(latent_t=2, latent_h=4, latent_w=6) + full = torch.arange(48, dtype=torch.float32).reshape(2, 4, 6) / 48 + cells = full.reshape(2, 2, 2, 3, 2).permute(0, 1, 3, 2, 4) + parsed = minimax_h3_video_edit_masks(full, **shape) + expected_model = torch.ceil(cells.amax(dim=(3, 4)).reshape(-1) * 256) / 256 + expected_restore = cells.unsqueeze(3).expand(2, 2, 3, 24, 2, 2).reshape(12, 96) + torch.testing.assert_close(parsed.model_mask_rows, expected_model) + torch.testing.assert_close(parsed.restore_mask_rows, expected_restore) + + +def test_audio_grid_uses_channel_major_order_and_preserves_raw_values(): + from vllm_omni.diffusion.models.minimax_h3.latent_mask import minimax_h3_audio_edit_masks + + grid = torch.tensor([[0.0, 0.1, 1.0], [0.25, 0.5, 0.75]]) + expected = grid.flatten() + parsed = minimax_h3_audio_edit_masks(grid, audio_t=3) + torch.testing.assert_close(parsed.model_mask_rows, torch.ceil(expected * 256) / 256) + torch.testing.assert_close(parsed.restore_mask_rows, expected) + + +def test_zero_fractional_one_semantics_and_cell_level_restore(): + from vllm_omni.diffusion.models.minimax_h3.latent_mask import ( + MiniMaxH3LatentEdit, + minimax_h3_video_edit_masks, + ) + + clean = torch.full((3, 2), 10.0) + anchor = torch.tensor([[8.0, 8.0], [6.0, 6.0], [4.0, 4.0]]) + state = torch.tensor([[2.0, 2.0], [4.0, 4.0], [6.0, 6.0]]) + velocity = torch.full((3, 2), 2.0) + mask = torch.tensor([0.0, 0.5, 1.0]) + edit = MiniMaxH3LatentEdit.from_rows(clean, anchor, mask, mask) + assert edit is not None + + model_rows, target_timesteps = edit.prepare(state, 0.25, 0.999, sigma=0.75) + torch.testing.assert_close(model_rows, torch.lerp(anchor, state, mask[:, None])) + torch.testing.assert_close( + target_timesteps, + torch.tensor([0.999, 0.625, 0.25]), + ) + predicted = model_rows + 0.75 * velocity * mask[:, None] + torch.testing.assert_close(edit.x0(model_rows, velocity, 0.25), torch.lerp(clean, predicted, mask[:, None])) + + # A single active cell drives the full model token, while raw restoration + # preserves the other three cells independently in every packed channel. + parsed = minimax_h3_video_edit_masks( + torch.tensor([[[1.0, 0.0], [0.0, 0.0]]]), + latent_t=1, + latent_h=2, + latent_w=2, + ) + restore = torch.tensor([1.0, 0.0, 0.0, 0.0]).repeat(24).reshape(1, 96) + torch.testing.assert_close(parsed.model_mask_rows, torch.ones(1)) + torch.testing.assert_close(parsed.restore_mask_rows, restore) + + cell_edit = MiniMaxH3LatentEdit.from_rows( + torch.full((1, 96), 10.0), + torch.full((1, 96), 8.0), + parsed.model_mask_rows, + parsed.restore_mask_rows, + ) + assert cell_edit is not None + cell_state = torch.full((1, 96), 2.0) + cell_velocity = torch.full((1, 96), 4.0) + expected = torch.where(restore.bool(), cell_state + 0.75 * cell_velocity, cell_edit.clean_rows) + torch.testing.assert_close(cell_edit.x0(cell_state, cell_velocity, 0.25), expected) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index e202c0d0e..77c66e8a0 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -27,6 +27,9 @@ class _SegmentMeanModel: assignment shows up as a numeric difference instead of passing silently. """ + def __init__(self): + self.calls: list[dict[str, torch.Tensor]] = [] + def __call__(self, **kwargs): x = kwargs["x"][0] audio_x = kwargs["audio_x"][0] @@ -45,6 +48,14 @@ class _SegmentMeanModel: pooled_audio[start:stop] = audio_x[start:stop].mean(dim=0, keepdim=True) row_timesteps = kwargs["unique_timesteps"][kwargs["inverse_indices"]].unsqueeze(-1) + self.calls.append( + { + "video_rows": x[img_pos].clone(), + "audio_rows": audio_x[audio_pos].clone(), + "video_timesteps": row_timesteps[img_pos, 0].clone(), + "audio_timesteps": row_timesteps[audio_pos, 0].clone(), + } + ) video = (pooled_video + 0.5 * x + row_timesteps)[img_pos] audio = (pooled_audio + 0.5 * audio_x + row_timesteps)[audio_pos] if not kwargs.get("skip_mask_out_condition", False): @@ -156,9 +167,76 @@ def test_step_execution_matches_request_mode_denoise_loop(): torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio) +def test_step_execution_matches_request_mode_with_latent_edits(): + """Masked model rows, row timesteps, and scheduler math match both paths.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop + from vllm_omni.diffusion.models.minimax_h3.latent_mask import MiniMaxH3LatentEdit + + model = _SegmentMeanModel() + branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=51) + sigmas_video = _sigmas(5, 12.0) + sigmas_audio = _sigmas(5, 3.0) + video_clean = torch.linspace(-1.0, 1.0, video_rows.numel()).reshape_as(video_rows) + audio_clean = torch.linspace(1.0, -1.0, audio_rows.numel()).reshape_as(audio_rows) + video_mask = torch.linspace(0.0, 1.0, video_rows.shape[0]) + video_edit = MiniMaxH3LatentEdit.from_rows( + video_clean, + 0.999 * video_clean + 0.001 * video_rows, + video_mask, + video_mask, + ) + audio_mask = torch.tensor([0.0, 0.2, 0.4, 0.6, 0.8, 1.0]) + audio_edit = MiniMaxH3LatentEdit.from_rows( + audio_clean, + audio_clean, + audio_mask, + audio_mask, + ) + + reference_video, reference_audio = minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=video_rows, + initial_audio_rows=audio_rows, + keyframe_cond_rows=None, + video_edit=video_edit, + audio_edit=audio_edit, + sigmas_video=sigmas_video, + sigmas_audio=sigmas_audio, + device=torch.device("cpu"), + ) + request_first_call = model.calls[0] + torch.testing.assert_close(request_first_call["video_rows"], video_edit.model_rows(video_rows)) + torch.testing.assert_close(request_first_call["audio_rows"], audio_edit.model_rows(audio_rows)) + torch.testing.assert_close( + request_first_call["video_timesteps"], + video_edit.target_timesteps(1.0 - sigmas_video[0], 0.999, sigma=sigmas_video[0]), + ) + torch.testing.assert_close( + request_first_call["audio_timesteps"], + audio_edit.target_timesteps(1.0 - sigmas_audio[0], 1.0, sigma=sigmas_audio[0]), + ) + model.calls.clear() + + state = _make_state("req-edit", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + state.extra[mod._STEP_VIDEO_EDIT] = video_edit + state.extra[mod._STEP_AUDIO_EDIT] = audio_edit + pipeline = _step_pipeline(model) + while not state.denoise_completed: + velocity = pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state]) + pipeline.step_scheduler(state, velocity) + + for name, expected in request_first_call.items(): + torch.testing.assert_close(model.calls[0][name], expected) + torch.testing.assert_close(state.latents, reference_video) + torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio) + + def test_batched_step_execution_matches_independent_requests(): """Two co-batched requests must land where they would have landed alone.""" from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + from vllm_omni.diffusion.models.minimax_h3.latent_mask import MiniMaxH3LatentEdit model = _SegmentMeanModel() specs = [ @@ -171,20 +249,37 @@ def test_batched_step_execution_matches_independent_requests(): pipeline = _step_pipeline(model) + def make_state(request_id: str, index: int): + branch, video_rows, audio_rows = _make_branch(**specs[index]) + sigmas_video, sigmas_audio = schedules[index] + state = _make_state(request_id, model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + video_mask = torch.linspace(0.1 * index, 0.8 + 0.1 * index, video_rows.shape[0]) + video_edit = MiniMaxH3LatentEdit.from_rows( + torch.full_like(video_rows, 10.0 + index), + torch.full_like(video_rows, 20.0 + index), + video_mask, + video_mask, + ) + audio_mask = torch.linspace(0.2 - 0.1 * index, 1.0 - 0.1 * index, audio_rows.shape[0]) + audio_edit = MiniMaxH3LatentEdit.from_rows( + torch.full_like(audio_rows, 30.0 + index), + torch.full_like(audio_rows, 40.0 + index), + audio_mask, + audio_mask, + ) + state.extra[mod._STEP_VIDEO_EDIT] = video_edit + state.extra[mod._STEP_AUDIO_EDIT] = audio_edit + return state + alone: list[tuple[torch.Tensor, torch.Tensor]] = [] - for spec, (sigmas_video, sigmas_audio) in zip(specs, schedules): - branch, video_rows, audio_rows = _make_branch(**spec) - state = _make_state("solo", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + for index in range(len(specs)): + state = make_state("solo", index) while not state.denoise_completed: pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state])) alone.append((state.latents, state.extra[mod._STEP_AUDIO_ROWS])) - states = [] - for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)): - branch, video_rows, audio_rows = _make_branch(**spec) - if index == 0: - assert branch.used_len == branch.seq_len - states.append(_make_state(f"req-{index}", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)) + states = [make_state(f"req-{index}", index) for index in range(len(specs))] + assert states[0].extra[mod._STEP_BRANCH].used_len == states[0].extra[mod._STEP_BRANCH].seq_len active = list(states) while active: @@ -291,6 +386,8 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames): from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8) + video_edit = object() + audio_edit = object() sigmas_video = _sigmas(6, 12.0) sigmas_audio = _sigmas(6, 3.0) context = { @@ -326,6 +423,8 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames): "audio_rows": audio_rows, "cond_anchor": None, "audio_anchor": None, + "video_edit": video_edit, + "audio_edit": audio_edit, "sigmas_video": sigmas_video, "sigmas_audio": sigmas_audio, }, @@ -347,6 +446,8 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames): assert state.do_true_cfg is False torch.testing.assert_close(state.current_timestep, torch.tensor(1.0 - sigmas_video[0])) assert state.extra[mod._STEP_BRANCH] is branch + assert state.extra[mod._STEP_VIDEO_EDIT] is video_edit + assert state.extra[mod._STEP_AUDIO_EDIT] is audio_edit assert state.extra[mod._STEP_SHAPE]["height"] == 96 pipeline.od_config = SimpleNamespace() diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 715c7ca04..f0e7b7d53 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -11,6 +11,7 @@ import json import os import threading import time +from concurrent.futures import CancelledError as FutureCancelledError from contextlib import asynccontextmanager from pathlib import Path from types import SimpleNamespace @@ -84,6 +85,7 @@ class FakeAsyncOmni: self.model_class_name = "WanPipeline" self.captured_prompt = None self.captured_reference_video_bytes = None + self.captured_latent_edit_source_bytes = {} self.captured_control_reference_bytes = {} self.captured_sampling_params_list = None @@ -104,6 +106,10 @@ class FakeAsyncOmni: and all(isinstance(item, str) for item in reference_videos) ): self.captured_reference_video_bytes = [Path(item).read_bytes() for item in reference_videos] + for field_name in ("source_video", "source_audio"): + source_path = prompt.get("multi_modal_data", {}).get(field_name) + if isinstance(source_path, str): + self.captured_latent_edit_source_bytes[field_name] = Path(source_path).read_bytes() num_outputs = sampling_params_list[0].num_outputs_per_prompt if sampling_params_list[0].emit_request_lifecycle: yield MockVideoResult( @@ -224,6 +230,10 @@ def test_resolve_diffusion_od_config_prefers_getter_over_attribute(): class BlockingVideoHandler: + supports_mixed_reference_inputs = False + supports_latent_mask_editing = False + supported_control_upload_types = frozenset() + def __init__(self): self.model_name = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" self.stage_configs = None @@ -243,8 +253,9 @@ class BlockingVideoHandler: reference_video=None, reference_audio=None, on_started=None, + latent_edit_input=None, ): - del request, reference_id, reference_image, reference_video, reference_audio + del request, reference_id, reference_image, reference_video, reference_audio, latent_edit_input if on_started is not None: await on_started() self.started.set() @@ -282,8 +293,9 @@ class CompletingDuringAbortHandler(BlockingVideoHandler): reference_video=None, reference_audio=None, on_started=None, + latent_edit_input=None, ): - del request, reference_image, reference_video, reference_audio + del request, reference_image, reference_video, reference_audio, latent_edit_input if on_started is not None: await on_started() self.started.set() @@ -318,8 +330,9 @@ class SchedulerQueuedVideoHandler(BlockingVideoHandler): reference_video=None, reference_audio=None, on_started=None, + latent_edit_input=None, ): - del request, reference_id, reference_image, reference_video, reference_audio + del request, reference_id, reference_image, reference_video, reference_audio, latent_edit_input self.started.set() try: while not self.admit.is_set(): @@ -496,6 +509,11 @@ def _make_test_video_bytes(size=(32, 24), num_frames=3) -> bytes: return mux_video_audio_bytes(frames, fps=8, video_codec_options={"preset": "ultrafast", "threads": "0"}) +def _mask_file(value) -> tuple[str, bytes, str]: + payload = value if isinstance(value, str) else json.dumps(value) + return "mask.json", payload.encode(), "application/json" + + def test_mux_video_audio_marks_aac_priming_timestamp(): frames = np.zeros((2, 16, 16, 3), dtype=np.uint8) audio = np.zeros((2, 2048), dtype=np.float32) @@ -1137,6 +1155,7 @@ def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_fal def test_typed_stage_drives_video_capability_checks(): from vllm_omni.config.config_factory import StageConfigFactory + from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata minimax_stage = StageConfigFactory.create_typed_default_diffusion( "minimax-h3", @@ -1148,6 +1167,8 @@ def test_typed_stage_drives_video_capability_checks(): stage_configs=[minimax_stage], ) assert minimax_handler.supports_mixed_reference_inputs + assert minimax_handler.supports_latent_mask_editing + assert get_diffusion_model_metadata("MiniMaxH3ModularPipeline").supports_latent_mask_editing cosmos_stage = StageConfigFactory.create_typed_default_diffusion( "cosmos3", @@ -1159,6 +1180,222 @@ def test_typed_stage_drives_video_capability_checks(): stage_configs=[cosmos_stage], ) assert cosmos_handler.supported_control_upload_types == frozenset({"edge", "blur", "depth", "seg", "wsm"}) + assert not cosmos_handler.supports_latent_mask_editing + + +@pytest.mark.parametrize("endpoint", ["/v1/videos", "/v1/videos/sync"]) +def test_h3_latent_edit_uploads_reach_generation_and_are_cleaned(endpoint, test_client, mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + engine = test_client.app.state.openai_serving_video._engine_client + engine.model_class_name = "MiniMaxH3Pipeline" + video_mask = [[[0.0, 1.0], [0.5, 1.0]]] + audio_mask = [0.0, 0.25, 1.0] + + response = test_client.post( + endpoint, + data={"prompt": "Replace only the masked regions."}, + files=[ + ("source_video", ("source.mov", b"source-video", "video/quicktime")), + ("source_audio", ("source.mp3", b"source-audio", "audio/mpeg")), + ("video_noise_mask", _mask_file(video_mask)), + ("audio_noise_mask", _mask_file(audio_mask)), + ], + ) + + assert response.status_code == 200 + if not endpoint.endswith("/sync"): + _wait_for_status(test_client, response.json()["id"], VideoGenerationStatus.COMPLETED.value) + + multi_modal_data = engine.captured_prompt["multi_modal_data"] + assert engine.captured_latent_edit_source_bytes == { + "source_video": b"source-video", + "source_audio": b"source-audio", + } + assert multi_modal_data["video_noise_mask"] == video_mask + assert multi_modal_data["audio_noise_mask"] == audio_mask + _wait_until( + lambda: ( + not Path(multi_modal_data["source_video"]).exists() and not Path(multi_modal_data["source_audio"]).exists() + ) + ) + assert not Path(multi_modal_data["source_video"]).exists() + assert not Path(multi_modal_data["source_audio"]).exists() + + +def test_latent_edit_fields_are_rejected_for_models_without_capability(test_client): + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "edit"}, + files=[ + ("source_video", ("source.mp4", b"video", "video/mp4")), + ("video_noise_mask", _mask_file("1")), + ], + ) + + assert response.status_code == 400 + assert "not supported" in response.json()["detail"].lower() + + +@pytest.mark.parametrize( + ("files", "message"), + [ + ( + [("source_video", ("source.mp4", b"video", "video/mp4"))], + "at least one", + ), + ( + [ + ("source_audio", ("source.wav", b"audio", "audio/wav")), + ("video_noise_mask", _mask_file("0")), + ], + "source_audio requires an audio_noise_mask", + ), + ], +) +def test_h3_latent_edit_cross_field_validation(test_client, files, message): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + + response = test_client.post("/v1/videos/sync", data={"prompt": "edit"}, files=files) + + assert response.status_code == 400 + assert message in response.json()["detail"] + + +@pytest.mark.parametrize( + ("field_name", "filename", "content_type", "message"), + [ + ("source_video", "source.webm", "video/webm", "MP4 or MOV"), + ("source_audio", "source.flac", "audio/flac", "WAV or MP3"), + ], +) +def test_h3_latent_edit_rejects_unsupported_source_formats( + test_client, + field_name, + filename, + content_type, + message, +): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + mask_name = "video_noise_mask" if field_name == "source_video" else "audio_noise_mask" + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "edit"}, + files=[ + (field_name, (filename, b"source", content_type)), + (mask_name, _mask_file("0")), + ], + ) + + assert response.status_code == 400 + assert message in response.json()["detail"] + + +@pytest.mark.parametrize( + ("mask", "message"), + [ + ("[not-json", "Invalid JSON"), + ("[" * 2000 + "0" + "]" * 2000, "nested too deeply"), + ("[" * 65 + "0" + "]" * 65, "nested too deeply"), + ], +) +def test_h3_latent_edit_rejects_malformed_mask_json(test_client, mask, message): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "edit"}, + files=[ + ("source_video", ("source.mp4", b"video", "video/mp4")), + ("video_noise_mask", _mask_file(mask)), + ], + ) + + assert response.status_code == 400 + assert message in response.json()["detail"] + + +def test_h3_latent_edit_rejects_mask_text_form_fields(test_client): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "edit", "video_noise_mask": "0"}, + ) + + assert response.status_code == 422 + + +def test_h3_latent_edit_masks_without_sources_are_forwarded_to_stage0(test_client, mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + engine = test_client.app.state.openai_serving_video._engine_client + engine.model_class_name = "MiniMaxH3Pipeline" + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "ordinary generation"}, + files=[ + ("video_noise_mask", _mask_file("1")), + ("audio_noise_mask", _mask_file("[1, 1]")), + ], + ) + + assert response.status_code == 200 + assert engine.captured_prompt["multi_modal_data"]["video_noise_mask"] == 1 + assert engine.captured_prompt["multi_modal_data"]["audio_noise_mask"] == [1, 1] + + +def test_h3_latent_edit_rejects_oversized_mask_json(test_client, monkeypatch): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + monkeypatch.setattr(video_generation_helpers, "LATENT_EDIT_MASK_FILE_MAX_BYTES", 4) + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "edit"}, + files=[ + ("source_video", ("source.mp4", b"video", "video/mp4")), + ("video_noise_mask", _mask_file("[0.0]")), + ], + ) + + assert response.status_code == 400 + assert "JSON exceeds" in response.json()["detail"] + + +def test_cancelled_latent_edit_persistence_cleans_prior_source(test_client, monkeypatch): + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + persist = video_generation_helpers._persist_latent_edit_source + persisted_paths: list[str] = [] + + async def persist_then_cancel(upload, *, field_name): + if field_name == "source_audio": + raise asyncio.CancelledError + path = await persist(upload, field_name=field_name) + persisted_paths.append(path) + return path + + monkeypatch.setattr(video_generation_helpers, "_persist_latent_edit_source", persist_then_cancel) + + with pytest.raises(FutureCancelledError): + test_client.post( + "/v1/videos/sync", + data={"prompt": "edit"}, + files=[ + ("source_video", ("source.mp4", b"source-video", "video/mp4")), + ("source_audio", ("source.wav", b"source-audio", "audio/wav")), + ("video_noise_mask", _mask_file("0")), + ("audio_noise_mask", _mask_file("0")), + ], + ) + + assert len(persisted_paths) == 1 + assert not Path(persisted_paths[0]).exists() @pytest.mark.parametrize("model_class_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) @@ -2030,6 +2267,33 @@ def test_failed_generation_awaits_storage_cleanup(test_client, isolated_video_ba delete_spy.assert_called_once_with(video_id) +def test_failed_async_latent_edit_cleans_source_upload(test_client, mocker: MockerFixture): + handler = test_client.app.state.openai_serving_video + handler._engine_client.model_class_name = "MiniMaxH3Pipeline" + captured_path: list[str] = [] + + async def fail_generation(*_args, latent_edit_input=None, **_kwargs): + assert latent_edit_input is not None + captured_path.append(latent_edit_input.source_video) + assert Path(latent_edit_input.source_video).read_bytes() == b"source-video" + raise RuntimeError("edit failed") + + mocker.patch.object(handler, "generate_video_bytes", side_effect=fail_generation) + response = test_client.post( + "/v1/videos", + data={"prompt": "will fail"}, + files=[ + ("source_video", ("source.mp4", b"source-video", "video/mp4")), + ("video_noise_mask", _mask_file("0")), + ], + ) + assert response.status_code == 200 + + _wait_for_status(test_client, response.json()["id"], VideoGenerationStatus.FAILED.value) + assert len(captured_path) == 1 + _wait_until(lambda: not Path(captured_path[0]).exists()) + + def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture): mocker.patch.object( OmniOpenAIServingVideo, 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 4c9a7137b..ce6c47ac7 100644 --- a/tests/model_executor/models/test_minimax_h3_encoder_payload.py +++ b/tests/model_executor/models/test_minimax_h3_encoder_payload.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import subprocess from contextlib import contextmanager from types import SimpleNamespace @@ -9,15 +10,82 @@ import torch from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MINIMAX_H3_ENCODER_LAYOUT_KEY, + STAGE_SCHEMA_VERSION, MiniMaxH3EncoderConditioning, MiniMaxH3EncoderMediaConditioning, MiniMaxH3EncoderMediaInput, ) +from vllm_omni.model_executor.models.minimax_h3.encoder_processing import ( + _canonical_audio_edit_mask, + _canonical_video_edit_mask, +) from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion 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} + 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( + _canonical_video_edit_mask(torch.tensor([[[0.25]]]), **video_kwargs), + torch.full((7, 8, 12), 0.25), + ) + + audio_temporal = torch.arange(37, dtype=torch.float32) / 37 + audio_full = audio_temporal.repeat(2, 1) + for value in (audio_temporal[None, None], audio_full.flatten()[None], audio_full[None]): + torch.testing.assert_close(_canonical_audio_edit_mask(value, audio_t=37), audio_full) + torch.testing.assert_close( + _canonical_audio_edit_mask(torch.tensor([[[0.25]]]), audio_t=37), + torch.full((2, 37), 0.25), + ) + + +@pytest.mark.parametrize( + ("canonicalize", "value", "kwargs", "message"), + [ + (_canonical_video_edit_mask, [True], {"latent_t": 7, "latent_h": 8, "latent_w": 12}, "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"), + ], +) +def test_edit_mask_boundary_rejects_invalid_values(canonicalize, value, kwargs, message) -> None: + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError, match=message): + canonicalize(value, **kwargs) + + +def test_encoder_elides_all_generate_masks_but_requires_sources(monkeypatch) -> None: + from vllm_omni.errors import OmniClientError + from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing + + monkeypatch.setattr(processing, "resolve_minimax_h3_shape", lambda *_args: (64, 96, 22, 7, 37)) + sampling = SimpleNamespace(extra_args={"task": "t2va"}) + prepared = processing.prepare_encoder_inputs( + { + "prompt": "generate", + "multi_modal_data": {"video_noise_mask": 1.0, "audio_noise_mask": 1.0}, + }, + sampling, + ) + assert prepared.media.video_edit is None + assert prepared.media.audio_edit is None + + with pytest.raises(OmniClientError, match="video_noise_mask requires source_video"): + processing.prepare_encoder_inputs( + {"prompt": "edit", "multi_modal_data": {"video_noise_mask": 0.5}}, + sampling, + ) + + def test_pipeline_owns_three_encoders_then_dit_and_decoders() -> None: from vllm_omni.model_executor.models.minimax_h3.pipeline import MINIMAX_H3_PIPELINE @@ -275,6 +343,43 @@ def test_prepare_encoder_inputs_keeps_reference_audio_budgets_separate(monkeypat assert prepared.condition_labels == [("audio", 1), ("video", 1), ("audio", 2)] +@pytest.mark.parametrize( + ("source_key", "mask_key", "mask_shape", "decoder"), + [ + ("source_video", "video_noise_mask", (7, 4, 6), "prepare_edit_video"), + ("source_audio", "audio_noise_mask", (37,), "load_audio_file"), + ("source_video", "audio_noise_mask", (37,), "load_video_audio"), + ], +) +def test_prepare_encoder_inputs_maps_corrupt_edit_media_to_client_error( + monkeypatch, + source_key, + mask_key, + mask_shape, + decoder, +) -> None: + from vllm_omni.errors import OmniClientError + from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing + + def fail_decode(*_args, **_kwargs): + raise subprocess.CalledProcessError(1, ["ffmpeg"]) + + monkeypatch.setattr(processing, "resolve_minimax_h3_shape", lambda *_args: (64, 96, 22, 7, 37)) + monkeypatch.setattr(processing, decoder, fail_decode) + + with pytest.raises(OmniClientError, match=f"could not decode {source_key}"): + processing.prepare_encoder_inputs( + { + "prompt": "edit", + "multi_modal_data": { + source_key: "corrupt.media", + mask_key: torch.zeros(mask_shape), + }, + }, + SimpleNamespace(extra_args={"task": "t2va"}), + ) + + @pytest.mark.parametrize( ("embedded_lengths", "standalone_lengths", "valid"), [((400,), (400,), True), ((320, 320), (400,), False), ((400,), (320, 320), False)], @@ -339,7 +444,7 @@ def test_encode_media_keeps_audio_budgets_and_component_residency_separate( assert scopes == [video_vae, audio_vae] -def test_encoder_output_reuses_encoder_handoff_and_round_trips() -> None: +def test_encoder_output_reuses_encoder_handoff_and_round_trips(monkeypatch) -> None: expected = MiniMaxH3EncoderConditioning( hidden_states=torch.randn(3, 5120, dtype=torch.bfloat16), token_tags=torch.tensor([1, 0, 1], dtype=torch.int64), @@ -349,8 +454,15 @@ def test_encoder_output_reuses_encoder_handoff_and_round_trips() -> None: num_frames=17, latent_t=5, audio_t=10, + video_edit_clean_rows=torch.zeros(5 * 8 * 14, 96), + video_edit_mask=torch.zeros(5, 16, 28), ) + def fail_value_scan(*_args, **_kwargs): + raise AssertionError("wire must trust encoder-boundary value validation") + + monkeypatch.setattr(torch, "isfinite", fail_value_scan) + source = SimpleNamespace( finished=True, request_id="request-1", @@ -366,8 +478,19 @@ def test_encoder_output_reuses_encoder_handoff_and_round_trips() -> None: assert result is not None payload = result["additional_information"]["encoder_output"] actual = MiniMaxH3EncoderConditioning.from_omni_payload(payload) + monkeypatch.undo() torch.testing.assert_close(actual.hidden_states, expected.hidden_states) torch.testing.assert_close(actual.token_tags, expected.token_tags) + torch.testing.assert_close(actual.video_edit_mask, expected.video_edit_mask) + + layout = payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY] + assert int(layout[1]) == STAGE_SCHEMA_VERSION + legacy_layout = layout.clone() + legacy_layout[1] = STAGE_SCHEMA_VERSION - 1 + payload["kv_metadata"][MINIMAX_H3_ENCODER_LAYOUT_KEY] = legacy_layout + + with pytest.raises(ValueError, match=rf"unsupported MiniMax H3 encoder wire schema 1:{STAGE_SCHEMA_VERSION - 1}"): + MiniMaxH3EncoderConditioning.from_omni_payload(payload) def test_encoder_releases_workspace_after_cpu_payload(monkeypatch) -> None: diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index 2ee1b1f4b..1059093b6 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -17,6 +17,10 @@ class DiffusionModelMetadata: supported_control_upload_types: tuple[str, ...] = () attention_mask_free: bool = False final_output_type: str | None = None + # Whether ``/v1/videos`` accepts source media plus per-token video/audio + # noise masks for latent initialization. Unknown pipelines must remain + # opted out so uploaded files never reach a model that cannot consume them. + supports_latent_mask_editing: bool = False # FLUX.2 Klein supports up to four reference images. @@ -52,6 +56,7 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { supports_multimodal_inputs=True, max_multimodal_image_inputs=9, supports_mixed_reference_inputs=True, + supports_latent_mask_editing=True, final_output_type="video", # H3 represents alignment padding as a second packed sequence. The # packed TRTLLM backend consumes cu_seqlens and isolates that padding. @@ -66,6 +71,7 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { supports_multimodal_inputs=True, max_multimodal_image_inputs=9, supports_mixed_reference_inputs=True, + supports_latent_mask_editing=True, final_output_type="video", attention_mask_free=True, ), diff --git a/vllm_omni/diffusion/models/minimax_h3/batched_packing.py b/vllm_omni/diffusion/models/minimax_h3/batched_packing.py index 4d41bde61..a9bf7cbba 100644 --- a/vllm_omni/diffusion/models/minimax_h3/batched_packing.py +++ b/vllm_omni/diffusion/models/minimax_h3/batched_packing.py @@ -38,6 +38,8 @@ def minimax_h3_batched_forward_kwargs( t_audio: Sequence[float], imgvid_cond_timesteps: Sequence[float], audio_ref_cond_timesteps: Sequence[float], + video_target_timesteps: Sequence[torch.Tensor | None] | None = None, + audio_target_timesteps: Sequence[torch.Tensor | None] | None = None, ) -> dict[str, Any]: """Build one DiT forward kwargs dict covering every request in the batch. @@ -46,6 +48,16 @@ def minimax_h3_batched_forward_kwargs( happen to agree -- which also keeps ``video_token_layout`` (the sparse- attention hint) on the single-request step path. """ + request_count = len(branches) + if video_target_timesteps is None: + video_target_timesteps = [None] * request_count + elif len(video_target_timesteps) != request_count: + raise ValueError(f"video_target_timesteps has {len(video_target_timesteps)} requests, expected {request_count}") + if audio_target_timesteps is None: + audio_target_timesteps = [None] * request_count + elif len(audio_target_timesteps) != request_count: + raise ValueError(f"audio_target_timesteps has {len(audio_target_timesteps)} requests, expected {request_count}") + if len(branches) == 1: return branches[0].forward_kwargs( video_rows=video_rows[0], @@ -54,6 +66,8 @@ def minimax_h3_batched_forward_kwargs( t_audio=t_audio[0], imgvid_cond_timestep=imgvid_cond_timesteps[0], audio_ref_cond_timestep=audio_ref_cond_timesteps[0], + video_target_timesteps=video_target_timesteps[0], + audio_target_timesteps=audio_target_timesteps[0], ) device = branches[0].device @@ -88,11 +102,15 @@ def minimax_h3_batched_forward_kwargs( position_id_parts.append(branch.img_position_ids_dev) text_embed_parts.append(branch.text_embeddings_dev) - timesteps[seq_offset : seq_offset + branch.seq_len] = float(t_video[index]) - timesteps[img_pos[branch.update_mask_dev]] = float(t_video[index]) - timesteps[img_pos[~branch.update_mask_dev]] = float(imgvid_cond_timesteps[index]) - timesteps[audio_pos[branch.audio_update_mask_dev]] = float(t_audio[index]) - timesteps[audio_pos[~branch.audio_update_mask_dev]] = float(audio_ref_cond_timesteps[index]) + branch.fill_timesteps( + timesteps[seq_offset : seq_offset + branch.seq_len], + t_video=t_video[index], + t_audio=t_audio[index], + imgvid_cond_timestep=imgvid_cond_timesteps[index], + audio_ref_cond_timestep=audio_ref_cond_timesteps[index], + video_target_timesteps=video_target_timesteps[index], + audio_target_timesteps=audio_target_timesteps[index], + ) # Empty documents are omitted because not every varlen kernel accepts # repeated interior boundaries. diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index 5f7579d07..ae6768f7a 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -23,6 +23,7 @@ from vllm_omni.diffusion.forward_context import ( ) from vllm_omni.platforms import current_omni_platform +from .latent_mask import MiniMaxH3LatentEdit, minimax_h3_prepare_edit_rows from .scheduling_minimax_h3_euler_ancestral import ( minimax_h3_euler_eta0_step, minimax_h3_rf_v_to_x0, @@ -185,25 +186,23 @@ class MiniMaxH3DenoiseBranch: t_audio: float, imgvid_cond_timestep: float, audio_ref_cond_timestep: float, + video_target_timesteps: torch.Tensor | None = None, + audio_target_timesteps: torch.Tensor | None = None, ) -> dict[str, Any]: x = self.x_base.clone() x[0].index_copy_(0, self.img_pos_dev, video_rows) audio_x = self.audio_x_base.clone() audio_x[0].index_copy_(0, self.audio_pos_dev, audio_rows) - # Packed-sequence timestep semantics: non-media rows (text and - # padding) inherit the current video timestep. Later steps must reuse - # the previous step's updated rows; re-initializing from zeros is only - # valid at step 0. - timesteps = torch.full( - (self.seq_len,), - float(t_video), - dtype=torch.float32, - device=x.device, + timesteps = torch.empty(self.seq_len, dtype=torch.float32, device=self.device) + self.fill_timesteps( + timesteps, + t_video=t_video, + t_audio=t_audio, + imgvid_cond_timestep=imgvid_cond_timestep, + audio_ref_cond_timestep=audio_ref_cond_timestep, + video_target_timesteps=video_target_timesteps, + audio_target_timesteps=audio_target_timesteps, ) - timesteps[self.img_pos_dev[self.update_mask_dev]] = t_video - timesteps[self.img_pos_dev[~self.update_mask_dev]] = imgvid_cond_timestep - timesteps[self.audio_pos_dev[self.audio_update_mask_dev]] = t_audio - timesteps[self.audio_pos_dev[~self.audio_update_mask_dev]] = audio_ref_cond_timestep unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True) return { **self.static_kwargs, @@ -213,6 +212,56 @@ class MiniMaxH3DenoiseBranch: "inverse_indices": inverse_indices, } + def fill_timesteps( + self, + timesteps: torch.Tensor, + *, + t_video: float, + t_audio: float, + imgvid_cond_timestep: float, + audio_ref_cond_timestep: float, + video_target_timesteps: torch.Tensor | None = None, + audio_target_timesteps: torch.Tensor | None = None, + ) -> None: + """Fill one request's packed row timesteps in an existing tensor.""" + if timesteps.shape != (self.seq_len,): + raise ValueError(f"timesteps must have shape ({self.seq_len},)") + # Packed-sequence timestep semantics: non-media rows (text and + # padding) inherit the current video timestep. Later steps must reuse + # the previous step's updated rows; re-initializing from zeros is only + # valid at step 0. + timesteps.fill_(float(t_video)) + timestep_groups = ( + ( + "video_target_timesteps", + self.img_pos_dev, + self.update_mask_dev, + video_target_timesteps, + t_video, + imgvid_cond_timestep, + ), + ( + "audio_target_timesteps", + self.audio_pos_dev, + self.audio_update_mask_dev, + audio_target_timesteps, + t_audio, + audio_ref_cond_timestep, + ), + ) + for name, positions, update_mask, target_timesteps, current_timestep, condition_timestep in timestep_groups: + target_positions = positions[update_mask] + if target_timesteps is None: + timesteps[target_positions] = current_timestep + else: + target_timesteps = target_timesteps.to(device=self.device, dtype=torch.float32).reshape(-1) + if target_timesteps.shape[0] != target_positions.shape[0]: + raise ValueError( + f"{name} rows {target_timesteps.shape[0]} != target rows {target_positions.shape[0]}" + ) + timesteps[target_positions] = target_timesteps + timesteps[positions[~update_mask]] = condition_timestep + def minimax_h3_prepare_denoise_rows( *, @@ -273,6 +322,8 @@ def minimax_h3_denoise_loop( initial_audio_rows: torch.Tensor, keyframe_cond_rows: torch.Tensor | None, audio_ref_rows: torch.Tensor | None = None, + video_edit: MiniMaxH3LatentEdit | None = None, + audio_edit: MiniMaxH3LatentEdit | None = None, sigmas_video: list[float], sigmas_audio: list[float], device: torch.device, @@ -307,6 +358,10 @@ def minimax_h3_denoise_loop( ) update = positive.update_mask_dev audio_update = positive.audio_update_mask_dev + if video_edit is not None: + video_edit = video_edit.to(device=device, dtype=torch.float32) + if audio_edit is not None: + audio_edit = audio_edit.to(device=device, dtype=torch.float32) num_steps = len(sigmas_video) - 1 for step in range(num_steps): @@ -323,35 +378,68 @@ def minimax_h3_denoise_loop( imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference)) audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference)) + model_video_rows, video_target_timesteps = minimax_h3_prepare_edit_rows( + video_rows, + update, + video_edit, + t_v, + imgvid_cond_t, + sigma=s_v, + ) + model_audio_rows, audio_target_timesteps = minimax_h3_prepare_edit_rows( + audio_rows, + audio_update, + audio_edit, + t_a, + audio_ref_cond_t, + sigma=s_a, + ) + fk = positive.forward_kwargs( - video_rows=video_rows, - audio_rows=audio_rows, + video_rows=model_video_rows, + audio_rows=model_audio_rows, t_video=t_v, t_audio=t_a, imgvid_cond_timestep=imgvid_cond_t, audio_ref_cond_timestep=audio_ref_cond_t, + video_target_timesteps=video_target_timesteps, + audio_target_timesteps=audio_target_timesteps, ) with torch.inference_mode(): v_video, v_audio = model(**fk) mv_video_t = v_video.float()[update] mv_audio_t = v_audio.float()[audio_update] - x0_video = minimax_h3_rf_v_to_x0( - video_rows[update], - mv_video_t, - torch.tensor(t_v, dtype=torch.float32, device=device), - ) + if video_edit is None: + x0_video = minimax_h3_rf_v_to_x0( + video_rows[update], + mv_video_t, + torch.tensor(t_v, dtype=torch.float32, device=device), + ) + else: + x0_video = video_edit.x0( + model_video_rows[update], + mv_video_t, + t_v, + ) new_target = minimax_h3_euler_eta0_step(video_rows[update], x0_video, sigma_curr=s_v, sigma_next=s_v_next) video_rows = video_rows.clone() video_rows[update] = new_target if cond_anchor is not None: video_rows[~update] = cond_anchor # per-step imgvid cond reset - x0_audio = minimax_h3_rf_v_to_x0( - audio_rows[audio_update], - mv_audio_t, - torch.tensor(t_a, dtype=torch.float32, device=device), - ) + if audio_edit is None: + x0_audio = minimax_h3_rf_v_to_x0( + audio_rows[audio_update], + mv_audio_t, + torch.tensor(t_a, dtype=torch.float32, device=device), + ) + else: + x0_audio = audio_edit.x0( + model_audio_rows[audio_update], + mv_audio_t, + t_a, + ) new_audio = minimax_h3_euler_eta0_step( audio_rows[audio_update], x0_audio, sigma_curr=s_a, sigma_next=s_a_next ) diff --git a/vllm_omni/diffusion/models/minimax_h3/latent_mask.py b/vllm_omni/diffusion/models/minimax_h3/latent_mask.py new file mode 100644 index 000000000..cd740284c --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/latent_mask.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax H3 latent-edit mask parsing and sampler math.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch + +from .scheduling_minimax_h3_euler_ancestral import minimax_h3_rf_v_to_x0 + +_MASK_LEVELS = 256.0 + + +def _quantize(mask: torch.Tensor) -> torch.Tensor: + """Match ComfyUI's token-grid mask used by H3's timestep embedding.""" + return (torch.ceil(mask * _MASK_LEVELS) / _MASK_LEVELS).clamp_(0, 1) + + +@dataclass(frozen=True) +class MiniMaxH3ParsedMask: + """Token mask for the model and raw mask for the final x0 restore.""" + + model_mask_rows: torch.Tensor + restore_mask_rows: torch.Tensor + + +def minimax_h3_video_edit_masks( + mask: torch.Tensor, + *, + latent_t: int, + latent_h: int, + latent_w: int, +) -> MiniMaxH3ParsedMask: + """Map the canonical full-grid video mask to model and restore rows. + + Full-grid masks are max-pooled over each 2x2 DiT token. Their raw four + cells are repeated in the packed 24-channel feature order for restoration. + """ + latent_t, latent_h, latent_w = int(latent_t), int(latent_h), int(latent_w) + if latent_t <= 0 or latent_h <= 0 or latent_w <= 0 or latent_h % 2 or latent_w % 2: + raise ValueError("video latent dimensions must be positive and spatially even") + + token_shape = (latent_t, latent_h // 2, latent_w // 2) + full_shape = (latent_t, latent_h, latent_w) + row_count = math.prod(token_shape) + if mask.dtype != torch.float32 or tuple(mask.shape) != full_shape or not mask.is_contiguous(): + raise ValueError(f"canonical video mask must be contiguous FP32 with shape {full_shape}") + cells = mask.reshape(latent_t, latent_h // 2, 2, latent_w // 2, 2).permute(0, 1, 3, 2, 4).contiguous() + model = _quantize(cells.amax(dim=(3, 4)).reshape(-1)) + restore = cells.unsqueeze(3).expand(*cells.shape[:3], 24, 2, 2).reshape(row_count, 96).clone() + return MiniMaxH3ParsedMask(model, restore) + + +def minimax_h3_audio_edit_masks(mask: torch.Tensor, *, audio_t: int) -> MiniMaxH3ParsedMask: + """Flatten the canonical channel-major audio mask to H3 row order.""" + audio_t = int(audio_t) + if audio_t <= 0: + raise ValueError("audio_t must be positive") + expected_shape = (2, audio_t) + if mask.dtype != torch.float32 or tuple(mask.shape) != expected_shape or not mask.is_contiguous(): + raise ValueError(f"canonical audio mask must be contiguous FP32 with shape {expected_shape}") + rows = mask.reshape(-1).clone() + return MiniMaxH3ParsedMask(_quantize(rows), rows) + + +@dataclass(frozen=True) +class MiniMaxH3LatentEdit: + """Clean source, model anchor, and masks for one packed latent stream.""" + + clean_rows: torch.Tensor + anchor_rows: torch.Tensor + mask_rows: torch.Tensor + restore_mask_rows: torch.Tensor + + @classmethod + def from_rows( + cls, + clean_rows: torch.Tensor, + anchor_rows: torch.Tensor, + mask_rows: torch.Tensor, + restore_mask_rows: torch.Tensor, + ) -> MiniMaxH3LatentEdit: + """Create an edit from already-validated encoder outputs. + + Value validation and all-generate elision belong at the encoder input + boundary. This constructor only normalizes devices/dtypes and checks + the cheap row-shape invariants needed by the sampler math. + """ + if clean_rows.ndim != 2 or anchor_rows.shape != clean_rows.shape: + raise ValueError("clean_rows and anchor_rows must have matching [rows, width] shape") + raw_mask = torch.as_tensor(mask_rows, dtype=torch.float32, device=clean_rows.device) + if raw_mask.ndim != 1 or raw_mask.numel() != clean_rows.shape[0]: + raise ValueError(f"mask_rows must have shape ({clean_rows.shape[0]},)") + restore = torch.as_tensor(restore_mask_rows, dtype=torch.float32, device=clean_rows.device) + valid_restore = restore.ndim == 1 and restore.shape == raw_mask.shape + valid_restore |= restore.ndim == 2 and tuple(restore.shape) in { + (clean_rows.shape[0], 1), + tuple(clean_rows.shape), + } + if not valid_restore: + raise ValueError("restore_mask_rows must have one scalar or one feature row per clean row") + return cls(clean_rows, anchor_rows, raw_mask, restore) + + def to(self, *, device: torch.device, dtype: torch.dtype = torch.float32) -> MiniMaxH3LatentEdit: + return MiniMaxH3LatentEdit( + self.clean_rows.to(device=device, dtype=dtype), + self.anchor_rows.to(device=device, dtype=dtype), + self.mask_rows.to(device=device), + self.restore_mask_rows.to(device=device), + ) + + def model_rows(self, state_rows: torch.Tensor) -> torch.Tensor: + return torch.lerp(self.anchor_rows, state_rows, self.mask_rows.to(state_rows.dtype).unsqueeze(-1)) + + def target_timesteps( + self, + timestep: float, + condition_timestep: float, + *, + sigma: float | None = None, + ) -> torch.Tensor: + sigma = 1.0 - timestep if sigma is None else float(sigma) + return torch.clamp(1.0 - self.mask_rows * sigma, max=condition_timestep) + + def prepare( + self, + state_rows: torch.Tensor, + timestep: float, + condition_timestep: float, + *, + sigma: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare the model rows and per-row timesteps for one denoise step.""" + return ( + self.model_rows(state_rows), + self.target_timesteps(timestep, condition_timestep, sigma=sigma), + ) + + def x0(self, model_rows: torch.Tensor, velocity: torch.Tensor, timestep: float) -> torch.Tensor: + weight = self.mask_rows.to(model_rows.dtype).unsqueeze(-1) + predicted = minimax_h3_rf_v_to_x0( + model_rows, + velocity.to(model_rows.dtype) * weight, + torch.tensor(float(timestep), dtype=model_rows.dtype, device=model_rows.device), + ) + restore = self.restore_mask_rows.to(model_rows.dtype) + if restore.ndim == 1: + restore = restore.unsqueeze(-1) + return torch.lerp(self.clean_rows, predicted, restore) + + +def minimax_h3_prepare_edit_rows( + rows: torch.Tensor, + update_mask: torch.Tensor, + edit: MiniMaxH3LatentEdit | None, + timestep: float, + condition_timestep: float, + *, + sigma: float, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Build one model input view without mutating the persistent sampler rows.""" + if edit is None: + return rows, None + model_rows = rows.clone() + model_rows[update_mask], target_timesteps = edit.prepare( + rows[update_mask], + timestep, + condition_timestep, + sigma=sigma, + ) + return model_rows, target_timesteps + + +__all__ = [ + "MiniMaxH3LatentEdit", + "MiniMaxH3ParsedMask", + "minimax_h3_audio_edit_masks", + "minimax_h3_prepare_edit_rows", + "minimax_h3_video_edit_masks", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 61c0ceecd..b297dae26 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -108,6 +108,12 @@ from .denoise_loop import ( from .encoder import MiniMaxH3Qwen3VLEncoder from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion from .fasth3_checkpoint import FastH3CheckpointSpec +from .latent_mask import ( + MiniMaxH3LatentEdit, + minimax_h3_audio_edit_masks, + minimax_h3_prepare_edit_rows, + minimax_h3_video_edit_masks, +) from .lora import TurboSpec, load_minimax_h3_turbo_lora from .minimax_h3_transformer import ( MiniMaxH3Attention, @@ -261,6 +267,12 @@ _MINIMAX_H3_DENOISE_INPUT_KEYS = ( "audio_condition_lengths", "keyframe_frame_indices", "pad_seq_len", + "video_edit_clean_rows", + "video_edit_mask_rows", + "video_edit_restore_mask_rows", + "audio_edit_clean_rows", + "audio_edit_mask_rows", + "audio_edit_restore_mask_rows", ) # ``StepRequestState.extra`` keys owned by the step-execution path. @@ -273,6 +285,8 @@ _STEP_COND_ANCHOR = "minimax_h3_cond_anchor" _STEP_AUDIO_ANCHOR = "minimax_h3_audio_anchor" _STEP_SHAPE = "minimax_h3_shape" _STEP_TRANSFORMER = "minimax_h3_transformer" +_STEP_VIDEO_EDIT = "minimax_h3_video_edit" +_STEP_AUDIO_EDIT = "minimax_h3_audio_edit" def _minimax_h3_step_schedule(state: StepRequestState) -> dict[str, float]: @@ -407,6 +421,27 @@ def get_minimax_h3_post_process_func( return _minimax_h3_post_process +def _expose_padded_audio_tail( + source_audio_t: int, + *, + target_audio_t: int, + mask_rows: torch.Tensor, + restore_mask_rows: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Mark the encoder-padded audio tail for generation.""" + if not 0 < source_audio_t <= target_audio_t: + raise ValueError(f"source_audio_t must be in [1, {target_audio_t}]") + expected_shape = (2 * target_audio_t,) + if tuple(mask_rows.shape) != expected_shape or tuple(restore_mask_rows.shape) != expected_shape: + raise ValueError(f"audio edit masks must have shape {expected_shape}") + model_mask = mask_rows.reshape(2, target_audio_t).clone() + restore_mask = restore_mask_rows.reshape(2, target_audio_t).clone() + if source_audio_t < target_audio_t: + model_mask[:, source_audio_t:] = 1.0 + restore_mask[:, source_audio_t:] = 1.0 + return model_mask.reshape(-1), restore_mask.reshape(-1) + + def _resolve_minimax_h3_num_outputs(value: Any) -> int: if value is None: return 1 @@ -1532,6 +1567,12 @@ class MiniMaxH3Pipeline( audio_condition_lengths: list[int] | None = None, keyframe_frame_indices: list[int] | None = None, pad_seq_len: int | None = None, + video_edit_clean_rows: torch.Tensor | None = None, + video_edit_mask_rows: torch.Tensor | None = None, + video_edit_restore_mask_rows: torch.Tensor | None = None, + audio_edit_clean_rows: torch.Tensor | None = None, + audio_edit_mask_rows: torch.Tensor | None = None, + audio_edit_restore_mask_rows: torch.Tensor | None = None, ) -> dict[str, Any]: """Build the packed layout, initial rows, anchors, and sigma schedules. @@ -1640,6 +1681,41 @@ class MiniMaxH3Pipeline( full_audio[branch.audio_update_mask] = initial_audio initial_audio = full_audio + video_edit = None + video_edit_values = (video_edit_clean_rows, video_edit_mask_rows, video_edit_restore_mask_rows) + if any(value is None for value in video_edit_values) and any(value is not None for value in video_edit_values): + raise ValueError("video edit clean, model-mask, and restore-mask rows must be provided together") + if ( + video_edit_clean_rows is not None + and video_edit_mask_rows is not None + and video_edit_restore_mask_rows is not None + ): + target_noise = initial_video[branch.update_mask].to(device=self.device, dtype=torch.float32) + clean = video_edit_clean_rows.to(device=self.device, dtype=torch.float32) + video_edit = MiniMaxH3LatentEdit.from_rows( + clean, + MINIMAX_H3_IMGVID_COND_TIMESTEP * clean + (1.0 - MINIMAX_H3_IMGVID_COND_TIMESTEP) * target_noise, + video_edit_mask_rows, + video_edit_restore_mask_rows, + ) + + audio_edit = None + audio_edit_values = (audio_edit_clean_rows, audio_edit_mask_rows, audio_edit_restore_mask_rows) + if any(value is None for value in audio_edit_values) and any(value is not None for value in audio_edit_values): + raise ValueError("audio edit clean, model-mask, and restore-mask rows must be provided together") + if ( + audio_edit_clean_rows is not None + and audio_edit_mask_rows is not None + and audio_edit_restore_mask_rows is not None + ): + clean = audio_edit_clean_rows.to(device=self.device, dtype=torch.float32) + audio_edit = MiniMaxH3LatentEdit.from_rows( + clean, + clean, + audio_edit_mask_rows, + audio_edit_restore_mask_rows, + ) + video_sigmas = minimax_h3_time_shift_sigmas( num_steps=num_steps, shift_scale=video_shift, @@ -1664,6 +1740,8 @@ class MiniMaxH3Pipeline( ), "sigmas_video": video_sigmas, "sigmas_audio": audio_sigmas, + "video_edit": video_edit, + "audio_edit": audio_edit, } def _unpack_denoised_rows( @@ -1722,6 +1800,12 @@ class MiniMaxH3Pipeline( audio_condition_lengths: list[int] | None = None, keyframe_frame_indices: list[int] | None = None, pad_seq_len: int | None = None, + video_edit_clean_rows: torch.Tensor | None = None, + video_edit_mask_rows: torch.Tensor | None = None, + video_edit_restore_mask_rows: torch.Tensor | None = None, + audio_edit_clean_rows: torch.Tensor | None = None, + audio_edit_mask_rows: torch.Tensor | None = None, + audio_edit_restore_mask_rows: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: inputs = self._build_denoise_inputs( task=task, @@ -1746,6 +1830,12 @@ class MiniMaxH3Pipeline( audio_condition_lengths=audio_condition_lengths, keyframe_frame_indices=keyframe_frame_indices, pad_seq_len=pad_seq_len, + video_edit_clean_rows=video_edit_clean_rows, + video_edit_mask_rows=video_edit_mask_rows, + video_edit_restore_mask_rows=video_edit_restore_mask_rows, + audio_edit_clean_rows=audio_edit_clean_rows, + audio_edit_mask_rows=audio_edit_mask_rows, + audio_edit_restore_mask_rows=audio_edit_restore_mask_rows, ) branch = inputs["branch"] transformer = self._transformer_for_task(task) @@ -1763,6 +1853,8 @@ class MiniMaxH3Pipeline( device=self.device, imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + video_edit=inputs["video_edit"], + audio_edit=inputs["audio_edit"], on_step=lambda step, video, audio: progress.update(), ) @@ -1997,7 +2089,14 @@ class MiniMaxH3Pipeline( if world_size == 1: assert conditioning is not None return conditioning - tensor_names = ("visual_condition", "audio_condition") + tensor_names = ( + "visual_condition", + "audio_condition", + "video_edit_clean_rows", + "video_edit_mask", + "audio_edit_clean_rows", + "audio_edit_mask", + ) header = [None] if rank == 0: assert conditioning is not None @@ -2187,6 +2286,50 @@ class MiniMaxH3Pipeline( visual_shapes = list(conditioning.visual_condition_shapes) or None audio_lengths = list(conditioning.audio_condition_lengths) or None + latent_edit: dict[str, torch.Tensor | None] = { + "video_edit_clean_rows": None, + "video_edit_mask_rows": None, + "video_edit_restore_mask_rows": None, + "audio_edit_clean_rows": None, + "audio_edit_mask_rows": None, + "audio_edit_restore_mask_rows": None, + } + try: + if conditioning.video_edit_clean_rows is not None: + if conditioning.video_edit_mask is None: + raise ValueError("MiniMax H3 video edit rows require a mask") + video_mask = minimax_h3_video_edit_masks( + conditioning.video_edit_mask, + latent_t=conditioning.latent_t, + latent_h=conditioning.height // 16, + latent_w=conditioning.width // 16, + ) + latent_edit.update( + video_edit_clean_rows=conditioning.video_edit_clean_rows, + video_edit_mask_rows=video_mask.model_mask_rows, + video_edit_restore_mask_rows=video_mask.restore_mask_rows, + ) + if conditioning.audio_edit_clean_rows is not None: + if conditioning.audio_edit_mask is None: + raise ValueError("MiniMax H3 audio edit rows require a mask") + audio_mask = minimax_h3_audio_edit_masks( + conditioning.audio_edit_mask, + audio_t=conditioning.audio_t, + ) + model_mask, restore_mask = _expose_padded_audio_tail( + conditioning.audio_edit_source_t, + target_audio_t=conditioning.audio_t, + mask_rows=audio_mask.model_mask_rows, + restore_mask_rows=audio_mask.restore_mask_rows, + ) + latent_edit.update( + audio_edit_clean_rows=conditioning.audio_edit_clean_rows, + audio_edit_mask_rows=model_mask, + audio_edit_restore_mask_rows=restore_mask, + ) + except ValueError as exc: + raise OmniClientError(str(exc)) from exc + base_schedule, num_steps = self._resolve_sigma_positions(task, sampling) quality_plan = self._quality_policy.resolve( quality=sampling.quality, @@ -2239,6 +2382,7 @@ class MiniMaxH3Pipeline( if extra.get("preencode_mp4", False) else None ), + **latent_edit, } @staticmethod @@ -2411,6 +2555,8 @@ class MiniMaxH3Pipeline( _STEP_AUDIO_ANCHOR: audio_anchor, _STEP_SIGMAS_VIDEO: sigmas_video, _STEP_SIGMAS_AUDIO: sigmas_audio, + _STEP_VIDEO_EDIT: inputs.get("video_edit"), + _STEP_AUDIO_EDIT: inputs.get("audio_edit"), _STEP_SHAPE: { "height": context["height"], "width": context["width"], @@ -2444,12 +2590,40 @@ class MiniMaxH3Pipeline( batch_states = list(states if states is not None else input_batch.states) branches = [state.extra[_STEP_BRANCH] for state in batch_states] - video_rows = [state.latents for state in batch_states] - audio_rows = [state.extra[_STEP_AUDIO_ROWS] for state in batch_states] schedules = [_minimax_h3_step_schedule(state) for state in batch_states] transformers = [state.extra[_STEP_TRANSFORMER] for state in batch_states] mixed_transformers = len({id(transformer) for transformer in transformers}) > 1 + video_rows: list[torch.Tensor] = [] + audio_rows: list[torch.Tensor] = [] + video_target_timesteps: list[torch.Tensor | None] = [] + audio_target_timesteps: list[torch.Tensor | None] = [] + for state, branch, schedule in zip(batch_states, branches, schedules, strict=True): + video_edit = state.extra.get(_STEP_VIDEO_EDIT) + request_video, request_video_timesteps = minimax_h3_prepare_edit_rows( + state.latents, + branch.update_mask_dev, + video_edit, + schedule["t_video"], + schedule["imgvid_cond_timestep"], + sigma=schedule["sigma_video"], + ) + video_target_timesteps.append(request_video_timesteps) + video_rows.append(request_video) + + state_audio = state.extra[_STEP_AUDIO_ROWS] + audio_edit = state.extra.get(_STEP_AUDIO_EDIT) + request_audio, request_audio_timesteps = minimax_h3_prepare_edit_rows( + state_audio, + branch.audio_update_mask_dev, + audio_edit, + schedule["t_audio"], + schedule["audio_ref_cond_timestep"], + sigma=schedule["sigma_audio"], + ) + audio_target_timesteps.append(request_audio_timesteps) + audio_rows.append(request_audio) + # Both execution modes must publish denoise progress for step-gated # attention features. Requests can differ in both step index and sigma # schedule, so a batch that is not at one single point has nothing to @@ -2495,6 +2669,8 @@ class MiniMaxH3Pipeline( t_audio=schedules[index]["t_audio"], imgvid_cond_timestep=schedules[index]["imgvid_cond_timestep"], audio_ref_cond_timestep=schedules[index]["audio_ref_cond_timestep"], + video_target_timesteps=video_target_timesteps[index], + audio_target_timesteps=audio_target_timesteps[index], ) request_video, request_audio = transformers[index](**forward_kwargs) video_parts.append(request_video) @@ -2510,6 +2686,8 @@ class MiniMaxH3Pipeline( t_audio=[schedule["t_audio"] for schedule in schedules], imgvid_cond_timesteps=[schedule["imgvid_cond_timestep"] for schedule in schedules], audio_ref_cond_timesteps=[schedule["audio_ref_cond_timestep"] for schedule in schedules], + video_target_timesteps=video_target_timesteps, + audio_target_timesteps=audio_target_timesteps, ) logger.debug( "MiniMax H3 denoise step: %d request(s) packed into %d rows", @@ -2543,11 +2721,19 @@ class MiniMaxH3Pipeline( audio_anchor = state.extra[_STEP_AUDIO_ANCHOR] device = video_rows.device - x0_video = minimax_h3_rf_v_to_x0( - video_rows[update], - noise_pred.float()[update], - torch.tensor(schedule["t_video"], dtype=torch.float32, device=device), - ) + video_edit = state.extra.get(_STEP_VIDEO_EDIT) + if video_edit is None: + x0_video = minimax_h3_rf_v_to_x0( + video_rows[update], + noise_pred.float()[update], + torch.tensor(schedule["t_video"], dtype=torch.float32, device=device), + ) + else: + x0_video = video_edit.x0( + video_edit.model_rows(video_rows[update]), + noise_pred.float()[update], + schedule["t_video"], + ) new_video = minimax_h3_euler_eta0_step( video_rows[update], x0_video, @@ -2559,11 +2745,19 @@ class MiniMaxH3Pipeline( if cond_anchor is not None: video_rows[~update] = cond_anchor # per-step imgvid cond reset - x0_audio = minimax_h3_rf_v_to_x0( - audio_rows[audio_update], - audio_noise_pred.float()[audio_update], - torch.tensor(schedule["t_audio"], dtype=torch.float32, device=device), - ) + audio_edit = state.extra.get(_STEP_AUDIO_EDIT) + if audio_edit is None: + x0_audio = minimax_h3_rf_v_to_x0( + audio_rows[audio_update], + audio_noise_pred.float()[audio_update], + torch.tensor(schedule["t_audio"], dtype=torch.float32, device=device), + ) + else: + x0_audio = audio_edit.x0( + audio_edit.model_rows(audio_rows[audio_update]), + audio_noise_pred.float()[audio_update], + schedule["t_audio"], + ) new_audio = minimax_h3_euler_eta0_step( audio_rows[audio_update], x0_audio, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 7100bf67f..6541c11df 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -168,6 +168,7 @@ from vllm_omni.entrypoints.openai.serving_rl_rollout import ServingRLRollout from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler from vllm_omni.entrypoints.openai.serving_video import ( + LatentEditInput, OmniOpenAIServingVideo, ReferenceAudio, ReferenceImage, @@ -2623,6 +2624,7 @@ async def create_video( ReferenceVideo | None, ReferenceAudio | None, str | None, + LatentEditInput | None, ] = Depends(_parse_video_form), ) -> VideoResponse: """Create an asynchronous video generation job. @@ -2638,6 +2640,7 @@ async def create_video( reference_video, reference_audio, control_path, + latent_edit_input, ) = ctx ref = video_response_from_request(effective_model_name, request) await VIDEO_STORE.upsert(ref.id, ref) @@ -2651,6 +2654,7 @@ async def create_video( reference_audio, control_path, app_state=raw_request.app.state, + latent_edit_input=latent_edit_input, ) ) await VIDEO_TASKS.upsert(ref.id, task) @@ -2676,6 +2680,7 @@ async def create_video_sync( ReferenceVideo | None, ReferenceAudio | None, str | None, + LatentEditInput | None, ] = Depends(_parse_video_form), ) -> Response: """Synchronous video generation endpoint. @@ -2695,6 +2700,7 @@ async def create_video_sync( reference_video, reference_audio, control_path, + latent_edit_input, ) = ctx request_id = f"video_sync-{random_uuid()}" raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id) @@ -2708,6 +2714,7 @@ async def create_video_sync( reference_image=reference_image, reference_video=reference_video, reference_audio=reference_audio, + latent_edit_input=latent_edit_input, ), timeout=VIDEO_SYNC_TIMEOUT_S, ), @@ -2731,7 +2738,7 @@ async def create_video_sync( detail=f"Video generation failed: {str(exc)}", ) from exc finally: - _cleanup_video_references(reference_video, reference_audio, control_path) + _cleanup_video_references(reference_video, reference_audio, control_path, latent_edit_input) inference_time_s = time.perf_counter() - started_at return Response( diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 717496362..c90a4d162 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -20,7 +20,7 @@ from vllm.engine.protocol import EngineClient from vllm.logger import init_logger from vllm_omni.diffusion.data import is_diffusion_request_started_output -from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata +from vllm_omni.diffusion.model_metadata import DiffusionModelMetadata, get_diffusion_model_metadata from vllm_omni.diffusion.utils.media_utils import count_mp4_frames, normalize_preencode_batch_frames from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni from vllm_omni.entrypoints.openai.protocol.videos import ( @@ -104,6 +104,17 @@ class ReferenceAudio: cleanup_paths: tuple[str, ...] = () +@dataclass +class LatentEditInput: + """Request-scoped source media and masks for latent-mask editing.""" + + source_video: str | None = None + source_audio: str | None = None + video_noise_mask: Any | None = None + audio_noise_mask: Any | None = None + cleanup_paths: tuple[str, ...] = () + + @dataclass class VideoGenerationArtifacts: """Normalized outputs and profiler metadata extracted from one request.""" @@ -167,6 +178,15 @@ class OmniOpenAIServingVideo: return get_od_config() return getattr(self._engine_client, "od_config", None) + def _resolve_diffusion_metadata( + self, + ) -> tuple[OmniDiffusionConfig | SimpleNamespace | None, tuple[DiffusionModelMetadata, ...]]: + """Resolve capability metadata from the runtime and diffusion stages.""" + od_config = self._resolve_diffusion_od_config() + model_archs = [None if od_config is None else getattr(od_config, "model_class_name", None)] + model_archs.extend(_stage_diffusion_model_class_name(stage) for stage in self.stage_configs or ()) + return od_config, tuple(get_diffusion_model_metadata(model_arch) for model_arch in model_archs) + def _resolve_video_generation_defaults( self, request: VideoGenerationRequest, @@ -211,19 +231,18 @@ class OmniOpenAIServingVideo: @property def supports_mixed_reference_inputs(self) -> bool: """Return whether the configured diffusion model accepts mixed refs.""" - od_config = self._resolve_diffusion_od_config() + od_config, metadata = self._resolve_diffusion_metadata() if od_config is None: return False capability = getattr(od_config, "supports_mixed_reference_inputs", None) - model_class_name = getattr(od_config, "model_class_name", None) - model_archs = [model_class_name] - for stage_config in self.stage_configs or (): - model_archs.append(_stage_diffusion_model_class_name(stage_config)) - metadata_capability = any( - get_diffusion_model_metadata(model_arch).supports_mixed_reference_inputs for model_arch in model_archs - ) - return capability is True or metadata_capability + return capability is True or any(item.supports_mixed_reference_inputs for item in metadata) + + @property + def supports_latent_mask_editing(self) -> bool: + """Return whether the configured diffusion model accepts latent edits.""" + _, metadata = self._resolve_diffusion_metadata() + return any(item.supports_latent_mask_editing for item in metadata) @property def supported_control_upload_types(self) -> frozenset[str]: @@ -233,14 +252,11 @@ class OmniOpenAIServingVideo: generic video API isolated from model-specific controls unless a model explicitly opts into the ``control_path`` contract in metadata. """ - od_config = self._resolve_diffusion_od_config() - model_archs = [None if od_config is None else getattr(od_config, "model_class_name", None)] - for stage_config in self.stage_configs or (): - model_archs.append(_stage_diffusion_model_class_name(stage_config)) + _, metadata = self._resolve_diffusion_metadata() supported: set[str] = set() - for model_arch in model_archs: - supported.update(get_diffusion_model_metadata(model_arch).supported_control_upload_types) + for item in metadata: + supported.update(item.supported_control_upload_types) return frozenset(supported) @classmethod @@ -271,6 +287,7 @@ class OmniOpenAIServingVideo: reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, on_started: Callable[[], Awaitable[None]] | None = None, + latent_edit_input: LatentEditInput | None = None, ) -> VideoGenerationArtifacts: """Run the generation pipeline and extract video/audio/profiler outputs.""" prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt, modalities=["video"]) @@ -351,6 +368,19 @@ class OmniOpenAIServingVideo: multi_modal_data["video"] = input_video if reference_audio is not None: multi_modal_data["audio"] = reference_audio.path + if latent_edit_input is not None: + multi_modal_data.update( + { + key: value + for key, value in { + "source_video": latent_edit_input.source_video, + "source_audio": latent_edit_input.source_audio, + "video_noise_mask": latent_edit_input.video_noise_mask, + "audio_noise_mask": latent_edit_input.audio_noise_mask, + }.items() + if value is not None + } + ) if multi_modal_data: prompt["multi_modal_data"] = multi_modal_data if vp.width is not None and vp.height is not None: @@ -489,6 +519,7 @@ class OmniOpenAIServingVideo: reference_image: ReferenceImage | None = None, reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, + latent_edit_input: LatentEditInput | None = None, ) -> VideoGenerationResponse: artifacts = await self._run_and_extract( request, @@ -496,6 +527,7 @@ class OmniOpenAIServingVideo: reference_image=reference_image, reference_video=reference_video, reference_audio=reference_audio, + latent_edit_input=latent_edit_input, ) video_codec_options = {"preset": "ultrafast", "threads": "0"} @@ -541,6 +573,7 @@ class OmniOpenAIServingVideo: reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, on_started: Callable[[], Awaitable[None]] | None = None, + latent_edit_input: LatentEditInput | None = None, ) -> tuple[bytes, dict[str, float], float, VideoAction | None, dict[str, object]]: """Generate a video and return raw MP4 bytes, bypassing base64 encoding.""" artifacts = await self._run_and_extract( @@ -550,6 +583,7 @@ class OmniOpenAIServingVideo: reference_video=reference_video, reference_audio=reference_audio, on_started=on_started, + latent_edit_input=latent_edit_input, ) if len(artifacts.videos) > 1: logger.warning( diff --git a/vllm_omni/entrypoints/openai/video/generation/helpers.py b/vllm_omni/entrypoints/openai/video/generation/helpers.py index 0ffd1ee83..5ed02443f 100644 --- a/vllm_omni/entrypoints/openai/video/generation/helpers.py +++ b/vllm_omni/entrypoints/openai/video/generation/helpers.py @@ -54,6 +54,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( VideoResponse, ) from vllm_omni.entrypoints.openai.serving_video import ( + LatentEditInput, OmniOpenAIServingVideo, ReferenceAudio, ReferenceImage, @@ -93,6 +94,9 @@ CONTROL_REFERENCE_IMAGE_SUFFIXES = frozenset({".bmp", ".gif", ".jpg", ".jpeg", " CONTROL_REFERENCE_VIDEO_SUFFIXES = frozenset({".mkv", ".mov", ".mp4", ".webm"}) CONTROL_REFERENCE_MAX_BYTES = 512 * 1024 * 1024 +LATENT_EDIT_SOURCE_MAX_BYTES = 512 * 1024 * 1024 +LATENT_EDIT_MASK_FILE_MAX_BYTES = 8 * 1024 * 1024 + VIDEO_SYNC_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_VIDEO_SYNC_TIMEOUT", 600.0)) @@ -126,6 +130,79 @@ def _parse_form_json(value: str | None, expected_type: type | None = None) -> An return parsed +def _parse_latent_edit_mask_json( + value: str, + *, + field_name: str, +) -> Any: + """Decode one request mask; Stage 0 owns its semantic validation.""" + try: + parsed = _parse_form_json(value) + # Bound container nesting at the parsing boundary without duplicating + # Stage 0's numeric, range, or shape validation. + pending = [(parsed, 0)] + while pending: + item, depth = pending.pop() + if not isinstance(item, (list, dict)): + continue + if depth >= 64: + raise ValueError("JSON nesting exceeds 64 levels") + children = item.values() if isinstance(item, dict) else item + pending.extend((child, depth + 1) for child in children) + return parsed + except (RecursionError, ValueError) as exc: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} contains invalid JSON or is nested too deeply.", + ) from exc + + +async def _read_latent_edit_mask_json( + value: UploadFile | None, + *, + field_name: str, +) -> Any | None: + """Read a mask from a bounded UTF-8 JSON file part.""" + if value is None: + return None + + try: + declared_size = value.size + if isinstance(declared_size, Integral) and int(declared_size) > LATENT_EDIT_MASK_FILE_MAX_BYTES: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} JSON exceeds the {LATENT_EDIT_MASK_FILE_MAX_BYTES // (1024 * 1024)} MiB limit.", + ) + + payload = await value.read(LATENT_EDIT_MASK_FILE_MAX_BYTES + 1) + if len(payload) > LATENT_EDIT_MASK_FILE_MAX_BYTES: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} JSON exceeds the {LATENT_EDIT_MASK_FILE_MAX_BYTES // (1024 * 1024)} MiB limit.", + ) + try: + text = payload.decode("utf-8") + except UnicodeDecodeError as exc: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} JSON file must be UTF-8 encoded.", + ) from exc + if not text: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} JSON file must not be empty.", + ) + parsed = _parse_latent_edit_mask_json(text, field_name=field_name) + if parsed is None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} JSON file must contain a mask value.", + ) + return parsed + finally: + await value.close() + + def _config_get(config: Any, key: str, default: Any = None) -> Any: if isinstance(config, dict): return config.get(key, default) @@ -278,6 +355,7 @@ def _cleanup_video_references( reference_video: ReferenceVideo | None, reference_audio: ReferenceAudio | None, control_path: str | None = None, + latent_edit_input: LatentEditInput | None = None, ) -> None: if reference_video is not None: for path in reference_video.cleanup_paths: @@ -290,6 +368,10 @@ def _cleanup_video_references( os.unlink(path) if control_path is not None and os.path.exists(control_path): os.unlink(control_path) + if latent_edit_input is not None: + for path in latent_edit_input.cleanup_paths: + if os.path.exists(path): + os.unlink(path) def _unpack_video_generation_result( @@ -320,11 +402,12 @@ async def _run_video_generation_job( reference_audio: ReferenceAudio | None = None, control_path: str | None = None, app_state: Any | None = None, + latent_edit_input: LatentEditInput | None = None, ) -> None: job = await VIDEO_STORE.get(video_id) if job is None: logger.warning("Video job %s missing before generation task started; skipping", video_id) - _cleanup_video_references(reference_video, reference_audio, control_path) + _cleanup_video_references(reference_video, reference_audio, control_path, latent_edit_input) return started_at = time.perf_counter() @@ -341,6 +424,7 @@ async def _run_video_generation_job( reference_video=reference_video, reference_audio=reference_audio, on_started=_mark_started, + latent_edit_input=latent_edit_input, ) ) @@ -400,7 +484,7 @@ async def _run_video_generation_job( await VIDEO_STORE.pop(video_id) raise finally: - _cleanup_video_references(reference_video, reference_audio, control_path) + _cleanup_video_references(reference_video, reference_audio, control_path, latent_edit_input) async def _persist_uploaded_video_references(uploads: list[UploadFile]) -> list[str]: @@ -423,6 +507,47 @@ async def _persist_uploaded_video_references(uploads: list[UploadFile]) -> list[ return paths +async def _persist_uploaded_reference( + upload: UploadFile, + *, + field_name: str, + suffix: str, + max_bytes: int, +) -> str: + """Stream one bounded upload to request-scoped local storage.""" + declared_size = upload.size + if isinstance(declared_size, Integral) and int(declared_size) > max_bytes: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + ) + + fd, path = tempfile.mkstemp(prefix=f"vllm_omni_{field_name}_", suffix=suffix) + size = 0 + persisted = False + try: + with os.fdopen(fd, "wb") as output: + while chunk := await upload.read(min(1024 * 1024, max_bytes - size + 1)): + size += len(chunk) + if size > max_bytes: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + ) + output.write(chunk) + if size == 0: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"{field_name} must not be empty.", + ) + persisted = True + return path + finally: + if not persisted: + with suppress(OSError): + os.unlink(path) + + async def _persist_uploaded_control_reference( upload: UploadFile, *, @@ -439,38 +564,70 @@ async def _persist_uploaded_control_reference( ) if not suffix: suffix = ".png" if kind == "image" else ".mp4" + return await _persist_uploaded_reference( + upload, + field_name="control_reference", + suffix=suffix, + max_bytes=max_bytes, + ) - declared_size = getattr(upload, "size", None) - if isinstance(declared_size, Integral) and int(declared_size) > max_bytes: + +async def _persist_latent_edit_source( + upload: UploadFile, + *, + field_name: Literal["source_video", "source_audio"], + max_bytes: int = LATENT_EDIT_SOURCE_MAX_BYTES, +) -> str: + """Stream one latent-edit source upload to request-scoped storage.""" + suffix = Path(upload.filename or "").suffix.lower() + if field_name == "source_video": + allowed_suffixes = MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES + expected = "an MP4 or MOV file" + else: + allowed_suffixes = MINIMAX_H3_REFERENCE_AUDIO_SUFFIXES + expected = "a WAV or MP3 file" + if suffix not in allowed_suffixes: raise HTTPException( status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"control_reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + detail=f"{field_name} must be {expected}.", ) - fd, path = tempfile.mkstemp(prefix="vllm_omni_control_reference_", suffix=suffix) - size = 0 - persisted = False - try: - with os.fdopen(fd, "wb") as output: - while chunk := await upload.read(min(1024 * 1024, max_bytes - size + 1)): - size += len(chunk) - if size > max_bytes: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"control_reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", - ) - output.write(chunk) - if size == 0: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="control_reference must not be empty.", - ) - persisted = True - return path - finally: - if not persisted: - with suppress(OSError): - os.unlink(path) + return await _persist_uploaded_reference( + upload, + field_name=field_name, + suffix=suffix, + max_bytes=max_bytes, + ) + + +def _validate_latent_edit_inputs( + handler: OmniOpenAIServingVideo, + *, + source_video: UploadFile | None, + source_audio: UploadFile | None, + video_noise_mask: Any | None, + audio_noise_mask: Any | None, +) -> bool: + """Validate cross-field latent-edit requirements before persisting files.""" + has_any_input = any(value is not None for value in (source_video, source_audio, video_noise_mask, audio_noise_mask)) + if not has_any_input: + return False + if not handler.supports_latent_mask_editing: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Latent-mask editing is not supported by this model.", + ) + if video_noise_mask is None and audio_noise_mask is None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Latent-mask editing requires at least one of video_noise_mask or audio_noise_mask.", + ) + if source_audio is not None and audio_noise_mask is None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="source_audio requires an audio_noise_mask upload.", + ) + return True def _validate_control_upload( @@ -494,7 +651,7 @@ def _validate_control_upload( ) normalized_type = control_type.strip().lower() - supported_types = frozenset(getattr(handler, "supported_control_upload_types", ())) + supported_types = handler.supported_control_upload_types if normalized_type not in supported_types: supported = ", ".join(sorted(supported_types)) or "none" raise HTTPException( @@ -732,6 +889,10 @@ async def _parse_video_form( lora: str | None = Form(default=None), extra_params: str | None = Form(default=None), return_stage_metrics: bool | None = Form(default=None), + source_video: UploadFile | None = File(default=None), + source_audio: UploadFile | None = File(default=None), + video_noise_mask: UploadFile | None = File(default=None), + audio_noise_mask: UploadFile | None = File(default=None), ) -> tuple[ VideoGenerationRequest, "OmniOpenAIServingVideo", @@ -740,6 +901,7 @@ async def _parse_video_form( ReferenceVideo | None, ReferenceAudio | None, str | None, + LatentEditInput | None, ]: """FastAPI dependency that parses video form data, validates inputs, resolves the handler, and decodes any reference image. @@ -751,6 +913,14 @@ async def _parse_video_form( parsed_image_reference = _parse_form_json(image_reference) parsed_video_reference = _parse_form_json(video_reference) parsed_audio_reference = _parse_form_json(audio_reference) + parsed_video_noise_mask = await _read_latent_edit_mask_json( + video_noise_mask, + field_name="video_noise_mask", + ) + parsed_audio_noise_mask = await _read_latent_edit_mask_json( + audio_noise_mask, + field_name="audio_noise_mask", + ) if input_references and any( item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference) @@ -835,6 +1005,14 @@ async def _parse_video_form( detail=f"Video generation setup failed: {str(e)}", ) + has_latent_edit = _validate_latent_edit_inputs( + handler, + source_video=source_video, + source_audio=source_audio, + video_noise_mask=parsed_video_noise_mask, + audio_noise_mask=parsed_audio_noise_mask, + ) + normalized_control_type = _validate_control_upload(handler, request, control_reference, control_type) if normalized_control_type is not None: # Make the selected transfer mode visible while choosing the model's @@ -842,7 +1020,7 @@ async def _parse_video_form( # until reference parsing succeeds, preserving failure cleanup. _attach_control_upload(request, normalized_control_type) - supports_mixed_reference_inputs = bool(getattr(handler, "supports_mixed_reference_inputs", False)) + supports_mixed_reference_inputs = handler.supports_mixed_reference_inputs if input_reference is not None: input_reference_bytes = await _read_upload_limited( input_reference, @@ -976,6 +1154,38 @@ async def _parse_video_form( _cleanup_video_references(reference_video, reference_audio, control_path) raise + latent_edit_input: LatentEditInput | None = None + if has_latent_edit: + source_paths: list[str] = [] + try: + source_video_path = None + source_audio_path = None + if source_video is not None: + source_video_path = await _persist_latent_edit_source( + source_video, + field_name="source_video", + ) + source_paths.append(source_video_path) + if source_audio is not None: + source_audio_path = await _persist_latent_edit_source( + source_audio, + field_name="source_audio", + ) + source_paths.append(source_audio_path) + latent_edit_input = LatentEditInput( + source_video=source_video_path, + source_audio=source_audio_path, + video_noise_mask=parsed_video_noise_mask, + audio_noise_mask=parsed_audio_noise_mask, + cleanup_paths=tuple(source_paths), + ) + except (asyncio.CancelledError, Exception): + for path in source_paths: + with suppress(OSError): + os.unlink(path) + _cleanup_video_references(reference_video, reference_audio, control_path) + raise + return ( request, handler, @@ -984,4 +1194,5 @@ async def _parse_video_form( reference_video, reference_audio, control_path, + latent_edit_input, ) diff --git a/vllm_omni/model_executor/models/minimax_h3/conditioning.py b/vllm_omni/model_executor/models/minimax_h3/conditioning.py index 761292dc3..b54d36751 100644 --- a/vllm_omni/model_executor/models/minimax_h3/conditioning.py +++ b/vllm_omni/model_executor/models/minimax_h3/conditioning.py @@ -4,6 +4,7 @@ from __future__ import annotations +import math from collections.abc import Mapping, Sequence from dataclasses import dataclass, fields from typing import Any @@ -123,7 +124,7 @@ class MiniMaxH3TextConditioning: MINIMAX_H3_ENCODER_REQUEST_KEY = "minimax_h3_encoder_request" MINIMAX_H3_ENCODER_LAYOUT_KEY = "minimax_h3_encoder_layout" -STAGE_SCHEMA_VERSION = 1 +STAGE_SCHEMA_VERSION = 2 _ENCODER_WIRE_SCHEMA_ID = 1 _TASK_TO_CODE = {"t2va": 1, "fl2va": 2, "ref2va": 3} @@ -144,6 +145,104 @@ _AUDIO_CONDITION_WIDTH = 32 _AUDIO_CONDITION_CHANNELS = 2 +def _validate_edit_mask_structure( + mask: torch.Tensor, + *, + name: str, + shape: tuple[int, ...], +) -> None: + """Validate only the internal tensor contract after encoder-boundary parsing.""" + if mask.dtype != torch.float32 or tuple(mask.shape) != shape: + raise ValueError(f"MiniMax H3 {name} edit mask must be FP32 with shape {shape}") + _validate_contiguous_strided_layout(f"{name} edit mask", mask) + + +def _validate_edit_tensors( + clean_rows: torch.Tensor | None, + raw_mask: torch.Tensor | None, + *, + name: str, + width: int, + mask_shape: tuple[int, ...], +) -> int: + values = (clean_rows, raw_mask) + if all(value is None for value in values): + return 0 + if clean_rows is None or raw_mask is None: + raise ValueError(f"MiniMax H3 {name} edit requires clean rows and a raw mask tensor") + if clean_rows.dtype != torch.float32 or clean_rows.ndim != 2 or clean_rows.shape[1] != width: + raise ValueError(f"MiniMax H3 {name} edit rows must be FP32 with shape [rows, {width}]") + row_count = int(clean_rows.shape[0]) + if row_count <= 0: + raise ValueError(f"MiniMax H3 {name} edit rows must not be empty") + _validate_edit_mask_structure(raw_mask, name=name, shape=mask_shape) + if raw_mask.device != clean_rows.device: + raise ValueError(f"MiniMax H3 {name} edit tensors must share one device") + _validate_contiguous_strided_layout(f"{name} edit rows", clean_rows) + return row_count + + +def _packed_vector_rows(value: torch.Tensor, *, width: int) -> torch.Tensor: + flat = value.reshape(-1) + padding = (-int(flat.numel())) % width + if padding: + flat = torch.cat((flat, flat.new_zeros(padding))) + return flat.reshape(-1, width) + + +def _pack_condition_slot( + condition: torch.Tensor | None, + clean_rows: torch.Tensor | None, + raw_mask: torch.Tensor | None, + *, + width: int, +) -> torch.Tensor: + if clean_rows is None: + return condition if condition is not None else torch.empty((0,), dtype=torch.float32) + if raw_mask is None: + raise ValueError("MiniMax H3 edit rows require a raw mask tensor") + values = [clean_rows, _packed_vector_rows(raw_mask, width=width)] + if condition is not None: + values.insert(0, condition) + return torch.cat(values).contiguous() + + +def _unpack_condition_slot( + value: torch.Tensor, + *, + condition_rows: int, + edit_rows: int, + mask_shape: tuple[int, ...], + width: int, + name: str, +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + mask_values = math.prod(mask_shape) if mask_shape else int(edit_rows > 0) + mask_wire_rows = (mask_values + width - 1) // width + expected_rows = condition_rows + edit_rows + mask_wire_rows + if expected_rows == 0: + if value.dtype != torch.float32 or tuple(value.shape) != (0,): + raise ValueError(f"MiniMax H3 empty {name} condition slot must be FP32 with shape [0]") + return None, None, None + if value.dtype != torch.float32 or value.ndim != 2 or tuple(value.shape) != (expected_rows, width): + raise ValueError( + f"MiniMax H3 packed {name} condition must be FP32 with shape " + f"[{expected_rows}, {width}], got {value.dtype} {tuple(value.shape)}" + ) + cursor = 0 + condition = value[:condition_rows].contiguous() if condition_rows else None + cursor += condition_rows + if edit_rows == 0: + return condition, None, None + clean_rows = value[cursor : cursor + edit_rows].contiguous() + cursor += edit_rows + raw_mask = value[cursor : cursor + mask_wire_rows].reshape(-1)[:mask_values].clone() + if mask_shape: + raw_mask = raw_mask.reshape(mask_shape) + else: + raw_mask = raw_mask.reshape(()) + return condition, clean_rows, raw_mask + + def _wire_layout(payload: Mapping[str, Any]) -> torch.Tensor: private_metadata = payload.get("kv_metadata") if not isinstance(private_metadata, Mapping): @@ -259,6 +358,10 @@ class MiniMaxH3EncoderMediaInput: video_audios: tuple[tuple[torch.Tensor, int] | None, ...] = () audios: tuple[tuple[torch.Tensor, int], ...] = () keyframe_frame_indices: tuple[int, ...] = () + video_edit: torch.Tensor | None = None + video_edit_mask: torch.Tensor | None = None + audio_edit: tuple[torch.Tensor, int] | None = None + audio_edit_mask: torch.Tensor | None = None @classmethod def from_mm_tensors( @@ -331,6 +434,37 @@ class MiniMaxH3EncoderMediaInput: if waveform.ndim not in (1, 2): raise ValueError("MiniMax H3 audio must have shape [samples] or [channels, samples]") audios.append((waveform, sample_rate)) + + video_edit = None + video_edit_mask = None + if bool(metadata.get("has_video_edit", False)): + if cursor + 2 > len(tensors): + raise ValueError("MiniMax H3 encoder video edit input is truncated") + video_edit, video_edit_mask = tensors[cursor : cursor + 2] + cursor += 2 + if tuple(video_edit.shape) != (num_frames, height, width, 3): + raise ValueError( + "MiniMax H3 encoder video edit must have shape " + f"[{num_frames}, {height}, {width}, 3], got {tuple(video_edit.shape)}" + ) + _validate_edit_mask_structure( + video_edit_mask, + name="video", + shape=(latent_t, height // 16, width // 16), + ) + + audio_edit = None + audio_edit_mask = None + audio_edit_sample_rate = int(metadata.get("audio_edit_sample_rate", 0)) + if audio_edit_sample_rate: + if cursor + 2 > len(tensors): + raise ValueError("MiniMax H3 encoder audio edit input is truncated") + waveform, audio_edit_mask = tensors[cursor : cursor + 2] + cursor += 2 + if waveform.ndim not in (1, 2) or audio_edit_sample_rate <= 0: + raise ValueError("MiniMax H3 encoder audio edit requires a waveform and positive sample rate") + _validate_edit_mask_structure(audio_edit_mask, name="audio", shape=(2, audio_t)) + audio_edit = (waveform, audio_edit_sample_rate) if cursor != len(tensors): raise ValueError(f"MiniMax H3 encoder media input has {len(tensors) - cursor} trailing tensors") return cls( @@ -345,6 +479,10 @@ class MiniMaxH3EncoderMediaInput: video_audios=tuple(video_audios), audios=tuple(audios), keyframe_frame_indices=tuple(int(value) for value in metadata.get("keyframe_frame_indices", ())), + video_edit=video_edit, + video_edit_mask=video_edit_mask, + audio_edit=audio_edit, + audio_edit_mask=audio_edit_mask, ) def to_mm_tensors(self) -> list[torch.Tensor]: @@ -362,6 +500,28 @@ class MiniMaxH3EncoderMediaInput: audio_items = [item for item in self.video_audios if item is not None] + list(self.audios) if any(waveform.ndim not in (1, 2) or int(sample_rate) <= 0 for waveform, sample_rate in audio_items): raise ValueError("MiniMax H3 encoder audio must have a waveform and positive sample rate") + video_edit_values = (self.video_edit, self.video_edit_mask) + if any(value is not None for value in video_edit_values): + if self.video_edit is None or self.video_edit_mask is None: + raise ValueError("MiniMax H3 encoder video edit requires a source and mask") + if tuple(self.video_edit.shape) != (self.num_frames, self.height, self.width, 3): + raise ValueError( + "MiniMax H3 encoder video edit must have shape " + f"[{self.num_frames}, {self.height}, {self.width}, 3], got {tuple(self.video_edit.shape)}" + ) + _validate_edit_mask_structure( + self.video_edit_mask, + name="video", + shape=(self.latent_t, self.height // 16, self.width // 16), + ) + audio_edit_values = (self.audio_edit, self.audio_edit_mask) + if any(value is not None for value in audio_edit_values): + if self.audio_edit is None or self.audio_edit_mask is None: + raise ValueError("MiniMax H3 encoder audio edit requires a source and mask") + waveform, sample_rate = self.audio_edit + if waveform.ndim not in (1, 2) or int(sample_rate) <= 0: + raise ValueError("MiniMax H3 encoder audio edit requires a waveform and positive sample rate") + _validate_edit_mask_structure(self.audio_edit_mask, name="audio", shape=(2, self.audio_t)) tensors = [*self.images, *self.videos] for item in self.video_audios: if item is None: @@ -369,10 +529,18 @@ class MiniMaxH3EncoderMediaInput: waveform, _sample_rate = item tensors.append(waveform) tensors.extend(waveform for waveform, _sample_rate in self.audios) + if self.video_edit is not None: + if self.video_edit_mask is None: + raise ValueError("MiniMax H3 encoder video edit requires a mask") + tensors.extend((self.video_edit, self.video_edit_mask)) + if self.audio_edit is not None: + if self.audio_edit_mask is None: + raise ValueError("MiniMax H3 encoder audio edit requires a mask") + tensors.extend((self.audio_edit[0], self.audio_edit_mask)) return tensors def to_metadata(self) -> dict[str, Any]: - return { + metadata = { "task": self.task, "height": self.height, "width": self.width, @@ -387,6 +555,11 @@ class MiniMaxH3EncoderMediaInput: "audio_sample_rates": [sample_rate for _waveform, sample_rate in self.audios], "keyframe_frame_indices": list(self.keyframe_frame_indices), } + if self.video_edit is not None: + metadata["has_video_edit"] = True + if self.audio_edit is not None: + metadata["audio_edit_sample_rate"] = int(self.audio_edit[1]) + return metadata @dataclass(frozen=True) @@ -403,6 +576,11 @@ class MiniMaxH3EncoderMediaConditioning: audio_condition_lengths: tuple[int, ...] = () ref_blocks: tuple[dict[str, Any], ...] = () keyframe_frame_indices: tuple[int, ...] = () + video_edit_clean_rows: torch.Tensor | None = None + video_edit_mask: torch.Tensor | None = None + audio_edit_clean_rows: torch.Tensor | None = None + audio_edit_mask: torch.Tensor | None = None + audio_edit_source_t: int = 0 def to_omni_components( self, @@ -418,7 +596,33 @@ class MiniMaxH3EncoderMediaConditioning: self.audio_condition, self.audio_condition_lengths, ) - empty = torch.empty((0,), dtype=torch.float32) + video_mask_shape = (self.latent_t, self.height // 16, self.width // 16) + audio_mask_shape = (2, self.audio_t) + video_edit_rows = _validate_edit_tensors( + self.video_edit_clean_rows, + self.video_edit_mask, + name="video", + width=_VIDEO_CONDITION_WIDTH, + mask_shape=video_mask_shape, + ) + audio_edit_rows = _validate_edit_tensors( + self.audio_edit_clean_rows, + self.audio_edit_mask, + name="audio", + width=_AUDIO_CONDITION_WIDTH, + mask_shape=audio_mask_shape, + ) + expected_video_rows = self.latent_t * (self.height // 32) * (self.width // 32) + if video_edit_rows not in (0, expected_video_rows): + raise ValueError(f"MiniMax H3 video edit has {video_edit_rows} rows, expected {expected_video_rows}") + expected_audio_rows = self.audio_t * _AUDIO_CONDITION_CHANNELS + if audio_edit_rows not in (0, expected_audio_rows): + raise ValueError(f"MiniMax H3 audio edit has {audio_edit_rows} rows, expected {expected_audio_rows}") + if audio_edit_rows == 0: + if self.audio_edit_source_t != 0: + raise ValueError("MiniMax H3 audio_edit_source_t requires audio edit rows") + elif not 0 < self.audio_edit_source_t <= self.audio_t: + raise ValueError(f"MiniMax H3 audio_edit_source_t must be in [1, {self.audio_t}]") ref_rows = _ref_block_rows(self.ref_blocks) layout = [ _ENCODER_WIRE_SCHEMA_ID, @@ -433,14 +637,27 @@ class MiniMaxH3EncoderMediaConditioning: len(self.audio_condition_lengths), len(ref_rows), len(self.keyframe_frame_indices), + video_edit_rows, + audio_edit_rows, + self.audio_edit_source_t, ] layout.extend(item for shape in self.visual_condition_shapes for item in shape) layout.extend(self.audio_condition_lengths) layout.extend(item for row in ref_rows for item in row) layout.extend(self.keyframe_frame_indices) return ( - self.visual_condition if self.visual_condition is not None else empty, - self.audio_condition if self.audio_condition is not None else empty, + _pack_condition_slot( + self.visual_condition, + self.video_edit_clean_rows, + self.video_edit_mask, + width=_VIDEO_CONDITION_WIDTH, + ), + _pack_condition_slot( + self.audio_condition, + self.audio_edit_clean_rows, + self.audio_edit_mask, + width=_AUDIO_CONDITION_WIDTH, + ), torch.tensor(layout, dtype=torch.int64), ) @@ -461,6 +678,11 @@ class MiniMaxH3EncoderConditioning: audio_condition_lengths: tuple[int, ...] = () ref_blocks: tuple[dict[str, Any], ...] = () keyframe_frame_indices: tuple[int, ...] = () + video_edit_clean_rows: torch.Tensor | None = None + video_edit_mask: torch.Tensor | None = None + audio_edit_clean_rows: torch.Tensor | None = None + audio_edit_mask: torch.Tensor | None = None + audio_edit_source_t: int = 0 @classmethod def from_components( @@ -509,11 +731,16 @@ class MiniMaxH3EncoderConditioning: if layout.dtype not in _INTEGER_DTYPES or layout.ndim != 1: raise ValueError("MiniMax H3 encoder layout must be a one-dimensional integer tensor") values = [int(item) for item in layout.detach().cpu().tolist()] - if len(values) < 12: + if len(values) < 2: + raise ValueError("MiniMax H3 encoder layout header is truncated") + schema_id, version = values[:2] + if schema_id != _ENCODER_WIRE_SCHEMA_ID or version != STAGE_SCHEMA_VERSION: + raise ValueError(f"unsupported MiniMax H3 encoder wire schema {schema_id}:{version}") + if len(values) < 15: raise ValueError("MiniMax H3 encoder layout header is truncated") ( - schema_id, - version, + _schema_id, + _version, task_code, height, width, @@ -524,9 +751,10 @@ class MiniMaxH3EncoderConditioning: audio_count, ref_count, keyframe_count, - ) = values[:12] - if schema_id != _ENCODER_WIRE_SCHEMA_ID or version != STAGE_SCHEMA_VERSION: - raise ValueError(f"unsupported MiniMax H3 encoder wire schema {schema_id}:{version}") + video_edit_rows, + audio_edit_rows, + audio_edit_source_t, + ) = values[:15] task = _CODE_TO_TASK.get(task_code) if task is None: raise ValueError(f"unsupported MiniMax H3 wire task code {task_code!r}") @@ -534,10 +762,12 @@ class MiniMaxH3EncoderConditioning: raise ValueError("MiniMax H3 encoder wire dimensions must be positive") if min(visual_count, audio_count, ref_count, keyframe_count) < 0: raise ValueError("MiniMax H3 encoder layout counts must be non-negative") - expected_size = 12 + 3 * visual_count + audio_count + 5 * ref_count + keyframe_count + if min(video_edit_rows, audio_edit_rows, audio_edit_source_t) < 0: + raise ValueError("MiniMax H3 encoder edit layout values must be non-negative") + expected_size = 15 + 3 * visual_count + audio_count + 5 * ref_count + keyframe_count if len(values) != expected_size: raise ValueError(f"MiniMax H3 encoder layout has {len(values)} integers, expected {expected_size}") - cursor = 12 + cursor = 15 visual_shapes = tuple( tuple(values[cursor + 3 * index : cursor + 3 * (index + 1)]) for index in range(visual_count) ) @@ -548,14 +778,55 @@ class MiniMaxH3EncoderConditioning: cursor += 5 * ref_count ref_blocks_tensor = torch.tensor(ref_rows, dtype=torch.int64).reshape(-1, 5) keyframe_frame_indices = tuple(values[cursor : cursor + keyframe_count]) - visual_condition = visual if visual.numel() else None - audio_condition = audio if audio.numel() else None + video_mask_shape = (latent_t, height // 16, width // 16) if video_edit_rows else () + audio_mask_shape = (2, audio_t) if audio_edit_rows else () + visual_condition_rows = sum(t * (h // 2) * (w // 2) for t, h, w in visual_shapes) + audio_condition_rows = _AUDIO_CONDITION_CHANNELS * sum(audio_lengths) + visual_condition, video_edit_clean_rows, video_edit_mask = _unpack_condition_slot( + visual, + condition_rows=visual_condition_rows, + edit_rows=video_edit_rows, + mask_shape=video_mask_shape, + width=_VIDEO_CONDITION_WIDTH, + name="visual", + ) + audio_condition, audio_edit_clean_rows, audio_edit_mask = _unpack_condition_slot( + audio, + condition_rows=audio_condition_rows, + edit_rows=audio_edit_rows, + mask_shape=audio_mask_shape, + width=_AUDIO_CONDITION_WIDTH, + name="audio", + ) _validate_condition_tensors( visual_condition, visual_shapes, audio_condition, audio_lengths, ) + _validate_edit_tensors( + video_edit_clean_rows, + video_edit_mask, + name="video", + width=_VIDEO_CONDITION_WIDTH, + mask_shape=(latent_t, height // 16, width // 16), + ) + _validate_edit_tensors( + audio_edit_clean_rows, + audio_edit_mask, + name="audio", + width=_AUDIO_CONDITION_WIDTH, + mask_shape=(2, audio_t), + ) + expected_video_edit_rows = latent_t * (height // 32) * (width // 32) + if video_edit_rows not in (0, expected_video_edit_rows): + raise ValueError(f"MiniMax H3 video edit has {video_edit_rows} rows, expected {expected_video_edit_rows}") + if audio_edit_rows not in (0, audio_t * _AUDIO_CONDITION_CHANNELS): + raise ValueError(f"MiniMax H3 audio edit has {audio_edit_rows} rows, expected {audio_t * 2}") + if (audio_edit_rows == 0 and audio_edit_source_t != 0) or ( + audio_edit_rows and not 0 < audio_edit_source_t <= audio_t + ): + raise ValueError("MiniMax H3 encoder audio edit source length is inconsistent") return cls( hidden_states=text.hidden_states, token_tags=text.token_tags, @@ -571,6 +842,11 @@ class MiniMaxH3EncoderConditioning: audio_condition_lengths=audio_lengths, ref_blocks=_decode_ref_blocks(ref_blocks_tensor), keyframe_frame_indices=keyframe_frame_indices, + video_edit_clean_rows=video_edit_clean_rows, + video_edit_mask=video_edit_mask, + audio_edit_clean_rows=audio_edit_clean_rows, + audio_edit_mask=audio_edit_mask, + audio_edit_source_t=audio_edit_source_t, ) def to_omni_payload(self) -> dict[str, Any]: @@ -593,6 +869,11 @@ class MiniMaxH3EncoderConditioning: audio_condition_lengths=self.audio_condition_lengths, ref_blocks=self.ref_blocks, keyframe_frame_indices=self.keyframe_frame_indices, + video_edit_clean_rows=self.video_edit_clean_rows, + video_edit_mask=self.video_edit_mask, + audio_edit_clean_rows=self.audio_edit_clean_rows, + audio_edit_mask=self.audio_edit_mask, + audio_edit_source_t=self.audio_edit_source_t, ).to_omni_components() return { "hidden_states": {"output": text.hidden_states}, 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 2d6998b7a..eb6361c13 100644 --- a/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py +++ b/vllm_omni/model_executor/models/minimax_h3/encoder_processing.py @@ -6,6 +6,7 @@ from __future__ import annotations import math import os +import subprocess import tempfile from collections.abc import Callable, Mapping, Sequence from contextlib import AbstractContextManager, nullcontext @@ -33,6 +34,7 @@ from vllm_omni.model_executor.models.minimax_h3.reference_video import ( load_audio_file, load_video_audio, load_video_frames, + prepare_edit_video, prepare_reference_videos, sample_reference_video_frames, validate_reference_audio_files, @@ -263,6 +265,101 @@ def _frames_to_tensor(frames: Sequence[Any]) -> torch.Tensor: ) +def _contains_bool(value: Any) -> bool: + if isinstance(value, bool): + return True + if isinstance(value, torch.Tensor): + return value.dtype == torch.bool + return isinstance(value, (list, tuple)) and any(_contains_bool(item) for item in value) + + +def _edit_mask(value: Any, *, name: str) -> torch.Tensor: + if _contains_bool(value): + raise OmniClientError(f"MiniMax H3 {name} must not contain booleans") + try: + mask = torch.as_tensor(value, dtype=torch.float32) + except (TypeError, ValueError, RuntimeError, OverflowError) as exc: + raise OmniClientError(f"MiniMax H3 {name} must be a finite numeric tensor") from exc + if mask.numel() == 0 or not bool(torch.isfinite(mask).all().item()): + raise OmniClientError(f"MiniMax H3 {name} must be non-empty and finite") + if bool(((mask < 0.0) | (mask > 1.0)).any().item()): + raise OmniClientError(f"MiniMax H3 {name} values must be in [0, 1]") + return mask + + +def _canonical_video_edit_mask( + value: Any, + *, + latent_t: int, + latent_h: int, + latent_w: int, +) -> torch.Tensor: + """Normalize every accepted request shape 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) + row_count = math.prod(token_shape) + candidate = mask + while True: + if candidate.ndim == 0: + return candidate.expand(full_shape).contiguous() + if candidate.ndim == 1 and candidate.numel() == row_count: + candidate = candidate.reshape(token_shape) + if tuple(candidate.shape) == token_shape: + return candidate.repeat_interleave(2, dim=1).repeat_interleave(2, dim=2).contiguous() + if tuple(candidate.shape) == full_shape: + return candidate.contiguous() + if not candidate.ndim or candidate.shape[0] != 1: + break + candidate = candidate.squeeze(0) + raise OmniClientError( + "MiniMax H3 video_noise_mask shape must be " + f"scalar, ({row_count},), {token_shape}, or {full_shape}; got {tuple(mask.shape)}" + ) + + +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") + row_count = 2 * audio_t + candidate = mask + while True: + if candidate.ndim == 0: + return candidate.expand(2, audio_t).contiguous() + if candidate.ndim == 1 and candidate.numel() == audio_t: + return candidate.expand(2, audio_t).contiguous() + if candidate.ndim == 1 and candidate.numel() == row_count: + return candidate.reshape(2, audio_t).contiguous() + if tuple(candidate.shape) == (2, audio_t): + return candidate.contiguous() + if not candidate.ndim or candidate.shape[0] != 1: + break + candidate = candidate.squeeze(0) + raise OmniClientError( + "MiniMax H3 audio_noise_mask shape must be " + f"scalar, ({audio_t},), (2, {audio_t}), or ({row_count},); got {tuple(mask.shape)}" + ) + + +def _fit_audio_edit_rows( + rows: torch.Tensor, + source_audio_t: int, + *, + target_audio_t: int, +) -> tuple[torch.Tensor, int]: + """Pad or trim channel-major clean source rows to the generated duration.""" + if rows.ndim != 2 or rows.shape[0] % 2 or rows.shape[1] != 32: + raise ValueError("MiniMax H3 audio edit rows must have shape [2 * audio_t, 32]") + encoded_audio_t = int(rows.shape[0]) // 2 + if not 0 < source_audio_t <= encoded_audio_t: + raise ValueError(f"MiniMax H3 source audio length must be in [1, {encoded_audio_t}]") + copied_t = min(source_audio_t, target_audio_t) + source = rows.reshape(2, encoded_audio_t, 32) + fitted = rows.new_zeros((2, target_audio_t, 32)) + fitted[:, :copied_t] = source[:, :copied_t] + return fitted.reshape(2 * target_audio_t, 32).contiguous(), copied_t + + def _reuse_prepared_reference_videos( prepared: list[dict[str, Any]] | None, *, @@ -356,6 +453,37 @@ def prepare_encoder_inputs( diffusion_sampling_params, raw_images[0] if raw_images else None, ) + source_video = multi_modal_data.get("source_video") + source_audio = multi_modal_data.get("source_audio") + raw_video_edit_mask = multi_modal_data.get("video_noise_mask") + raw_audio_edit_mask = multi_modal_data.get("audio_noise_mask") + if (source_video is not None or source_audio is not None) and ( + raw_video_edit_mask is None and raw_audio_edit_mask is None + ): + raise OmniClientError("MiniMax H3 edit sources require video_noise_mask or audio_noise_mask") + video_edit_mask = ( + _canonical_video_edit_mask( + raw_video_edit_mask, + latent_t=latent_t, + latent_h=height // 16, + latent_w=width // 16, + ) + if raw_video_edit_mask is not None + else None + ) + audio_edit_mask = ( + _canonical_audio_edit_mask(raw_audio_edit_mask, audio_t=audio_t) if raw_audio_edit_mask is not None else None + ) + # An all-generate mask has no clean source dependency and remains a no-op. + if video_edit_mask is not None and bool(torch.all(video_edit_mask == 1.0).item()): + video_edit_mask = None + if audio_edit_mask is not None and bool(torch.all(audio_edit_mask == 1.0).item()): + audio_edit_mask = None + if video_edit_mask is not None and source_video is None: + raise OmniClientError("MiniMax H3 video_noise_mask requires source_video") + if audio_edit_mask is not None and source_audio is None and source_video is None: + raise OmniClientError("MiniMax H3 audio_noise_mask requires source_audio or source_video") + images = _prepare_encoder_images( task, raw_images, @@ -431,6 +559,48 @@ def prepare_encoder_inputs( audio_index += 1 condition_labels.append(("audio", audio_index)) + video_edit: torch.Tensor | None = None + edit_video_metadata: dict[str, Any] | None = None + if video_edit_mask is not None: + with tempfile.TemporaryDirectory(prefix="minimax_h3_edit_") as workdir: + try: + edit_video_metadata = prepare_edit_video( + source_video, + target_width=width, + target_height=height, + target_frame_count=num_frames, + workdir=workdir, + ) + edit_video_frames = load_video_frames(edit_video_metadata["prepared_path"]) + except subprocess.CalledProcessError as exc: + raise OmniClientError("MiniMax H3 could not decode source_video for editing") from exc + video_edit = _frames_to_tensor(edit_video_frames) + + audio_edit: tuple[torch.Tensor, int] | None = None + if audio_edit_mask is not None: + if source_audio is not None: + try: + waveform, sample_rate = _load_audio(source_audio) + except subprocess.CalledProcessError as exc: + raise OmniClientError("MiniMax H3 could not decode source_audio for editing") from exc + else: + if edit_video_metadata is not None and not edit_video_metadata["input_has_audio"]: + raise OmniClientError("MiniMax H3 source_video has no audio track for audio editing") + try: + waveform, sample_rate = load_video_audio( + str(source_video), + duration_seconds=float(num_frames) / MINIMAX_H3_FPS, + ) + except subprocess.CalledProcessError as exc: + raise OmniClientError("MiniMax H3 could not decode source_video audio for editing") from exc + if waveform.ndim not in (1, 2) or int(sample_rate) <= 0: + raise OmniClientError("MiniMax H3 edit audio requires a waveform and positive sample rate") + max_samples = max(1, int(round(float(num_frames) / MINIMAX_H3_FPS * int(sample_rate)))) + waveform = waveform[..., :max_samples].float().contiguous() + if waveform.shape[-1] == 0: + raise OmniClientError("MiniMax H3 edit audio must not be empty") + audio_edit = (waveform, int(sample_rate)) + if raw_audio is not None: validate_reference_audio_files(raw_audio) standalone_audios = _load_audios(raw_audio) if raw_audio is not None else [] @@ -447,6 +617,10 @@ def prepare_encoder_inputs( video_audios=tuple(video_audio_inputs), audios=tuple((waveform.float().contiguous(), int(sample_rate)) for waveform, sample_rate in standalone_audios), keyframe_frame_indices=tuple(keyframe_indices), + video_edit=video_edit, + video_edit_mask=video_edit_mask, + audio_edit=audio_edit, + audio_edit_mask=audio_edit_mask, ) audio_inputs = _effective_audio_inputs( media_input.video_audios, @@ -491,7 +665,8 @@ def encode_media( """ visual_rows: list[torch.Tensor] = [] visual_shapes: list[tuple[int, int, int]] = [] - if media.images or media.videos: + video_edit_clean_rows: torch.Tensor | None = None + if media.images or media.videos or media.video_edit is not None: if video_vae is None: raise RuntimeError("MiniMax H3 video VAE is not resident on this rank") with component_scope(video_vae): @@ -504,19 +679,34 @@ def encode_media( rows, shape = video_vae.encode_video(frames) visual_rows.append(rows) visual_shapes.append(tuple(int(item) for item in shape)) + if media.video_edit is not None: + frames = np.asarray(media.video_edit.detach().cpu().to(torch.uint8).contiguous().numpy()) + rows, shape = video_vae.encode_video(frames) + shape = tuple(int(item) for item in shape) + expected_shape = (media.latent_t, media.height // 16, media.width // 16) + expected_rows = media.latent_t * (media.height // 32) * (media.width // 32) + if shape != expected_shape or tuple(rows.shape) != (expected_rows, 96): + raise ValueError( + "MiniMax H3 video edit encoder shape mismatch: " + f"got latent shape {shape} and rows {tuple(rows.shape)}, " + f"expected {expected_shape} and ({expected_rows}, 96)" + ) + video_edit_clean_rows = rows.to(dtype=torch.float32).contiguous() if not emit_conditioning: return None audio_rows: list[torch.Tensor] = [] audio_lengths: list[int] = [] + audio_edit_clean_rows: torch.Tensor | None = None + audio_edit_source_t = 0 embedded_audio_count = sum(item is not None for item in media.video_audios) audio_inputs = _effective_audio_inputs( media.video_audios, media.audios, max_standalone_seconds=float(media.num_frames) / MINIMAX_H3_FPS, ) - if audio_inputs: + if audio_inputs or media.audio_edit is not None: if audio_vae is None: raise RuntimeError("MiniMax H3 audio WVAE is not resident on the encoder leader") with component_scope(audio_vae): @@ -524,6 +714,15 @@ def encode_media( rows, length = audio_vae.encode_waveform(waveform, sample_rate) audio_rows.append(rows) audio_lengths.append(int(length)) + if media.audio_edit is not None: + waveform, sample_rate = media.audio_edit + rows, source_audio_t = audio_vae.encode_waveform(waveform, sample_rate) + rows = rows.to(dtype=torch.float32).contiguous() + audio_edit_clean_rows, audio_edit_source_t = _fit_audio_edit_rows( + rows, + int(source_audio_t), + target_audio_t=media.audio_t, + ) if audio_lengths: if any(length < 80 or length > 600 for length in audio_lengths): raise ValueError("MiniMax H3 audio references must each be between 2 and 15 seconds") @@ -553,6 +752,22 @@ def encode_media( } for length in audio_lengths[embedded_audio_count:] ) + video_edit_mask = None + if video_edit_clean_rows is not None: + if media.video_edit_mask is None: + raise ValueError("MiniMax H3 video edit rows require a mask") + video_edit_mask = media.video_edit_mask.to( + device=video_edit_clean_rows.device, + dtype=torch.float32, + ).contiguous() + audio_edit_mask = None + if audio_edit_clean_rows is not None: + if media.audio_edit_mask is None: + raise ValueError("MiniMax H3 audio edit rows require a mask") + audio_edit_mask = media.audio_edit_mask.to( + device=audio_edit_clean_rows.device, + dtype=torch.float32, + ).contiguous() return MiniMaxH3EncoderMediaConditioning( task=media.task, height=media.height, @@ -566,4 +781,9 @@ def encode_media( audio_condition_lengths=tuple(audio_lengths), ref_blocks=tuple(ref_blocks), keyframe_frame_indices=media.keyframe_frame_indices, + video_edit_clean_rows=video_edit_clean_rows, + video_edit_mask=video_edit_mask, + audio_edit_clean_rows=audio_edit_clean_rows, + audio_edit_mask=audio_edit_mask, + audio_edit_source_t=audio_edit_source_t, ) diff --git a/vllm_omni/model_executor/models/minimax_h3/reference_video.py b/vllm_omni/model_executor/models/minimax_h3/reference_video.py index 4b361b26c..6531f1265 100644 --- a/vllm_omni/model_executor/models/minimax_h3/reference_video.py +++ b/vllm_omni/model_executor/models/minimax_h3/reference_video.py @@ -324,10 +324,21 @@ def _transcode_reference_video( workdir: str, start_time_seconds: float = 0.0, duration_seconds: float | None = None, + pad_last_frame: bool = False, ) -> str: output = str(Path(workdir) / "prepared.mp4") duration_args = ["-t", f"{float(duration_seconds):.6f}"] if duration_seconds is not None else [] frame_count_args = ["-frames:v", str(int(target_frame_count))] if target_frame_count > 0 else [] + filters = [ + f"fps={MINIMAX_H3_FPS:g}", + f"scale={target_width}:{target_height}:flags=lanczos", + "setsar=1", + ] + if pad_last_frame: + # Keep cloning the final converted frame until -frames:v reaches the + # requested target; longer inputs are deterministically trimmed by the + # same -frames:v limit without enabling this padding filter. + filters.append("tpad=stop_mode=clone:stop=-1") subprocess.run( [ "ffmpeg", @@ -342,7 +353,7 @@ def _transcode_reference_video( "0:v:0", "-an", "-vf", - (f"fps={MINIMAX_H3_FPS:g},scale={target_width}:{target_height}:flags=lanczos,setsar=1"), + ",".join(filters), *duration_args, *frame_count_args, "-metadata:s:v:0", @@ -366,6 +377,41 @@ def _transcode_reference_video( return output +def prepare_edit_video( + value: Any, + target_width: int, + target_height: int, + target_frame_count: int, + workdir: str, +) -> dict[str, Any]: + """Resize an edit source and clone its last frame to the target length.""" + if not isinstance(value, str | os.PathLike): + raise OmniClientError("MiniMax H3 edit video input must be a single file path") + source = str(value) + source_meta = _probe_video(source) + if int(source_meta.get("frame_count", 0)) <= 0: + raise OmniClientError(f"video has no frames: {source}") + + Path(workdir).mkdir(parents=True, exist_ok=True) + prepared_path = _transcode_reference_video( + source, + target_width=target_width, + target_height=target_height, + target_frame_count=target_frame_count, + workdir=workdir, + pad_last_frame=True, + ) + + return { + "original_path": source, + "prepared_path": prepared_path, + "input_has_audio": bool(source_meta.get("audio_codecs")), + "width": target_width, + "height": target_height, + "frame_count": target_frame_count, + } + + def prepare_reference_videos( values: Any, *, @@ -674,6 +720,7 @@ __all__ = [ "load_audio_file", "load_video_audio", "load_video_frames", + "prepare_edit_video", "prepare_reference_videos", "sample_reference_video_frames", "validate_reference_audio_files",