[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>

docs/serving/videos_api.md

@@ -251,6 +251,7 @@ These `extra_params` control how the server turns decoded frames into MP4 bytes.
 | Field | Type | Default | Description |
 | --- | --- | --- | --- |
 | `preencode_mp4` | boolean | false | Encode the MP4 on the worker while the VAE is still decoding, instead of after the full video is materialized |
+| `preencode_batch_frames` | positive integer | 17 (H3, Wan T2V/I2V); 1 (Wan S2V) | Minimum accumulated frames per worker transfer/encoding batch; used only with `preencode_mp4=true` |
 | `video_codec_options` | object | null | Encoder options passed through to the H.264 encoder, such as `{"preset": "ultrafast", "threads": "0"}` |
 
 With `preencode_mp4` enabled, each committed VAE chunk leaves the accelerator and
@@ -261,16 +262,29 @@ unchanged: the same complete MP4, byte-for-byte equivalent frames.
 ```bash
 curl -X POST http://localhost:8091/v1/videos/sync \
   -F "prompt=A small robot walking through a neon city" \
-  -F 'extra_params={"preencode_mp4": true, "video_codec_options": {"preset": "ultrafast"}}' \
+  -F 'extra_params={"preencode_mp4": true, "preencode_batch_frames": 33, "video_codec_options": {"preset": "ultrafast"}}' \
   -o output.mp4
 ```
 
+Set `preencode_batch_frames` in `extra_params` (or `extra_args` for offline
+sampling) to tune batching. The worker accumulates complete VAE chunks until
+it has at least this many frames, then transfers and encodes them together.
+It always flushes the final partial batch. This is a threshold, not an exact
+chunk length: a value of 1 submits every native chunk immediately, and a value
+smaller than a native chunk does not split it. Larger values reduce transfers
+but retain more frames on the accelerator and delay encoding. The VAE decode
+window and output frame count stay unchanged. Wan S2V keeps its existing
+per-clip behavior by default. Zero, negative, fractional, boolean, string, and
+null values are rejected when pre-encoding is enabled.
+
 `preencode_mp4` applies to the complete-MP4 response paths only. The
 `/v1/realtime/video` WebSocket endpoint rejects it, because that path already
-overlaps encoding through its own incremental fragmented-MP4 encoder.
+overlaps encoding through its own incremental fragmented-MP4 encoder. Wan also
+rejects it together with `enable_frame_interpolation`, which needs the decoded
+frames the pre-encoded path no longer materializes.
 
-Support is per model: MiniMax-H3 implements it, and other models ignore the flag
-and take the full-decode path.
+Support is per model: MiniMax-H3 and Wan 2.2 (T2V, I2V, and S2V) implement it,
+and other models ignore the flag and take the full-decode path.
 
 ## Storage
 

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -2852,7 +2852,7 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
 
     class FakeEncoder:
-        instances = []
+        instances: list["FakeEncoder"] = []
 
         def __init__(self, **kwargs):
             self.pushes = []
@@ -2873,6 +2873,8 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
             return torch.zeros(1, 1, 2)
 
     class FakeVideoVAE:
+        chunk_value_range = (0.0, 1.0)
+
         def decode_with_chunks(self, latent, *, on_chunk):
             for value in (0.0, 0.25, 0.5):
                 on_chunk(torch.full((1, 3, 3, 2, 2), value))
@@ -2884,7 +2886,7 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
     pipeline.device = torch.device("cpu")
     monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
     monkeypatch.setattr(
-        "vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder",
+        "vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder",
         FakeEncoder,
     )
 
@@ -2904,12 +2906,81 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
     assert encoder.pushes[1].shape == (3, 2, 2, 3)
 
 
+def _peer_rank_preencode_pipeline(monkeypatch):
+    """An H3 pipeline whose video VAE is a distributed rank that owns no output.
+
+    Every rank of the VAE group must call ``decode_with_chunks`` to keep the
+    temporal collectives in lockstep, but only the owner receives chunks, so a
+    peer's callback is never invoked.
+    """
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+
+    class FakeAudioVAE:
+        def decode_latent(self, latent):
+            return torch.zeros(1, 1, 2)
+
+    class PeerRankVideoVAE:
+        chunk_value_range = (0.0, 1.0)
+
+        def decode_with_chunks(self, latent, *, on_chunk):
+            del latent, on_chunk
+
+    pipeline = object.__new__(mod.MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.audio_vae = FakeAudioVAE()
+    pipeline.video_vae = PeerRankVideoVAE()
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config = SimpleNamespace()
+    monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
+    return pipeline
+
+
+def test_peer_vae_rank_returns_no_preencoded_output_instead_of_failing(monkeypatch):
+    """A rank that owns no output must not turn its silence into a request error."""
+    pipeline = _peer_rank_preencode_pipeline(monkeypatch)
+
+    output = pipeline.decode_to_mp4(torch.zeros(1), torch.zeros(1), height=2, width=2)
+
+    assert output == b""
+
+
+def test_peer_vae_rank_post_decode_reaches_post_processing(monkeypatch):
+    """The real H3 post_decode chain, not just the shared consumer, tolerates a peer rank."""
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+
+    pipeline = _peer_rank_preencode_pipeline(monkeypatch)
+    monkeypatch.setattr(pipeline, "_unpack_denoised_rows", lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)))
+    state = StepRequestState(
+        request_id="peer-rank",
+        sampling=SimpleNamespace(num_outputs_per_prompt=1),
+        prompt="a prompt",
+    )
+    state.latents = torch.zeros(1)
+    state.extra[mod._STEP_BRANCH] = object()
+    state.extra[mod._STEP_AUDIO_ROWS] = torch.zeros(1)
+    state.extra[mod._STEP_SHAPE] = {
+        "latent_t": 1,
+        "latent_h": 1,
+        "latent_w": 1,
+        "audio_t": 1,
+        "height": 2,
+        "width": 2,
+        "preencode_mp4": True,
+    }
+
+    output = pipeline.post_decode(state)
+
+    assert output.output == (b"", None)
+    assert mod._minimax_h3_post_process(output.output)["video"] == [b""]
+
+
 def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch):
     """A client's encoder options must survive the worker-side pre-encode path."""
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
 
     class FakeEncoder:
