[4/N] Unify worker-side MP4 encoding and wire Wan into it (#7048)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import io
+from types import SimpleNamespace
+
+import av
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import (
+    WAN_DEFAULT_OUTPUT_FPS,
+    resolve_wan_output_fps,
+    resolve_wan_preencode_batch_frames,
+    resolve_wan_preencode_mp4,
+    resolve_wan_video_codec_options,
+    wan_preencoded_mp4_payload,
+)
+from vllm_omni.diffusion.utils.chunked_video import decode_to_mp4
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeWanVAE:
+    """Publish temporal chunks the way the Wan VAE seam does."""
+
+    chunk_value_range = (-1.0, 1.0)
+
+    def __init__(self, *, batch: int, chunks: int, frames_per_chunk: int, height: int, width: int):
+        self.batch = batch
+        self.chunks = chunks
+        self.frames_per_chunk = frames_per_chunk
+        self.height = height
+        self.width = width
+        self.chunk_calls = 0
+
+    def decode_with_chunks(self, latents, *, on_chunk):
+        del latents
+        assert on_chunk is not None
+        for index in range(self.chunks):
+            self.chunk_calls += 1
+            value = -1.0 + 2.0 * index / max(self.chunks - 1, 1)
+            on_chunk(torch.full((self.batch, 3, self.frames_per_chunk, self.height, self.width), value))
+
+
+def _decoded_frame_count(data: bytes) -> int:
+    with av.open(io.BytesIO(data)) as container:
+        return sum(1 for _ in container.decode(video=0))
+
+
+def test_decode_wan_latents_to_mp4_emits_one_playable_mp4_per_batch_entry():
+    vae = _FakeWanVAE(batch=2, chunks=6, frames_per_chunk=2, height=16, width=32)
+
+    videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=4)
+
+    assert len(videos) == 2
+    for data in videos:
+        assert data[:4] and _decoded_frame_count(data) == 12
+
+
+def test_decode_wan_latents_to_mp4_coalesces_transfers_without_losing_frames():
+    """batch_frames only changes how many transfers carry the same frames."""
+    unbatched = _FakeWanVAE(batch=1, chunks=8, frames_per_chunk=1, height=16, width=32)
+    batched = _FakeWanVAE(batch=1, chunks=8, frames_per_chunk=1, height=16, width=32)
+
+    per_chunk = decode_to_mp4(unbatched, torch.zeros(1), fps=24, batch_frames=1)
+    coalesced = decode_to_mp4(batched, torch.zeros(1), fps=24, batch_frames=4)
+
+    assert _decoded_frame_count(per_chunk[0]) == _decoded_frame_count(coalesced[0]) == 8
+
+
+def test_decode_wan_latents_to_mp4_returns_nothing_for_a_rank_without_output():
+    class SilentVAE:
+        chunk_value_range = (-1.0, 1.0)
+
+        def decode_with_chunks(self, latents, *, on_chunk):
+            del latents, on_chunk
+
+    assert decode_to_mp4(SilentVAE(), torch.zeros(1), fps=24) == []
+
+
+def test_decode_wan_latents_to_mp4_rejects_a_non_positive_batch():
+    with pytest.raises(ValueError, match="batch_frames"):
+        decode_to_mp4(
+            _FakeWanVAE(batch=1, chunks=1, frames_per_chunk=1, height=16, width=16), None, fps=24, batch_frames=0
+        )
+
+
+def test_decode_wan_latents_to_mp4_rejects_a_vae_without_the_capability():
+    class PlainVAE:
+        chunk_value_range = (-1.0, 1.0)
+
+        def decode(self, z, return_dict=True):
+            del z, return_dict
+
+    with pytest.raises(TypeError, match="chunked VAE decode capability"):
+        decode_to_mp4(PlainVAE(), torch.zeros(1), fps=24)
+
+
+def test_preencode_flag_is_off_unless_the_request_asks_for_it():
+    assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args=None), output_type="np") is False
+    assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args={}), output_type="np") is False
+    assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args={"preencode_mp4": True}), output_type="np") is True
+
+
+def test_preencode_rejects_combinations_it_cannot_serve():
+    interpolating = SimpleNamespace(extra_args={"preencode_mp4": True}, enable_frame_interpolation=True)
+    with pytest.raises(ValueError, match="enable_frame_interpolation"):
+        resolve_wan_preencode_mp4(interpolating, output_type="np")
+
+    with pytest.raises(ValueError, match="output_type"):
+        resolve_wan_preencode_mp4(SimpleNamespace(extra_args={"preencode_mp4": True}), output_type="pil")
+
+
+def test_output_fps_mirrors_the_serving_fallback():
+    assert resolve_wan_output_fps(SimpleNamespace(fps=16)) == 16
+    assert resolve_wan_output_fps(SimpleNamespace(fps=[16])) == 16
+    assert resolve_wan_output_fps(SimpleNamespace(fps=None)) == WAN_DEFAULT_OUTPUT_FPS
+    assert resolve_wan_output_fps(SimpleNamespace()) == WAN_DEFAULT_OUTPUT_FPS
+
+
+def test_codec_options_reach_the_encoder_as_strings():
+    params = SimpleNamespace(extra_args={"video_codec_options": {"preset": "ultrafast", "threads": 0}})
+    assert resolve_wan_video_codec_options(params) == {"preset": "ultrafast", "threads": "0"}
+    assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={})) == {"preset": "ultrafast", "threads": "0"}
+    assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={"video_codec_options": None})) is None
+    assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={"video_codec_options": {}})) == {}
+
+
+def test_preencoded_payload_passes_bytes_through_and_ignores_tensors():
+    assert wan_preencoded_mp4_payload(torch.zeros(1, 3, 2, 4, 4)) is None
+    assert wan_preencoded_mp4_payload([b"one", b"two"]) == {"payload": {"video": [b"one", b"two"]}, "metadata": {}}
+    assert wan_preencoded_mp4_payload(b"one") == {"payload": {"video": [b"one"]}, "metadata": {}}
+
+
+@pytest.mark.parametrize("value", [1, 5, 33])
+def test_request_batch_threshold_controls_transfers_without_losing_frames(monkeypatch, value):
+    from vllm_omni.diffusion.utils import chunked_video
+
+    transfers = []
+    quantize = chunked_video.chunk_to_uint8_frames
+
+    def capture(chunk, value_range):
+        transfers.append(chunk.shape[2])
+        return quantize(chunk, value_range)
+
+    monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
+    params = SimpleNamespace(extra_args={"preencode_mp4": True, "preencode_batch_frames": value})
+    vae = _FakeWanVAE(batch=1, chunks=4, frames_per_chunk=3, height=16, width=32)
+    videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=resolve_wan_preencode_batch_frames(params))
+
+    assert transfers == {1: [3, 3, 3, 3], 5: [6, 6], 33: [12]}[value]
+    assert _decoded_frame_count(videos[0]) == 12
+
+
+@pytest.mark.parametrize("value", [0, -1, True, 1.5, "17", None])
+def test_request_batch_threshold_rejects_invalid_values(value):
+    with pytest.raises(ValueError, match="preencode_batch_frames"):
+        resolve_wan_preencode_batch_frames(SimpleNamespace(extra_args={"preencode_batch_frames": value}))
+
+
+def test_request_batch_threshold_preserves_model_defaults():
+    params = SimpleNamespace(extra_args=None)
+    assert resolve_wan_preencode_batch_frames(params) == 17
+    assert resolve_wan_preencode_batch_frames(params, default=1) == 1

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -647,3 +647,245 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
     assert video.shape[1] == 3  # channels
     assert audio_waveform is not None
     assert audio_sr == 16000
