[Model][Frontend] MiniMax-H3: Add latent-mask editing (#7465)

Signed-off-by: xiaoyu <1259084489@qq.com>

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):

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)

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)

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()