-        instances = []
+        instances: list["FakeEncoder"] = []
 
         def __init__(self, **kwargs):
             self.kwargs = kwargs
@@ -2929,6 +3000,8 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc
             return torch.zeros(1, 1, 2)
 
     class FakeVideoVAE:
+        chunk_value_range = (0.0, 1.0)
+
         def decode_with_chunks(self, latent, *, on_chunk):
             on_chunk(torch.zeros(1, 3, 1, 2, 2))
 
@@ -2938,7 +3011,7 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc
     pipeline.video_vae = FakeVideoVAE()
     pipeline.device = torch.device("cpu")
     monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
-    monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder)
+    monkeypatch.setattr("vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder", FakeEncoder)
 
     pipeline.decode_to_mp4(
         torch.zeros(1),

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)

tests/diffusion/utils/test_chunked_video.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import io
+
+import av
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.utils.chunked_video import (
+    ChunkedVideoMP4Session,
+    chunk_to_uint8_frames,
+    decode_to_mp4,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeVAE:
+    """Publish committed chunks the way a chunked-decode VAE does."""
+
+    def __init__(self, value_range, *, batch=1, chunks=4, frames=2, height=16, width=32):
+        self.chunk_value_range = value_range
+        self._shape = (batch, 3, frames, height, width)
+        self._chunks = chunks
+
+    def decode_with_chunks(self, z, *, on_chunk):
+        del z
+        low, high = self.chunk_value_range
+        for index in range(self._chunks):
+            value = low + (high - low) * index / max(self._chunks - 1, 1)
+            on_chunk(torch.full(self._shape, value))
+
+
+def _decoded(data: bytes):
+    with av.open(io.BytesIO(data)) as container:
+        return np.stack([f.to_ndarray(format="rgb24") for f in container.decode(video=0)])
+
+
+def test_quantization_maps_each_producer_range_onto_the_same_pixels():
+    """A model's published range is what makes chunks comparable across models."""
+    wan = torch.tensor([-1.0, 0.0, 1.0]).view(1, 1, 3, 1, 1).expand(1, 3, 3, 1, 1)
+    h3 = torch.tensor([0.0, 0.5, 1.0]).view(1, 1, 3, 1, 1).expand(1, 3, 3, 1, 1)
+
+    from_wan = chunk_to_uint8_frames(wan.contiguous(), (-1.0, 1.0))
+    from_h3 = chunk_to_uint8_frames(h3.contiguous(), (0.0, 1.0))
+
+    # Same physical pixels, different published ranges.
+    assert np.array_equal(from_wan, from_h3)
+    assert from_wan.reshape(-1, 3)[:, 0].tolist() == [0, 128, 255]
+
+
+def test_quantization_rejects_a_degenerate_range():
+    with pytest.raises(ValueError, match="increasing"):
+        chunk_to_uint8_frames(torch.zeros(1, 3, 1, 2, 2), (1.0, 1.0))
+
+
+@pytest.mark.parametrize("value_range", [(-1.0, 1.0), (0.0, 1.0)])
+def test_decode_to_mp4_drives_any_declared_range(value_range):
+    videos = decode_to_mp4(_FakeVAE(value_range), torch.zeros(1), fps=24, batch_frames=4)
+
+    assert len(videos) == 1
+    assert _decoded(videos[0]).shape[0] == 8
+
+
+def test_decode_to_mp4_emits_one_container_per_batch_entry():
+    videos = decode_to_mp4(_FakeVAE((-1.0, 1.0), batch=3), torch.zeros(1), fps=24)
+
+    assert len(videos) == 3
+    assert all(_decoded(v).shape[0] == 8 for v in videos)
+
+
+def test_decode_to_mp4_rejects_a_vae_without_the_capability():
+    class PlainVAE:
+        def decode(self, z):
+            del z
+
+    with pytest.raises(TypeError, match="chunked VAE decode capability"):
+        decode_to_mp4(PlainVAE(), torch.zeros(1), fps=24)
+
+
+def test_decode_to_mp4_returns_nothing_for_a_rank_without_output():
+    class SilentVAE:
+        chunk_value_range = (-1.0, 1.0)
+
+        def decode_with_chunks(self, z, *, on_chunk):
+            del z, on_chunk
+
+    assert decode_to_mp4(SilentVAE(), torch.zeros(1), fps=24) == []
+
+
+def test_session_crops_the_decoder_padding_away():
+    session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, crop=(8, 16))
+    session.push(torch.zeros(1, 3, 4, 16, 32))
+
+    assert _decoded(session.finish()[0]).shape[1:3] == (8, 16)
+
+
+def test_session_muxes_one_waveform_per_batch_entry():
+    audio = [np.zeros(8000, dtype=np.float32), np.full(8000, 0.5, dtype=np.float32)]
+    session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, audio_waveforms=audio, audio_sample_rate=16000)
+    session.push(torch.zeros(2, 3, 4, 16, 16))
+
+    for data in session.finish():
+        with av.open(io.BytesIO(data)) as container:
+            assert container.streams.audio
+
+
+def test_session_rejects_a_waveform_count_that_does_not_match_the_batch():
+    session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, audio_waveforms=[np.zeros(8000, dtype=np.float32)])
+    with pytest.raises(ValueError, match="one audio waveform per batch entry"):
+        session.push(torch.zeros(2, 3, 4, 16, 16))
+
+
+def test_session_rejects_a_non_positive_batch():
+    with pytest.raises(ValueError, match="batch_frames"):
+        ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, batch_frames=0)

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -1726,66 +1726,47 @@ class MiniMaxH3Pipeline(
     ) -> bytes:
         """Decode and encode one output on the worker without full-video materialization.
 
-        Audio is decoded first so the incremental mux session can attach its audio
-        stream before temporal video chunks arrive. The callback receives committed
-        float32 ``BCTHW`` frames, performs the requested-size crop and uint8
-        conversion in the worker, then applies bounded backpressure to the encoder.
-        """
-        from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder
+        Audio is decoded first so the incremental mux session can attach its
+        audio stream before temporal video chunks arrive. Everything after a
+        chunk is committed -- crop, quantization, transfer, encoding -- is the
+        shared consumer's job; this method only supplies what is specific to
+        H3: the audio waveform, the requested-size crop, and the fixed rate.
 