+
+
+def _make_s2v_preencode_pipeline() -> Wan22S2VPipeline:
+    """Build the same stub pipeline the batching test drives, for preencode runs."""
+    pipeline = object.__new__(Wan22S2VPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = MagicMock()
+    pipeline.transformer.dtype = torch.float32
+    # A fresh iterator per call: the clip loop reads this once per clip.
+    pipeline.transformer.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float32)])
+    pipeline.transformer.casual_audio_encoder = None
+    pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
+    pipeline.vae = MagicMock()
+    pipeline.vae.dtype = torch.float32
+    # The shared consumer reads the published pixel range off the VAE.
+    pipeline.vae.chunk_value_range = (-1.0, 1.0)
+    # One decoded clip: [B, C, T, H, W] for a batch of four (2 requests x 2 outputs).
+    pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
+    pipeline.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
+    pipeline._flow_shift = 3.0
+    pipeline.vae_scale_factor_spatial = 8
+    pipeline.resolution_divisor = 16
+    pipeline.motion_frames = 7
+    pipeline.drop_first_motion = True
+    pipeline._DEFAULT_INFER_FRAMES = 8
+    pipeline._guidance_scale = None
+    pipeline._num_timesteps = None
+    pipeline.check_inputs = lambda *args, **kwargs: None
+    pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
+    pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
+    pipeline.prepare_motion_latents = MagicMock(
+        side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
+    )
+    pipeline._denormalize_latents = lambda latents: latents
+    pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])
+    return pipeline
+
+
+def _make_s2v_preencode_batch(audio_a, audio_b, **sampling_overrides) -> DiffusionRequestBatch:
+    image = PIL.Image.new("RGB", (16, 16))
+    return DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="a",
+                prompt={"prompt": "first", "multi_modal_data": {"image": image, "audio": audio_a}},
+                sampling_params=_make_s2v_sampling(
+                    num_outputs_per_prompt=2,
+                    latents=torch.zeros(2, 16, 2, 2, 2),
+                    output_type="np",
+                    **sampling_overrides,
+                ),
+            ),
+            SimpleNamespace(
+                request_id="b",
+                prompt={"prompt": "second", "multi_modal_data": {"image": image, "audio": audio_b}},
+                sampling_params=_make_s2v_sampling(
+                    num_outputs_per_prompt=2,
+                    latents=torch.ones(2, 16, 2, 2, 2),
+                    output_type="np",
+                    **sampling_overrides,
+                ),
+            ),
+        ]
+    )
+
+
+def _decode_mp4(data: bytes):
+    """Return (video frame count, first audio samples) for one encoded container."""
+    import io
+
+    import av
+
+    with av.open(io.BytesIO(data)) as container:
+        frames = sum(1 for _ in container.decode(video=0))
+    with av.open(io.BytesIO(data)) as container:
+        assert container.streams.audio, "the worker must mux the waveform into the container"
+        samples = np.concatenate(
+            [frame.to_ndarray().reshape(-1) for frame in container.decode(audio=0)][:4],
+        )
+    return frames, samples
+
+
+@pytest.mark.parametrize("batch_frames", [1, 1000])
+def test_s2v_preencode_returns_playable_mp4_bytes_per_request(monkeypatch, batch_frames) -> None:
+    """The clip loop hands finished clips to the encoder instead of concatenating."""
+    pipeline = _make_s2v_preencode_pipeline()
+    # Two clips, so the autoregressive motion feedback runs between pushes.
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
+    audio_a = np.zeros(16000, dtype=np.float32)
+    audio_b = np.full(16000, 0.5, dtype=np.float32)
+    batch = _make_s2v_preencode_batch(
+        audio_a, audio_b, extra_args={"preencode_mp4": True, "preencode_batch_frames": batch_frames}
+    )
+    from vllm_omni.diffusion.utils import chunked_video
+
+    transfers = []
+    quantize = chunked_video.chunk_to_uint8_frames
+
+    def capture(chunk, value_range):
+        transfers.append(chunk.shape[2])
+        return quantize(chunk, value_range)
+
+    monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        outputs = pipeline.forward(batch)
+
+    assert len(outputs) == 2
+    assert pipeline.vae.decode.call_count == 2, "both clips must reach the encoder"
+    assert len(transfers) == (2 if batch_frames == 1 else 1)
+
+    decoded = []
+    for request_output in outputs:
+        # Two outputs per prompt, each already a complete container.
+        assert isinstance(request_output.output, list)
+        assert len(request_output.output) == 2
+        for data in request_output.output:
+            assert isinstance(data, bytes)
+            frames, samples = _decode_mp4(data)
+            assert frames > 0
+            decoded.append(samples)
+
+    # Four entries: request a's waveform for its two outputs, then request b's.
+    # A silent and a non-silent waveform must not land on the same request.
+    assert np.abs(decoded[0]).max() == pytest.approx(0.0, abs=1e-3)
+    assert np.abs(decoded[1]).max() == pytest.approx(0.0, abs=1e-3)
+    assert np.abs(decoded[2]).max() > 0.1
+    assert np.abs(decoded[3]).max() > 0.1
+
+
+def test_s2v_preencode_aborts_encoders_when_a_later_clip_fails(monkeypatch) -> None:
+    """A failure after the first push must not strand encoder worker threads."""
+    pipeline = _make_s2v_preencode_pipeline()
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
+    batch = _make_s2v_preencode_batch(
+        np.zeros(16000, dtype=np.float32),
+        np.ones(16000, dtype=np.float32),
+        extra_args={"preencode_mp4": True, "preencode_batch_frames": 1},
+    )
+    sessions = []
+
+    from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
+
+    def capture_session(**kwargs):
+        session = ChunkedVideoMP4Session(**kwargs)
+        sessions.append(session)
+        return session
+
+    def fail_on_second_clip(**kwargs):
+        if pipeline.vae.decode.call_count:
+            raise RuntimeError("second clip failed")
+        return kwargs["latents"]
+
+    pipeline.diffuse = fail_on_second_clip
+    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.ChunkedVideoMP4Session", capture_session)
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        with pytest.raises(RuntimeError, match="second clip failed"):
+            pipeline.forward(batch)
+
+    assert len(sessions) == 1
+    assert len(sessions[0]._encoders) == 4
+    assert all(not encoder._thread.is_alive() for encoder in sessions[0]._encoders)
+
+
+def test_s2v_preencode_skips_encoding_on_a_vae_patch_parallel_peer(monkeypatch) -> None:
+    """A peer still receives every clip for the motion loop but encodes none of them."""
+    pipeline = _make_s2v_preencode_pipeline()
+    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
+    # A peer's patch-parallel decode returns an empty placeholder.
+    pipeline.vae.decode.return_value = (torch.empty(0),)
+    pipeline.vae._vae_pp_group = object()
+    broadcasts = []
+    monkeypatch.setattr(
+        "torch.distributed.broadcast",
+        lambda tensor, src, group: broadcasts.append(tensor.zero_()),
+    )
+    batch = _make_s2v_preencode_batch(
+        np.zeros(16000, dtype=np.float32),
+        np.ones(16000, dtype=np.float32),
+        extra_args={"preencode_mp4": True, "preencode_batch_frames": 1},
+    )
+    sessions = []
+
+    from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
+
+    def capture_session(**kwargs):
+        session = ChunkedVideoMP4Session(**kwargs)
+        sessions.append(session)
+        return session
+
+    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.ChunkedVideoMP4Session", capture_session)
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        outputs = pipeline.forward(batch)
+
+    assert len(broadcasts) == 2, "the peer must still join the broadcast for both clips"
+    # Once for the initial motion latents, once more from the peer's broadcast clip.
+    assert pipeline.prepare_motion_latents.call_count == 2, "the motion loop still consumes the peer's frames"
+    assert sessions[0]._encoders == []
+    assert [request_output.output for request_output in outputs] == [[], []]
+
+
+@pytest.mark.parametrize("output_type", ["pil", "pt", "latent"])
+def test_s2v_preencode_rejects_request_output_types_it_cannot_serve(output_type) -> None:
+    """Pre-encoding returns MP4 bytes, so a request asking for frames must fail."""
+    pipeline = _make_s2v_preencode_pipeline()
+    batch = _make_s2v_preencode_batch(
+        np.zeros(16000, dtype=np.float32),
+        np.ones(16000, dtype=np.float32),
+        extra_args={"preencode_mp4": True},
+    )
+    for request in batch.requests:
+        request.sampling_params.output_type = output_type
+
+    with pytest.raises(ValueError, match="output_type"):
+        pipeline.forward(batch)
+
+
+def test_s2v_preencode_keeps_the_full_decode_path_untouched() -> None:
+    """Without the flag the loop still returns the (video, audio, rate) tuple."""
+    pipeline = _make_s2v_preencode_pipeline()
+    audio_a = np.zeros(16000, dtype=np.float32)
+    audio_b = np.ones(16000, dtype=np.float32)
+    batch = _make_s2v_preencode_batch(audio_a, audio_b)
+
+    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+        platform.is_available.return_value = False
+        outputs = pipeline.forward(batch)
+
+    assert outputs[0].output[0].shape[0] == 2
+    np.testing.assert_array_equal(outputs[0].output[1], audio_a)
+    np.testing.assert_array_equal(outputs[1].output[1], audio_b)