-        if batch_frames <= 0:
-            raise ValueError("batch_frames must be positive")
+        Every rank of a distributed VAE group drives the temporal collectives,
+        but only the output owner receives chunks, so a peer rank returns empty
+        bytes -- the pre-encoded counterpart of the empty tensor the full
+        decode leaves there.
+        """
+        from vllm_omni.diffusion.utils.chunked_video import decode_to_mp4 as decode_chunks_to_mp4
 
         with self._component_on_device(self.audio_vae):
             audio = self.audio_vae.decode_latent(audio_latent)
         audio_np = audio.detach().float().cpu().numpy()
         if audio_np.ndim == 3 and audio_np.shape[0] == 1:
             audio_np = audio_np[0]
-        encoder = ChunkedMP4Encoder(
-            width=width,
-            height=height,
-            fps=MINIMAX_H3_FPS,
-            audio_waveform=audio_np,
-            audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE,
-            max_pending=max_pending,
-            video_codec_options=video_codec_options,
-        )
 
-        pending_chunks: list[torch.Tensor] = []
-        pending_frames = 0
-
-        def flush_pending() -> None:
-            nonlocal pending_frames
-            if not pending_chunks:
-                return
-            batched = torch.cat(pending_chunks, dim=1)
-            encoder.push(batched[0].cpu().numpy())
-            pending_chunks.clear()
-            pending_frames = 0
-
-        def on_chunk(frames: torch.Tensor) -> None:
-            nonlocal pending_frames
-            prepared = _prepare_minimax_h3_video_output(frames[..., :height, :width])
-            if prepared.shape[0] != 1:
-                raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder")
-            pending_chunks.append(prepared)
-            pending_frames += int(prepared.shape[1])
-            if pending_frames >= batch_frames:
-                flush_pending()
-
-        try:
-            with self._component_on_device(self.video_vae):
-                with current_omni_platform.create_autocast_context(
-                    device_type=self.device.type,
-                    dtype=torch.float16,
-                    enabled=True,
-                ):
-                    self.video_vae.decode_with_chunks(video_latent, on_chunk=on_chunk)
-            flush_pending()
-            return encoder.finish()
-        except BaseException:
-            encoder.abort()
-            raise
+        with self._component_on_device(self.video_vae):
+            with current_omni_platform.create_autocast_context(
+                device_type=self.device.type,
+                dtype=torch.float16,
+                enabled=True,
+            ):
+                videos = decode_chunks_to_mp4(
+                    self.video_vae,
+                    video_latent,
+                    fps=MINIMAX_H3_FPS,
+                    audio_waveforms=[audio_np],
+                    audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE,
+                    batch_frames=batch_frames,
+                    max_pending=max_pending,
+                    video_codec_options=video_codec_options,
+                    crop=(height, width),
+                )
+        if not videos:
+            return b""
+        if len(videos) != 1:
+            raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder")
+        return videos[0]
 
     def decode(
         self,

vllm_omni/diffusion/models/wan2_2/chunked_mp4.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Worker-side MP4 encoding for Wan VAE temporal chunks."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm_omni.diffusion.utils.media_utils import normalize_preencode_batch_frames, normalize_video_codec_options
+
+# Matches the fallback the video serving layer applies when a request omits fps,
+# so the bytes the worker encodes carry the rate the response advertises.
+WAN_DEFAULT_OUTPUT_FPS = 24
+
+# Keep worker-side encoding consistent with the one-shot serving encoder when
+# a request does not choose codec options explicitly.
+WAN_DEFAULT_VIDEO_CODEC_OPTIONS = {"preset": "ultrafast", "threads": "0"}
+
+# Wan publishes one latent frame group per callback, which is far finer than one
+# host transfer is worth. 17 frames matches MiniMax-H3's native clip size.
+WAN_DEFAULT_BATCH_FRAMES = 17
+
+
+def wan_preencoded_mp4_payload(video: Any) -> dict[str, Any] | None:
+    """Return a post-process payload when the worker already encoded the MP4.
+
+    Returns ``None`` for the ordinary frame-tensor output so callers fall
+    through to their existing post-processing.
+
+    No fps is reported here: :func:`resolve_wan_output_fps` mirrors the fallback
+    the serving layer applies, so the rate the response advertises already
+    matches the rate these bytes were encoded at.
+    """
+    if isinstance(video, (bytes, bytearray, memoryview)):
+        video = [video]
+    if not isinstance(video, list) or not all(isinstance(item, (bytes, bytearray, memoryview)) for item in video):
+        return None
+    return {"payload": {"video": [bytes(item) for item in video]}, "metadata": {}}
+
+
+def resolve_wan_preencode_mp4(sampling_params: Any, *, output_type: str) -> bool:
+    """Return whether this request wants the worker to pre-encode the MP4.
+
+    Reject rather than silently downgrade the combinations the pre-encoded path
+    cannot serve: it emits MP4 bytes, so nothing downstream can interpolate
+    frames or hand back an array.
+    """
+    extra_args = getattr(sampling_params, "extra_args", None) or {}
+    if not extra_args.get("preencode_mp4", False):
+        return False
+    if getattr(sampling_params, "enable_frame_interpolation", False):
+        raise ValueError(
+            "preencode_mp4 cannot be combined with enable_frame_interpolation: interpolation "
+            "needs the decoded frames that the pre-encoded path no longer materializes"
+        )
+    if output_type != "np":
+        raise ValueError(f"preencode_mp4 returns MP4 bytes and cannot serve output_type={output_type!r}")
+    return True
+
+
+def resolve_wan_preencode_batch_frames(sampling_params: Any, *, default: int = WAN_DEFAULT_BATCH_FRAMES) -> int:
+    """Read the MP4 batching threshold before starting expensive generation."""
+    extra_args = getattr(sampling_params, "extra_args", None) or {}
+    return normalize_preencode_batch_frames(extra_args.get("preencode_batch_frames", default))
+
+
+def resolve_wan_video_codec_options(sampling_params: Any) -> dict[str, str] | None:
+    """Read the request's encoder options, preserving serving defaults.
+
+    An absent key inherits the one-shot serving encoder's defaults. Explicit
+    ``None`` and an empty mapping remain distinct caller choices.
+    """
+    extra_args = getattr(sampling_params, "extra_args", None) or {}
+    return normalize_video_codec_options(extra_args.get("video_codec_options", WAN_DEFAULT_VIDEO_CODEC_OPTIONS))
+
+
+def resolve_wan_output_fps(sampling_params: Any) -> int:
+    """Resolve the fps the worker encodes with.
+
+    The serving layer resolves output fps from the request's ``fps`` field, not
+    from ``frame_rate``, so pre-encoding has to read the same field or the bytes
+    would carry a rate the response does not advertise.
+    """
+    fps = getattr(sampling_params, "fps", None)
+    if isinstance(fps, list):
+        fps = fps[0] if fps else None
+    return int(fps) if fps else WAN_DEFAULT_OUTPUT_FPS

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -40,12 +40,20 @@ from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
+from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import (
+    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.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
 from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention, WanTransformer3DModel
 from vllm_omni.diffusion.offloader import OffloadPlan
 from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.chunked_video import decode_to_mp4
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.platforms import current_omni_platform
@@ -269,6 +277,9 @@ def get_wan22_post_process_func(
             output_type = sampling_params.output_type
         if output_type == "latent":
             return video
+        encoded = wan_preencoded_mp4_payload(video)
+        if encoded is not None:
+            return encoded
         video_metadata = {}
         if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
             video, multiplier = interpolate_video_tensor(
@@ -711,6 +722,8 @@ class Wan22Pipeline(
             num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps
 
         output_type = common.output_type or "np"
+        preencode_mp4 = resolve_wan_preencode_mp4(common, output_type=output_type)
+        preencode_batch_frames = resolve_wan_preencode_batch_frames(common) if preencode_mp4 else 17
         num_outputs_per_prompt = common.num_outputs_per_prompt or 1
         attention_kwargs: dict | None = None
 
@@ -965,27 +978,36 @@ class Wan22Pipeline(
                 latents.device, latents.dtype
             )
             latents = latents / latents_std + latents_mean
-            decoded = self.vae.decode(latents, return_dict=False)[0]
-            # Distributed VAE decode uses broadcast_result=False, so only the
-            # output-owning rank receives the full [B, C, T, H, W] video; other
-            # ranks get an empty placeholder. Emit typed media only from the
-            # owning rank and keep the placeholder on the legacy output field, so
-            # the media batch-dimension check in split_diffusion_output_by_request
-            # does not trip on every non-owner rank.
-            if decoded.dim() == 5:
-                output = None
-                media = DiffusionMediaOutput(
-                    video=VideoMediaOutput(
-                        tensor=decoded,
-                        spec=VideoTensorSpec(
-                            layout=VideoTensorLayout.BCTHW,
-                            encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
-                            value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
-                        ),
-                    )
+            if preencode_mp4:
+                output = decode_to_mp4(
+                    self.vae,
+                    latents,
+                    fps=resolve_wan_output_fps(common),
+                    batch_frames=preencode_batch_frames,
+                    video_codec_options=resolve_wan_video_codec_options(common),
                 )
             else:
-                output = decoded
+                decoded = self.vae.decode(latents, return_dict=False)[0]
+                # Distributed VAE decode uses broadcast_result=False, so only the
+                # output-owning rank receives the full [B, C, T, H, W] video; other
+                # ranks get an empty placeholder. Emit typed media only from the
+                # owning rank and keep the placeholder on the legacy output field, so
+                # the media batch-dimension check in split_diffusion_output_by_request
+                # does not trip on every non-owner rank.
+                if decoded.dim() == 5:
+                    output = None
+                    media = DiffusionMediaOutput(
+                        video=VideoMediaOutput(
+                            tensor=decoded,
+                            spec=VideoTensorSpec(
+                                layout=VideoTensorLayout.BCTHW,
+                                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+                            ),
+                        )
+                    )
+                else:
+                    output = decoded
 
         if DEBUG_PERF:
             current_omni_platform.synchronize()

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -31,6 +31,13 @@ from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.utils import _load_json
+from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import (
+    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.models.wan2_2.pipeline_wan2_2 import (
     _WAN_TEXT_ENCODER_OFFLOAD_PLAN,
     build_wan_scheduler,
@@ -46,6 +53,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3
 from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.chunked_video import decode_to_mp4
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.platforms import current_omni_platform
@@ -70,6 +78,9 @@ def get_wan22_i2v_post_process_func(
             output_type = sampling_params.output_type
         if output_type == "latent":
             return video
+        encoded = wan_preencoded_mp4_payload(video)
+        if encoded is not None:
+            return encoded
         video_metadata = {}
         if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
             video, multiplier = interpolate_video_tensor(
@@ -507,6 +518,8 @@ class Wan22I2VPipeline(
         num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps
 
         output_type = common.output_type or "np"
+        preencode_mp4 = resolve_wan_preencode_mp4(common, output_type=output_type)
+        preencode_batch_frames = resolve_wan_preencode_batch_frames(common) if preencode_mp4 else 17
         num_outputs_per_prompt = common.num_outputs_per_prompt or 1
         attention_kwargs: dict | None = None
 
@@ -736,7 +749,16 @@ class Wan22I2VPipeline(
                 latents.device, latents.dtype
             )
             latents = latents / latents_std + latents_mean
-            output = self.vae.decode(latents, return_dict=False)[0]
+            if preencode_mp4:
+                output = decode_to_mp4(
+                    self.vae,
+                    latents,
+                    fps=resolve_wan_output_fps(common),
+                    batch_frames=preencode_batch_frames,
+                    video_codec_options=resolve_wan_video_codec_options(common),
+                )
+            else:
+                output = self.vae.decode(latents, return_dict=False)[0]
 
         if DEBUG_PERF:
             current_omni_platform.synchronize()

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -33,6 +33,11 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL
 from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
+from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import (
+    resolve_wan_preencode_batch_frames,
+    resolve_wan_preencode_mp4,
+    resolve_wan_video_codec_options,
+)
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     load_transformer_config,
     load_wan_weights_with_optional_gate,
@@ -44,6 +49,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import (
 from vllm_omni.diffusion.offloader.config import DIT_COMPONENT, selected_offload_components
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.platforms import current_omni_platform
@@ -63,6 +69,10 @@ _S2V_DEFAULT_NEG_PROMPT = (
 
 _AUDIO_VIDEO_RATE = 30
 
+# S2V generates at a fixed rate; the worker-side encoder and the post-process
+# payload must report the same one.
+S2V_OUTPUT_FPS = 16
+
 
 def _s2v_audio_condition_key(audio: str | np.ndarray | None) -> tuple[Any, ...]:
     if isinstance(audio, str):
@@ -280,6 +290,10 @@ def get_wan22_s2v_post_process_func(
         if output_type == "latent":
             return video
 
+        if isinstance(video, list) and all(isinstance(item, (bytes, bytearray, memoryview)) for item in video):
+            # The worker already muxed the audio into these bytes.
+            return {"video": [bytes(item) for item in video], "audio": None, "fps": S2V_OUTPUT_FPS}
+
         processed_video = video_processor.postprocess_video(video, output_type=output_type)
 
         if audio_waveform is not None:
@@ -287,7 +301,7 @@ def get_wan22_s2v_post_process_func(
                 "video": processed_video,
                 "audio": audio_waveform,
                 "audio_sample_rate": audio_sr,
-                "fps": 16,
+                "fps": S2V_OUTPUT_FPS,
             }
         return processed_video
 
@@ -1177,6 +1191,13 @@ class Wan22S2VPipeline(
         )
         num_steps = common.num_inference_steps or num_inference_steps
         num_outputs_per_prompt = common.num_outputs_per_prompt or 1
+        # The request's output_type is what serving honours downstream; the keyword
+        # argument only carries a default for direct callers. Reading it here is what
+        # lets pre-encoding reject pil, pt, and latent instead of returning MP4 bytes
+        # to a caller that asked for frames.
+        output_type = common.output_type or output_type or "np"
+        preencode_mp4 = resolve_wan_preencode_mp4(common, output_type=output_type)
+        preencode_batch_frames = resolve_wan_preencode_batch_frames(common, default=1) if preencode_mp4 else 1
 
         if common.guidance_scale_provided:
             guidance_scale = common.guidance_scale
@@ -1315,170 +1336,207 @@ class Wan22S2VPipeline(
 
         # ---- 6. Multi-clip autoregressive denoising ----
         clips: list[torch.Tensor] = []
+        # Encode each finished clip while the next one denoises, instead of
+        # concatenating the whole video first.
+        mp4_session = (
+            ChunkedVideoMP4Session(
+                value_range=self.vae.chunk_value_range,
+                audio_waveforms=[raw_audio_waveforms[index // num_outputs_per_prompt] for index in range(batch_size)],
+                audio_sample_rate=raw_audio_sr,
+                fps=S2V_OUTPUT_FPS,
+                batch_frames=preencode_batch_frames,
+                video_codec_options=resolve_wan_video_codec_options(common),
+            )
+            if preencode_mp4
+            else None
+        )
         # Keep a pixel-space buffer of the trailing motion_frames for the
         # autoregressive connection between clips.
         videos_last_frames = torch.zeros([batch_size, 3, motion_frames, height, width], dtype=dtype, device=device)
 
-        for r in range(num_repeat):
-            clip_generators = _make_clip_generators(seeds, generators, r, device)
-            # -- Noise --
-            if r == 0 and request_latents is not None:
-                latents = request_latents.to(device=device, dtype=dtype)
-            else:
-                latents = torch.stack(
-                    [
-                        self.prepare_latents(
-                            infer_frames=infer_frames,
-                            height=height,
-                            width=width,
-                            motion_frames=motion_frames,
-                            lat_motion_frames=lat_motion_frames,
-                            dtype=dtype,
-                            device=device,
-                            generator=clip_generator,
-                        )
-                        for clip_generator in clip_generators
-                    ]
+        try:
+            for r in range(num_repeat):
+                clip_generators = _make_clip_generators(seeds, generators, r, device)
+                # -- Noise --
+                if r == 0 and request_latents is not None:
+                    latents = request_latents.to(device=device, dtype=dtype)
+                else:
+                    latents = torch.stack(
+                        [
+                            self.prepare_latents(
+                                infer_frames=infer_frames,
+                                height=height,
+                                width=width,
+                                motion_frames=motion_frames,
+                                lat_motion_frames=lat_motion_frames,
+                                dtype=dtype,
+                                device=device,
+                                generator=clip_generator,
+                            )
+                            for clip_generator in clip_generators
+                        ]
+                    )
+
+                # -- Scheduler --
+                self.scheduler.set_timesteps(num_steps, device=device, shift=self._flow_shift)
+                timesteps = self.scheduler.timesteps
+                self._num_timesteps = len(timesteps)
+
+                # -- Slice audio for this clip --
+                left_idx = r * infer_frames
+                right_idx = left_idx + infer_frames
+                audio_input = audio_emb[..., left_idx:right_idx]
+
+                # -- Pose condition for this clip --
+                # Default: zero condition (no pose driving)
+                # Shape: [B, C=16, T, H, W]
+                lat_h = height // self.vae_scale_factor_spatial
+                lat_w = width // self.vae_scale_factor_spatial
+                lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames
+                cond_latents = torch.zeros(
+                    [batch_size, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device
                 )
 
-            # -- Scheduler --
-            self.scheduler.set_timesteps(num_steps, device=device, shift=self._flow_shift)
-            timesteps = self.scheduler.timesteps
-            self._num_timesteps = len(timesteps)
+                # -- Clone motion latents for this clip --
+                input_motion_latents = motion_latents.clone()
 
-            # -- Slice audio for this clip --
-            left_idx = r * infer_frames
-            right_idx = left_idx + infer_frames
-            audio_input = audio_emb[..., left_idx:right_idx]
+                # Max sequence length for the transformer
+                max_seq_len = int(np.prod([lat_target_frames, lat_h, lat_w]) // 4)
 
-            # -- Pose condition for this clip --
-            # Default: zero condition (no pose driving)
-            # Shape: [B, C=16, T, H, W]
-            lat_h = height // self.vae_scale_factor_spatial
-            lat_w = width // self.vae_scale_factor_spatial
-            lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames
-            cond_latents = torch.zeros([batch_size, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device)
-
-            # -- Clone motion latents for this clip --
-            input_motion_latents = motion_latents.clone()
-
-            # Max sequence length for the transformer
-            max_seq_len = int(np.prod([lat_target_frames, lat_h, lat_w]) // 4)
-
-            # -- Precompute audio embeddings once per clip --
-            mf = [motion_frames, lat_motion_frames]
-            # When CPU offload is active the transformer lives on CPU until
-            # its __call__ hook fires.  encode_audio() bypasses __call__,
-            # so move the audio sub-module to GPU explicitly.
-            _audio_enc = getattr(self.transformer, "casual_audio_encoder", None)
-            _audio_on_cpu = _audio_enc is not None and next(_audio_enc.parameters()).device.type == "cpu"
-            if _audio_on_cpu:
-                _audio_enc.to(device)
-            param_dtype = next(self.transformer.parameters()).dtype
-            with torch.amp.autocast(device.type, dtype=param_dtype):
-                positive_audio_emb = self.transformer.encode_audio(audio_input, mf)
-            do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None
-            if do_true_cfg:
+                # -- Precompute audio embeddings once per clip --
+                mf = [motion_frames, lat_motion_frames]
+                # When CPU offload is active the transformer lives on CPU until
+                # its __call__ hook fires.  encode_audio() bypasses __call__,
+                # so move the audio sub-module to GPU explicitly.
+                _audio_enc = getattr(self.transformer, "casual_audio_encoder", None)
+                _audio_on_cpu = _audio_enc is not None and next(_audio_enc.parameters()).device.type == "cpu"
+                if _audio_on_cpu:
+                    _audio_enc.to(device)
+                param_dtype = next(self.transformer.parameters()).dtype
                 with torch.amp.autocast(device.type, dtype=param_dtype):
-                    negative_audio_emb = self.transformer.encode_audio(0.0 * audio_input, mf)
-            else:
-                negative_audio_emb = None
-            if _audio_on_cpu:
-                _audio_enc.to("cpu")
+                    positive_audio_emb = self.transformer.encode_audio(audio_input, mf)
+                do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None
+                if do_true_cfg:
+                    with torch.amp.autocast(device.type, dtype=param_dtype):
+                        negative_audio_emb = self.transformer.encode_audio(0.0 * audio_input, mf)
+                else:
+                    negative_audio_emb = None
+                if _audio_on_cpu:
+                    _audio_enc.to("cpu")
 
-            # -- Denoising loop --
-            latents = self.diffuse(
-                latents=latents,
-                timesteps=timesteps,
-                prompt_embeds=prompt_embeds,
-                negative_prompt_embeds=negative_prompt_embeds,
-                guidance_scale=guidance_scale,
-                clip_generator=clip_generators,
-                dtype=dtype,
-                device=device,
-                max_seq_len=max_seq_len,
-                cond_latents=cond_latents,
-                input_motion_latents=input_motion_latents,
-                ref_latents=ref_latents,
-                motion_frames=mf,
-                drop_first_motion=drop_first_motion and r == 0,
-                positive_audio_emb=positive_audio_emb,
-                negative_audio_emb=negative_audio_emb,
-            )
-
-            # ---- Decode this clip ----
-            if self._should_release_dit_before_decode():
-                self.transformer.to("cpu")
-                current_omni_platform.empty_cache()
-
-            if not (drop_first_motion and r == 0):
-                decode_latents = torch.cat([motion_latents, latents], dim=2)
-            else:
-                decode_latents = torch.cat([ref_latents, latents], dim=2)
-
-            decode_latents = self._denormalize_latents(decode_latents)
-            decode_latents = decode_latents.to(self.vae.dtype)
-            clip_video = self.vae.decode(decode_latents, return_dict=False)[0]  # [1, C, T, H, W]
-
-            # Handle VAE patch parallel: only rank 0 gets result, broadcast to all ranks
-            # This is needed for S2V's autoregressive loop where all ranks need the decoded frames
-            if clip_video.numel() == 0:
-                # Non-rank0 received empty tensor from patch parallel decode
-                # Use the same broadcast mechanism as the VAE patch parallel code
-                import torch.distributed as dist
-
-                total_frames = decode_latents.shape[2]
-
-                # Create buffer for broadcast
-                clip_video = torch.empty(
-                    (batch_size, 3, total_frames, height, width),
-                    device=decode_latents.device,
-                    dtype=decode_latents.dtype,
+                # -- Denoising loop --
+                latents = self.diffuse(
+                    latents=latents,
+                    timesteps=timesteps,
+                    prompt_embeds=prompt_embeds,
+                    negative_prompt_embeds=negative_prompt_embeds,
+                    guidance_scale=guidance_scale,
+                    clip_generator=clip_generators,
+                    dtype=dtype,
+                    device=device,
+                    max_seq_len=max_seq_len,
+                    cond_latents=cond_latents,
+                    input_motion_latents=input_motion_latents,
+                    ref_latents=ref_latents,
+                    motion_frames=mf,
+                    drop_first_motion=drop_first_motion and r == 0,
+                    positive_audio_emb=positive_audio_emb,
+                    negative_audio_emb=negative_audio_emb,
                 )
 
-                # Get the VAE's patch parallel group (same one used in decode)
-                vae_pp_group = getattr(self.vae, "_vae_pp_group", None)
-                if vae_pp_group is not None:
-                    # Broadcast using the same group as VAE patch parallel
-                    dist.broadcast(clip_video, src=0, group=vae_pp_group)
+                # ---- Decode this clip ----
+                if self._should_release_dit_before_decode():
+                    self.transformer.to("cpu")
+                    current_omni_platform.empty_cache()
 
-            # Trim to the infer_frames of interest
-            clip_video = clip_video[:, :, -infer_frames:]
-            if drop_first_motion and r == 0:
-                # Drop the first 3 frames (artifact from ref_latents prepend)
-                clip_video = clip_video[:, :, 3:]
+                if not (drop_first_motion and r == 0):
+                    decode_latents = torch.cat([motion_latents, latents], dim=2)
+                else:
+                    decode_latents = torch.cat([ref_latents, latents], dim=2)
 
-            # ---- Update motion for next clip (autoregressive) ----
-            overlap_frames_num = min(motion_frames, clip_video.shape[2])
-            videos_last_frames = torch.cat(
-                [videos_last_frames[:, :, overlap_frames_num:], clip_video[:, :, -overlap_frames_num:]],
-                dim=2,
+                decode_latents = self._denormalize_latents(decode_latents)
+                decode_latents = decode_latents.to(self.vae.dtype)
+                clip_video = self.vae.decode(decode_latents, return_dict=False)[0]  # [1, C, T, H, W]
+                # Under VAE patch parallelism only the output-owning rank decodes a
+                # real clip; peers get an empty placeholder filled by the broadcast.
+                owns_output = clip_video.numel() > 0
+
+                # Handle VAE patch parallel: only rank 0 gets result, broadcast to all ranks
+                # This is needed for S2V's autoregressive loop where all ranks need the decoded frames
+                if clip_video.numel() == 0:
+                    # Non-rank0 received empty tensor from patch parallel decode
+                    # Use the same broadcast mechanism as the VAE patch parallel code
+                    import torch.distributed as dist
+
+                    total_frames = decode_latents.shape[2]
+
+                    # Create buffer for broadcast
+                    clip_video = torch.empty(
+                        (batch_size, 3, total_frames, height, width),
+                        device=decode_latents.device,
+                        dtype=decode_latents.dtype,
+                    )
+
+                    # Get the VAE's patch parallel group (same one used in decode)
+                    vae_pp_group = getattr(self.vae, "_vae_pp_group", None)
+                    if vae_pp_group is not None:
+                        # Broadcast using the same group as VAE patch parallel
+                        dist.broadcast(clip_video, src=0, group=vae_pp_group)
+
+                # Trim to the infer_frames of interest
+                clip_video = clip_video[:, :, -infer_frames:]
+                if drop_first_motion and r == 0:
+                    # Drop the first 3 frames (artifact from ref_latents prepend)
+                    clip_video = clip_video[:, :, 3:]
+
+                # ---- Update motion for next clip (autoregressive) ----
+                overlap_frames_num = min(motion_frames, clip_video.shape[2])
+                videos_last_frames = torch.cat(
+                    [videos_last_frames[:, :, overlap_frames_num:], clip_video[:, :, -overlap_frames_num:]],
+                    dim=2,
+                )
+                videos_last_frames = videos_last_frames.to(dtype=dtype, device=device)
+                # Only prepare motion latents if there's another clip coming
+                if r < num_repeat - 1:
+                    motion_latents = self.prepare_motion_latents(videos_last_frames, device=device).to(dtype=dtype)
+
+                if mp4_session is not None:
+                    # Every rank needs the frames for the motion loop above, but the
+                    # executor keeps only the owner's response, so a peer skips the
+                    # encode and finishes with no containers, like T2V's non-owners.
+                    if owns_output:
+                        mp4_session.push(clip_video)
+                else:
+                    clips.append(clip_video.cpu())
+
+                # Free VRAM between clips
+                if current_omni_platform.is_available():
+                    current_omni_platform.empty_cache()
+
+            # ---- Assemble the output ----
+            if mp4_session is not None:
+                output = mp4_session.finish()  # one progressive MP4 per batch entry
+            else:
+                output = torch.cat(clips, dim=2)  # [B, C, T_total, H, W]
+
+            outputs = split_diffusion_output_by_request(
+                DiffusionOutput(
+                    output=output,
+                    stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+                ),
+                req,
+                num_outputs_per_prompt=num_outputs_per_prompt,
             )
-            videos_last_frames = videos_last_frames.to(dtype=dtype, device=device)
-            # Only prepare motion latents if there's another clip coming
-            if r < num_repeat - 1:
-                motion_latents = self.prepare_motion_latents(videos_last_frames, device=device).to(dtype=dtype)
-
-            clips.append(clip_video.cpu())
-
-            # Free VRAM between clips
-            if current_omni_platform.is_available():
-                current_omni_platform.empty_cache()
-
-        # ---- Concatenate all clips ----
-        output = torch.cat(clips, dim=2)  # [B, C, T_total, H, W]
-
-        outputs = split_diffusion_output_by_request(
-            DiffusionOutput(
-                output=output,
-                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
-            ),
-            req,
-            num_outputs_per_prompt=num_outputs_per_prompt,
-        )
-        for request_output, raw_audio_waveform in zip(outputs, raw_audio_waveforms):
-            request_output.output = (request_output.output, raw_audio_waveform, raw_audio_sr)
-        return outputs
+            if mp4_session is not None:
+                # The waveform is already muxed into each container.
+                return outputs
+            for request_output, raw_audio_waveform in zip(outputs, raw_audio_waveforms):
+                request_output.output = (request_output.output, raw_audio_waveform, raw_audio_sr)
+            return outputs
+        except BaseException:
+            if mp4_session is not None:
+                mp4_session.abort()
+            raise
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         return load_wan_weights_with_optional_gate(self, weights)

vllm_omni/diffusion/utils/chunked_video.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Turn committed VAE chunks into MP4 bytes without materializing the video.
+
+The producer side is the model's business: a Wan VAE walks a causal
+frame-by-frame loop, a MiniMax-H3 VAE walks overlapping clips, and Wan S2V
+decodes one clip per autoregressive iteration. What happens to a chunk once it
+is final is not: every one of them quantizes to uint8, moves to the host once,
+and lands in a bounded encoder. That consumer lives here so a model only has to
+publish finished chunks -- :class:`SupportsChunkedVAEDecode` for producers that
+can be driven, :class:`ChunkedVideoMP4Session` for producers that drive
+themselves.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode
+from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder
+
+
+def chunk_to_uint8_frames(chunk: torch.Tensor, value_range: tuple[float, float]) -> np.ndarray:
+    """Quantize a ``BCTHW`` chunk to ``BTHWC`` uint8, transferring once.
+
+    ``value_range`` is the interval the producer publishes, which differs per
+    checkpoint (see ``SupportsChunkedVAEDecode.chunk_value_range``). Quantizing
+    on the accelerator means one transfer moves the final bytes rather than
+    float frames.
+
+    That transfer is synchronous: ``.cpu()`` blocks inside the producer's chunk
+    callback, so what overlaps the remaining decode today is the CPU H.264
+    encode of earlier chunks, not this copy. Splitting device preparation,
+    device-to-host copy, and encode into three independently scheduled stages --
+    a pool of reusable pinned host slots with non-blocking copies on a dedicated
+    stream, leases released by the encoder -- fits behind this same callback
+    contract, but is not implemented here.
+    """
+    low, high = float(value_range[0]), float(value_range[1])
+    if high <= low:
+        raise ValueError(f"value_range must be increasing, got {value_range!r}")
+    scale = 255.0 / (high - low)
+    frames = chunk.clamp(low, high).sub(low).mul(scale).round().to(torch.uint8)
+    return frames.permute(0, 2, 3, 4, 1).cpu().numpy()
+
+
+class ChunkedVideoMP4Session:
+    """Encode committed video chunks into one progressive MP4 per batch entry.
+
+    Push finished ``BCTHW`` chunks as the producer commits them; each batch
+    entry gets its own bounded encoder, so host transfer and H.264 encoding
+    overlap whatever the producer is still decoding.
+
+    ``audio_waveforms`` holds one waveform per batch entry (``None`` for a
+    silent entry), so a caller with several outputs per prompt repeats a
+    request's waveform across its entries. ``batch_frames`` coalesces transfers
+    for producers that publish finer than a transfer is worth; ``crop`` trims
+    the decoder's padding to the requested output size.
+    """
+
+    def __init__(
+        self,
+        *,
+        value_range: tuple[float, float],
+        fps: float,
+        audio_waveforms: list[np.ndarray | None] | None = None,
+        audio_sample_rate: int | None = None,
+        batch_frames: int = 1,
+        max_pending: int = 2,
+        video_codec_options: dict[str, str] | None = None,
+        crop: tuple[int, int] | None = None,
+    ) -> None:
+        if batch_frames <= 0:
+            raise ValueError("batch_frames must be positive")
+        self._value_range = value_range
+        self._fps = fps
+        self._audio_waveforms = audio_waveforms
+        self._audio_sample_rate = audio_sample_rate
+        self._batch_frames = batch_frames
+        self._max_pending = max_pending
+        self._video_codec_options = video_codec_options
+        self._crop = crop
+        self._encoders: list[ChunkedMP4Encoder] = []
+        self._pending: list[torch.Tensor] = []
+        self._pending_frames = 0
+
+    def push(self, chunk: torch.Tensor) -> None:
+        """Queue one committed ``BCTHW`` chunk."""
+        if self._crop is not None:
+            height, width = self._crop
+            chunk = chunk[..., :height, :width]
+        self._pending.append(chunk)
+        self._pending_frames += int(chunk.shape[2])
+        if self._pending_frames >= self._batch_frames:
+            self._flush()
+
+    def finish(self) -> list[bytes]:
+        """Flush what is pending and return one MP4 per batch entry."""
+        self._flush()
+        return [encoder.finish() for encoder in self._encoders]
+
+    def abort(self) -> None:
+        for encoder in self._encoders:
+            encoder.abort()
+
+    def _flush(self) -> None:
+        if not self._pending:
+            return
+        frames = chunk_to_uint8_frames(torch.cat(self._pending, dim=2), self._value_range)
+        if not self._encoders:
+            self._encoders = [
+                ChunkedMP4Encoder(
+                    width=frames.shape[3],
+                    height=frames.shape[2],
+                    fps=self._fps,
+                    audio_waveform=self._waveform_for(index, frames.shape[0]),
+                    audio_sample_rate=self._audio_sample_rate,
+                    max_pending=self._max_pending,
+                    video_codec_options=self._video_codec_options,
+                )
+                for index in range(frames.shape[0])
+            ]
+        for index, encoder in enumerate(self._encoders):
+            encoder.push(np.ascontiguousarray(frames[index]))
+        self._pending.clear()
+        self._pending_frames = 0
+
+    def _waveform_for(self, index: int, batch_size: int) -> np.ndarray | None:
+        if self._audio_waveforms is None:
+            return None
+        if len(self._audio_waveforms) != batch_size:
+            raise ValueError(
+                f"expected one audio waveform per batch entry, got "
+                f"{len(self._audio_waveforms)} for {batch_size} entries"
+            )
+        return self._audio_waveforms[index]
+
+
+def decode_to_mp4(vae: Any, z: torch.Tensor, **session_kwargs: Any) -> list[bytes]:
+    """Decode ``z`` straight into one progressive MP4 per batch entry.
+
+    Drives a VAE that declares :class:`SupportsChunkedVAEDecode`, so host
+    transfer and encoding overlap the remaining decode and the full video is
+    never materialized. Ranks that own no decode output receive no chunks and
+    get an empty list, matching the empty tensor the full-decode path returns
+    there.
+    """
+    if not supports_chunked_vae_decode(vae):
+        raise TypeError(f"{type(vae).__name__} does not expose the chunked VAE decode capability")
+
+    session = ChunkedVideoMP4Session(value_range=vae.chunk_value_range, **session_kwargs)
+    try:
+        vae.decode_with_chunks(z, on_chunk=session.push)
+        return session.finish()
+    except BaseException:
+        session.abort()
+        raise

vllm_omni/diffusion/utils/media_utils.py

@@ -51,8 +51,11 @@ def _validate_video_chunk(chunk: np.ndarray, *, width: int, height: int) -> None
 class ChunkedMP4Encoder:
     """Encode temporal video chunks while the producer is still decoding.
 
-    A bounded queue and one muxing worker provide ordered backpressure while
-    keeping host memory bounded by ``max_pending`` chunks. Chunks use the same
+    A bounded queue and one muxing worker provide ordered backpressure. What
+    that bounds is pending raw-frame memory: at most ``max_pending`` chunks wait
+    to be encoded, so the bound is in chunks rather than bytes and a taller
+    frame or a longer chunk raises it. The muxed container is not bounded -- it
+    accumulates in memory and grows with output duration. Chunks use the same
     ``(T, H, W, 3)`` uint8 contract as :func:`mux_video_audio_bytes`.
     """