[Core][Diffusion] Add a typed pre-D2H video media contract (#6615)

Signed-off-by: NancyFyong <NancyFyong@users.noreply.github.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: NancyFyong <NancyFyong@users.noreply.github.com>
Co-authored-by: Samit <285365963@qq.com>

.buildkite/cuda/test-nightly.yml

@@ -389,6 +389,14 @@ steps:
           - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · Wan2.2 Device-Postprocess Equivalence Test"
+        timeout_in_minutes: 90
+        commands:
+          - |
+            VLLM_OMNI_DEVICE_POSTPROCESS_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+              pytest -s -v tests/e2e/offline_inference/test_device_postprocess_equivalence.py -m "full_model and cuda" --run-level "full_model"
+        mirror_hardwares: h100_1
+
       - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
         source_file_dependencies: diffusion_wan22_accuracy
         timeout_in_minutes: 180

docs/.nav.yml

@@ -71,6 +71,7 @@ nav:
       - Overview and Compatibility:
         - Overview: user_guide/diffusion_features.md
         - Feature Compatibility: user_guide/feature_compatibility.md
+      - Device-side Video Postprocessing: user_guide/diffusion/device_side_video_postprocess.md
       - CPU Offloading:
         - Overview: user_guide/diffusion/cpu_offload.md
         - Model-Level: user_guide/diffusion/offloader/module_offload.md

docs/mkdocs/hooks/generate_argparse.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import ast
+import json
 import logging
 import re
 import sys
@@ -124,6 +125,7 @@ def extract_omni_serve_subparser_init():
                         "make_arg_parser": lambda parser: parser,  # no-op for doc
                         "_ensure_vllm_platform": lambda: None,  # no-op for doc
                         "_nonneg_finite_float": float,
+                        "_json_object": json.loads,
                         "_parse_stage_overrides": __import__("json").loads,
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,

docs/user_guide/diffusion/device_side_video_postprocess.md

@@ -0,0 +1,72 @@
+# Device-side video postprocessing
+
+vLLM-Omni can convert decoded WAN2.2 video from normalized floating-point
+`[B, C, T, H, W]` tensors to contiguous `uint8` `[B, T, H, W, C]` frames before
+the worker-to-engine device-to-host transfer. This reduces the transported video
+payload by four times when the original IPC path widens bfloat16 to float32.
+
+The optimization is disabled by default. Enable it with a structured diffusion
+configuration:
+
+```python
+from vllm_omni import Omni
+
+engine = Omni(
+    model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+    video_output_transport={"enable_device_postprocess": True},
+)
+```
+
+For the server CLI, pass the same object as JSON:
+
+```bash
+vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni \
+  --video-output-transport '{"enable_device_postprocess": true}'
+```
+
+For a deploy configuration, place it on the diffusion stage:
+
+```yaml
+stages:
+  - stage_id: 0
+    video_output_transport:
+      enable_device_postprocess: true
+```
+
+## Runtime contract
+
+WAN2.2 emits decoded video through the typed `DiffusionOutput.media` field. The
+model runner validates and splits the batch into request-local tensors before
+preparing each video for transport. The worker then performs the D2H copy and
+shared-memory IPC using the prepared representation. The engine uses the generic
+media finalizer instead of a model-specific postprocessor.
+
+The runtime leaves the video in normalized floating-point form when:
+
+- `enable_device_postprocess` is false;
+- the requested output type is not `np`; or
+- frame interpolation still requires the floating-point tensor.
+
+These are policy fallbacks, not malformed-contract errors. If the device cannot
+allocate the temporary conversion buffer, the runner logs a warning and prepares
+the request-local float representation instead. Invalid tensor layout, encoding,
+value range, lifecycle state, and non-memory runtime failures still fail before
+IPC.
+
+## Precision
+
+Device preparation converts to float32 before denormalization and quantization.
+For bfloat16 WAN output this is more precise than the historical host path,
+which denormalizes in bfloat16. Relative to that path, values may differ by at
+most one level out of 255. Float32 inputs remain byte-identical after
+quantization.
+
+## Scope
+
+WAN2.2 is the reference migration for the first contract version. Other video
+pipelines continue using their existing output paths until they can declare the
+same typed media contract without bypassing model-specific float consumers such
+as safety checks or audio/video packaging.
+
+See [RFC #6541](https://github.com/vllm-project/vllm-omni/issues/6541) for the
+contract, lifecycle, batching rules, and migration plan.

tests/config/test_omni_config.py

@@ -1542,6 +1542,61 @@ def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
     assert cfg.quantization_config.get_name() == "int8"
 
 
+def test_video_output_transport_mapping_is_normalized() -> None:
+    from vllm_omni.diffusion.data import VideoOutputTransportConfig
+
+    cfg = omni_config_module._DiffusionConfigProjection(
+        video_output_transport={"enable_device_postprocess": True},
+    )
+
+    assert isinstance(cfg.video_output_transport, VideoOutputTransportConfig)
+    assert cfg.video_output_transport.enable_device_postprocess is True
+
+
+def test_omni_diffusion_config_normalizes_video_output_transport_mapping() -> None:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig, VideoOutputTransportConfig
+
+    cfg = OmniDiffusionConfig(
+        model=None,
+        video_output_transport={"enable_device_postprocess": True},
+    )
+
+    assert isinstance(cfg.video_output_transport, VideoOutputTransportConfig)
+    assert cfg.video_output_transport.enable_device_postprocess is True
+
+
+def test_video_output_transport_rejects_non_boolean_flag() -> None:
+    from vllm_omni.diffusion.data import VideoOutputTransportConfig
+
+    with pytest.raises(TypeError, match="enable_device_postprocess must be a bool"):
+        VideoOutputTransportConfig(enable_device_postprocess="true")  # type: ignore[arg-type]
+
+
+def test_video_output_transport_survives_stage_override_filtering() -> None:
+    from vllm_omni.config.stage_config import build_stage_runtime_overrides, deploy_runtime_override_keys
+
+    transport = {"enable_device_postprocess": True}
+    overrides = build_stage_runtime_overrides(0, {"video_output_transport": transport})
+
+    assert "video_output_transport" in deploy_runtime_override_keys()
+    assert overrides["video_output_transport"] == transport
+
+
+def test_video_output_transport_reaches_default_diffusion_stage() -> None:
+    from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+    transport = {"enable_device_postprocess": True}
+    stages = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "model": "unused",
+            "model_class_name": "UnknownPipeline",
+            "video_output_transport": transport,
+        }
+    )
+
+    assert stages[0]["engine_args"]["video_output_transport"] == transport
+
+
 def test_compact_offload_config_reaches_terminal_config(monkeypatch):
     from vllm_omni.diffusion.data import OmniDiffusionConfig
     from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy

tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -9,6 +9,7 @@ import pytest
 import torch
 from torch import nn
 
+from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
 from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -30,6 +31,22 @@ class _StubTextEncoder(nn.Module):
         return torch.float32
 
 
+class _StubVaeConfig:
+    latents_mean = [0.0, 0.0, 0.0, 0.0]
+    latents_std = [1.0, 1.0, 1.0, 1.0]
+    z_dim = 4
+
+
+class _StubVae(nn.Module):
+    dtype = torch.float32
+    config = _StubVaeConfig()
+
+    def decode(self, latents, return_dict=False):
+        del return_dict
+        batch, _, frames, height, width = latents.shape
+        return (torch.zeros(batch, 3, frames, height, width),)
+
+
 class _StubScheduler:
     def __init__(self, timesteps: list[int]) -> None:
         self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
@@ -75,6 +92,7 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline.transformer = _StubTransformer()
     pipeline.transformer_2 = None
     pipeline.text_encoder = _StubTextEncoder()
+    pipeline.vae = _StubVae()
     pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4)
     pipeline.scheduler = _StubScheduler([9, 5])
     pipeline.od_config = SimpleNamespace(flow_shift=5.0)
@@ -227,6 +245,99 @@ def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
     torch.testing.assert_close(outputs[1].output, latents_b)
 
 
+def test_forward_emits_request_local_typed_media_after_vae_decode() -> None:
+    pipeline = _make_pipeline()
+
+    def _fake_diffuse(
+        *,
+        latents,
+        timesteps,
+        prompt_embeds,
+        negative_prompt_embeds,
+        guidance_low,
+        guidance_high,
+        boundary_timestep,
+        dtype,
+        attention_kwargs,
+        latent_condition,
+        first_frame_mask,
+        generator,
+    ):
+        del (
+            timesteps,
+            prompt_embeds,
+            negative_prompt_embeds,
+            guidance_low,
+            guidance_high,
+            boundary_timestep,
+            dtype,
+            attention_kwargs,
+            latent_condition,
+            first_frame_mask,
+            generator,
+        )
+        return torch.zeros_like(latents)
+
+    pipeline.diffuse = _fake_diffuse  # type: ignore[method-assign]
+    batch = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt="prompt",
+                request_id="request-0",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_frames=1,
+                    num_inference_steps=2,
+                    max_sequence_length=32,
+                    output_type="np",
+                ),
+            )
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert len(outputs) == 1
+    assert outputs[0].output is None
+    assert outputs[0].media is not None
+    assert outputs[0].media.prepared_for_transport is False
+    assert outputs[0].media.video.tensor.shape == (1, 3, 1, 8, 8)
+    assert outputs[0].media.video.spec.layout is VideoTensorLayout.BCTHW
+    assert outputs[0].media.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+    assert outputs[0].media.video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE
+
+
+def test_forward_keeps_legacy_output_on_non_owner_vae_rank() -> None:
+    # Distributed VAE decode uses broadcast_result=False, so non-owner ranks get
+    # an empty placeholder instead of the full video. Wrapping that as typed media
+    # would fail split_diffusion_output_by_request's batch check on every non-owner
+    # rank, so the pipeline must keep the placeholder on the legacy output field.
+    pipeline = _make_pipeline()
+    pipeline.vae.decode = lambda latents, return_dict=False: (torch.empty(0),)  # type: ignore[assignment]
+    pipeline.diffuse = lambda **kwargs: torch.zeros_like(kwargs["latents"])  # type: ignore[method-assign]
+
+    batch = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt="prompt",
+                request_id="request-0",
+                sampling_params=OmniDiffusionSamplingParams(
+                    num_frames=1,
+                    num_inference_steps=2,
+                    max_sequence_length=32,
+                    output_type="np",
+                ),
+            )
+        ]
+    )
+
+    outputs = pipeline.forward(batch)
+
+    assert len(outputs) == 1
+    assert outputs[0].media is None
+    assert outputs[0].output is not None
+    assert outputs[0].output.numel() == 0
+
+
 def test_forward_batches_precomputed_prompt_embeddings() -> None:
     pipeline = _make_pipeline()
     diffuse_call = {}

tests/diffusion/test_async_output_worker.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for WorkerProc async output logic."""
 
 import queue
@@ -8,13 +8,23 @@ import time
 from unittest.mock import MagicMock
 
 import pytest
+import torch
 
+import vllm_omni.diffusion.ipc as diffusion_ipc
 from vllm_omni.diffusion.data import (
     AsyncDiffusionOutput,
     AsyncOutputKind,
     DiffusionOutput,
     OmniACK,
 )
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    VideoMediaOutput,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTensorSpec,
+    VideoValueRange,
+)
 from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
 from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 
@@ -182,6 +192,66 @@ class TestReturnResultSyncPath:
         mock_pack.assert_called_once_with(output)
         proc.result_mq.enqueue.assert_called_once_with(output)
 
+    def test_media_contract_error_is_not_sent_as_an_unpacked_result(self, mocker):
+        proc = _make_worker_proc(step_execution=True)
+        proc._async_output_queue = None
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+            side_effect=ValueError("media was not prepared"),
+        )
+
+        with pytest.raises(ValueError, match="media was not prepared"):
+            proc._return_result(DiffusionOutput())
+
+        proc.result_mq.enqueue.assert_not_called()
+
+    def test_media_packing_memory_failure_is_reraised_for_typed_media(self, mocker):
+        """A non-ValueError packing failure on typed media must re-raise, not
+        swallow it and enqueue a half-packed payload."""
+        proc = _make_worker_proc(step_execution=True)
+        proc._async_output_queue = None
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+            side_effect=RuntimeError("SHM pack out of memory"),
+        )
+
+        output = DiffusionOutput(media=MagicMock(spec=DiffusionMediaOutput))
+        with pytest.raises(RuntimeError, match="SHM pack out of memory"):
+            proc._return_result(output)
+
+        proc.result_mq.enqueue.assert_not_called()
+
+    def test_media_packing_failure_is_reraised_for_dp_container(self, mocker):
+        """Container dispatch and typed-media failure detection share a walker."""
+        proc = _make_worker_proc(step_execution=True)
+        proc._async_output_queue = None
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+            side_effect=RuntimeError("SHM pack out of memory"),
+        )
+
+        output = DiffusionOutput(media=MagicMock(spec=DiffusionMediaOutput))
+        tagged_output = {"dp_rank": 0, "output": output}
+        with pytest.raises(RuntimeError, match="SHM pack out of memory"):
+            proc._return_result(tagged_output)
+
+        proc.result_mq.enqueue.assert_not_called()
+
+    def test_legacy_output_packing_failure_is_still_swallowed(self, mocker):
+        """Non-typed (legacy output) failures keep the historic swallow-and-
+        enqueue behavior; only typed media re-raises."""
+        proc = _make_worker_proc(step_execution=True)
+        proc._async_output_queue = None
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+            side_effect=RuntimeError("SHM pack out of memory"),
+        )
+
+        proc._return_result(DiffusionOutput(output="data"))
+
+        # Legacy path: warning logged, output still enqueued once.
+        proc.result_mq.enqueue.assert_called_once_with(DiffusionOutput(output="data"))
+
     def test_non_diffusion_output_skips_async_in_request_mode(self, mocker):
         """Non-DiffusionOutput/BatchRunnerOutput in request mode uses sync path."""
         proc = _make_worker_proc(step_execution=False)
@@ -243,6 +313,70 @@ class TestAsyncOutputLoopLogic:
         )
         assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"
 
+    def test_late_field_failure_sends_error_without_partial_payload_or_shm_leak(
+        self,
+        mocker,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        from multiprocessing import shared_memory
+
+        proc = _make_worker_proc(step_execution=False)
+        media = DiffusionMediaOutput(
+            video=VideoMediaOutput(
+                tensor=torch.zeros(1, 3, 2, 4, 5),
+                spec=VideoTensorSpec(
+                    layout=VideoTensorLayout.BCTHW,
+                    encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                    value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+                ),
+            ),
+            prepared_for_transport=True,
+        )
+        output = DiffusionOutput(media=media, trajectory_latents=torch.ones(30))
+        original_latents = output.trajectory_latents
+
+        real_array_to_shm = diffusion_ipc._array_to_shm
+        created_handles: list[dict[str, object]] = []
+
+        def _fail_on_second(array):
+            if created_handles:
+                raise RuntimeError("late-field packing failure")
+            handle = real_array_to_shm(array)
+            created_handles.append(handle)
+            return handle
+
+        monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1)
+        monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second)
+
+        def _pack_without_stream(value, d2h_stream=None):
+            # The worker loop behavior is under test here. CPU tensors avoid
+            # introducing a device dependency into this unit test.
+            del d2h_stream
+            return diffusion_ipc.pack_diffusion_output_shm(value)
+
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
+            side_effect=_pack_without_stream,
+        )
+        mock_torch = mocker.MagicMock()
+        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
+
+        proc._async_output_queue.put((output, "late-failure", None))
+        proc._async_output_queue.put(None)
+        proc._async_output_loop()
+
+        messages = [call.args[0] for call in proc.result_mq.enqueue.call_args_list]
+        assert len(messages) == 1
+        assert messages[0].kind is AsyncOutputKind.OUTPUT_READY
+        assert messages[0].async_output_id == "late-failure"
+        assert messages[0].error == "Background D2H/SHM packing failed"
+        assert messages[0].output is None
+        assert output.media is media
+        assert output.trajectory_latents is original_latents
+        with pytest.raises(FileNotFoundError):
+            shared_memory.SharedMemory(name=str(created_handles[0]["name"]))
+
 
 class TestDrainAsyncOutputs:
     """Async outputs must finish packing before device memory is released."""

tests/diffusion/test_device_output_preparation.py

@@ -0,0 +1,460 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import FrozenInstanceError
+
+import numpy as np
+import pytest
+import torch
+
+import vllm_omni.diffusion.ipc as diffusion_ipc
+import vllm_omni.diffusion.postprocess.device_reduction as device_reduction
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig
+from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    FloatVideoConsumer,
+    VideoMediaOutput,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTensorSpec,
+    VideoTransportConstraints,
+    VideoValueRange,
+    ensure_request_owned_tensor,
+    slice_diffusion_media_output,
+)
+from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _config(*, enabled: bool) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model=None,
+        video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled),
+    )
+
+
+def _sampling(*, output_type: str | None = "np", interpolate: bool = False) -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        output_type=output_type,
+        enable_frame_interpolation=interpolate,
+    )
+
+
+def _media(
+    tensor: torch.Tensor,
+    *,
+    value_range: VideoValueRange = VideoValueRange.NEGATIVE_ONE_TO_ONE,
+    consumers: frozenset[FloatVideoConsumer] = frozenset(),
+) -> DiffusionMediaOutput:
+    return DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=tensor,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=value_range,
+            ),
+            constraints=VideoTransportConstraints(pending_float_consumers=consumers),
+        )
+    )
+
+
+def test_runtime_prepares_request_local_video_as_uint8() -> None:
+    video = torch.linspace(-1.0, 1.0, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
+
+    prepared = prepare_diffusion_media_for_transport(
+        _media(video),
+        od_config=_config(enabled=True),
+        sampling_params=_sampling(),
+    )
+
+    assert prepared.prepared_for_transport is True
+    assert prepared.video.tensor.dtype == torch.uint8
+    assert prepared.video.tensor.shape == (2, 2, 4, 5, 3)
+    assert prepared.video.spec == VideoTensorSpec(
+        layout=VideoTensorLayout.BTHWC,
+        encoding=VideoTensorEncoding.UINT8_FRAMES,
+        value_range=VideoValueRange.ZERO_TO_255,
+    )
+    assert prepared.video.constraints.pending_float_consumers == frozenset()
+
+
+def test_disabled_runtime_preserves_float_but_marks_media_prepared() -> None:
+    video = torch.randn(1, 3, 2, 4, 5)
+
+    prepared = prepare_diffusion_media_for_transport(
+        _media(video),
+        od_config=_config(enabled=False),
+        sampling_params=_sampling(),
+    )
+
+    assert prepared.prepared_for_transport is True
+    assert prepared.video.tensor is video
+    assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+
+
+def test_device_oom_falls_back_to_request_owned_float_media(monkeypatch: pytest.MonkeyPatch) -> None:
+    source = torch.randn(2, 3, 2, 4, 5)
+    sliced = slice_diffusion_media_output(_media(source), 0, 1)
+
+    def raise_oom(video: torch.Tensor, *, do_denormalize: bool) -> torch.Tensor:
+        raise torch.OutOfMemoryError("injected")
+
+    monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", raise_oom)
+
+    prepared = prepare_diffusion_media_for_transport(
+        sliced,
+        od_config=_config(enabled=True),
+        sampling_params=_sampling(),
+    )
+
+    assert prepared.prepared_for_transport is True
+    assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+    assert prepared.video.tensor.dtype == source.dtype
+    assert prepared.video.tensor.shape == (1, 3, 2, 4, 5)
+    assert prepared.video.tensor.is_contiguous()
+    assert prepared.video.tensor._base is None
+
+
+def test_device_non_oom_error_is_not_swallowed(monkeypatch: pytest.MonkeyPatch) -> None:
+    def raise_runtime_error(video: torch.Tensor, *, do_denormalize: bool) -> torch.Tensor:
+        raise RuntimeError("injected")
+
+    monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", raise_runtime_error)
+
+    with pytest.raises(RuntimeError, match="injected"):
+        prepare_diffusion_media_for_transport(
+            _media(torch.randn(1, 3, 2, 4, 5)),
+            od_config=_config(enabled=True),
+            sampling_params=_sampling(),
+        )
+
+
+def test_declared_and_request_float_consumers_are_unioned() -> None:
+    video = torch.randn(1, 3, 2, 4, 5)
+    media = _media(
+        video,
+        consumers=frozenset({FloatVideoConsumer.VIDEO_GUARDRAILS}),
+    )
+
+    prepared = prepare_diffusion_media_for_transport(
+        media,
+        od_config=_config(enabled=True),
+        sampling_params=_sampling(interpolate=True),
+    )
+
+    assert prepared.video.tensor is video
+    assert prepared.video.constraints.pending_float_consumers == frozenset(
+        {
+            FloatVideoConsumer.FRAME_INTERPOLATION,
+            FloatVideoConsumer.VIDEO_GUARDRAILS,
+        }
+    )
+
+
+def test_non_numpy_presentation_preserves_float() -> None:
+    prepared = prepare_diffusion_media_for_transport(
+        _media(torch.randn(1, 3, 2, 4, 5)),
+        od_config=_config(enabled=True),
+        sampling_params=_sampling(output_type="pil"),
+    )
+
+    assert prepared.prepared_for_transport is True
+    assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+
+
+def test_zero_to_one_video_is_not_denormalized() -> None:
+    video = torch.tensor([[[[[0.0, 0.5, 1.0]]]]]).expand(1, 3, 1, 1, 3)
+
+    prepared = prepare_diffusion_media_for_transport(
+        _media(video, value_range=VideoValueRange.ZERO_TO_ONE),
+        od_config=_config(enabled=True),
+        sampling_params=_sampling(),
+    )
+
+    expected = np.rint(video.permute(0, 2, 3, 4, 1).numpy() * 255.0).astype(np.uint8)
+    np.testing.assert_array_equal(prepared.video.tensor.numpy(), expected)
+
+
+def test_invalid_contract_raises_even_when_feature_is_off() -> None:
+    media = _media(torch.randn(3, 2, 4, 5))
+
+    with pytest.raises(ValueError, match="rank 5"):
+        prepare_diffusion_media_for_transport(
+            media,
+            od_config=_config(enabled=False),
+            sampling_params=_sampling(),
+        )
+
+
+@pytest.mark.parametrize(
+    ("tensor", "spec", "message"),
+    [
+        (
+            torch.randn(1, 2, 4, 5, 3),
+            VideoTensorSpec(
+                layout=VideoTensorLayout.BTHWC,
+                encoding=VideoTensorEncoding.UINT8_FRAMES,
+                value_range=VideoValueRange.ZERO_TO_255,
+            ),
+            "torch.uint8",
+        ),
+        (
+            torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8),
+            VideoTensorSpec(
+                layout=VideoTensorLayout.BTHWC,
+                encoding=VideoTensorEncoding.UINT8_FRAMES,
+                value_range=VideoValueRange.ZERO_TO_ONE,
+            ),
+            "incompatible value range",
+        ),
+        (
+            torch.randn(1, 4, 2, 4, 5),
+            VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            ),
+            "3 channels",
+        ),
+    ],
+)
+def test_invalid_spec_tensor_combinations_are_contract_errors(
+    tensor: torch.Tensor,
+    spec: VideoTensorSpec,
+    message: str,
+) -> None:
+    media = DiffusionMediaOutput(video=VideoMediaOutput(tensor=tensor, spec=spec))
+
+    with pytest.raises(ValueError, match=message):
+        media.validate()
+
+
+def test_pipeline_cannot_emit_unprepared_uint8_media() -> None:
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8),
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BTHWC,
+                encoding=VideoTensorEncoding.UINT8_FRAMES,
+                value_range=VideoValueRange.ZERO_TO_255,
+            ),
+        )
+    )
+
+    with pytest.raises(ValueError, match="must emit unprepared media as NORMALIZED_FLOAT"):
+        prepare_diffusion_media_for_transport(
+            media,
+            od_config=_config(enabled=True),
+            sampling_params=_sampling(),
+        )
+
+
+def test_descriptors_are_immutable() -> None:
+    media = _media(torch.randn(1, 3, 2, 4, 5))
+
+    with pytest.raises(FrozenInstanceError):
+        media.prepared_for_transport = True  # type: ignore[misc]
+    with pytest.raises(FrozenInstanceError):
+        media.video.spec.layout = VideoTensorLayout.BTHWC  # type: ignore[misc]
+
+
+def test_float_fallback_detaches_a_request_view_before_transport() -> None:
+    source = torch.randn(4, 3, 2, 4, 5)
+    sliced = slice_diffusion_media_output(_media(source), 1, 3)
+
+    prepared = prepare_diffusion_media_for_transport(
+        sliced,
+        od_config=_config(enabled=False),
+        sampling_params=_sampling(),
+    )
+
+    tensor = prepared.video.tensor
+    assert tensor.shape == (2, 3, 2, 4, 5)
+    assert tensor.is_contiguous()
+    assert tensor._base is None
+    assert tensor.untyped_storage().nbytes() == tensor.numel() * tensor.element_size()
+    assert tensor.untyped_storage().data_ptr() != source.untyped_storage().data_ptr()
+
+
+def test_full_batch_slice_does_not_duplicate_owned_storage() -> None:
+    source = torch.randn(2, 3, 2, 4, 5)
+
+    sliced = slice_diffusion_media_output(_media(source), 0, 2)
+
+    assert sliced.video.tensor.untyped_storage().data_ptr() == source.untyped_storage().data_ptr()
+
+
+def test_diffusion_output_to_cpu_moves_prepared_typed_media() -> None:
+    prepared = prepare_diffusion_media_for_transport(
+        _media(torch.randn(1, 3, 2, 4, 5)),
+        od_config=_config(enabled=False),
+        sampling_params=_sampling(),
+    )
+    output = DiffusionOutput(media=prepared, to_cpu=True)
+
+    assert output.media is not None
+    assert output.media.video.tensor.device.type == "cpu"
+
+
+def test_diffusion_output_to_cpu_rejects_unprepared_media() -> None:
+    with pytest.raises(ValueError, match="prepared before to_cpu"):
+        DiffusionOutput(media=_media(torch.randn(1, 3, 2, 4, 5)), to_cpu=True)
+
+
+def test_prepared_media_round_trips_through_typed_ipc(monkeypatch: pytest.MonkeyPatch) -> None:
+    expected = prepare_diffusion_media_for_transport(
+        _media(torch.randn(1, 3, 2, 256, 256)),
+        od_config=_config(enabled=True),
+        sampling_params=_sampling(),
+    )
+    output = DiffusionOutput(media=expected)
+
+    monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1)
+    pack_diffusion_output_shm(output)
+
+    assert isinstance(output.media, dict)
+    assert output.media["__type__"] == "diffusion_media_v1"
+    assert output.media["video"]["tensor"]["__tensor_shm__"] is True
+
+    unpack_diffusion_output_shm(output)
+
+    assert isinstance(output.media, DiffusionMediaOutput)
+    assert output.media.prepared_for_transport is True
+    assert output.media.video.spec == expected.video.spec
+    torch.testing.assert_close(output.media.video.tensor, expected.video.tensor)
+
+
+def test_float_media_constraints_round_trip_through_typed_ipc(monkeypatch: pytest.MonkeyPatch) -> None:
+    expected = prepare_diffusion_media_for_transport(
+        _media(torch.randn(1, 3, 2, 256, 256, dtype=torch.bfloat16)),
+        od_config=_config(enabled=True),
+        sampling_params=_sampling(interpolate=True),
+    )
+    output = DiffusionOutput(media=expected)
+
+    monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1)
+    pack_diffusion_output_shm(output)
+    unpack_diffusion_output_shm(output)
+
+    assert isinstance(output.media, DiffusionMediaOutput)
+    assert output.media.video.tensor.dtype is torch.bfloat16
+    assert output.media.video.constraints.pending_float_consumers == frozenset({FloatVideoConsumer.FRAME_INTERPOLATION})
+    torch.testing.assert_close(output.media.video.tensor, expected.video.tensor)
+
+
+def test_ipc_rejects_media_that_skipped_runner_preparation() -> None:
+    output = DiffusionOutput(media=_media(torch.randn(1, 3, 2, 4, 5)))
+
+    with pytest.raises(ValueError, match="must be prepared"):
+        pack_diffusion_output_shm(output)
+
+
+def test_ipc_rejects_unknown_media_schema_version() -> None:
+    prepared = prepare_diffusion_media_for_transport(
+        _media(torch.randn(1, 3, 2, 4, 5)),
+        od_config=_config(enabled=False),
+        sampling_params=_sampling(),
+    )
+    output = DiffusionOutput(media=prepared)
+    pack_diffusion_output_shm(output)
+    assert isinstance(output.media, dict)
+    output.media["schema_version"] = 2
+
+    with pytest.raises(ValueError, match="schema version"):
+        unpack_diffusion_output_shm(output)
+
+
+def test_prepare_rejects_media_a_pipeline_already_prepared(monkeypatch: pytest.MonkeyPatch) -> None:
+    # A pipeline must emit unprepared media so the runner applies request policy
+    # (e.g. frame interpolation). Accepting a forged prepared uint8 payload would
+    # skip the interpolation constraint; interpolator never runs and gets dropped.
+    forged = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8),
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BTHWC,
+                encoding=VideoTensorEncoding.UINT8_FRAMES,
+                value_range=VideoValueRange.ZERO_TO_255,
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+
+    with pytest.raises(ValueError, match="already prepared"):
+        prepare_diffusion_media_for_transport(
+            forged,
+            od_config=_config(enabled=True),
+            sampling_params=_sampling(interpolate=True),
+        )
+
+
+def test_prepare_rejects_prepared_media_despite_pending_interpolation(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # The forged payload is the exact case the reviewer described: device
+    # postprocessing enabled plus frame_interpolation, but the media already being
+    # prepared. It must be rejected, not silently returned with interpolation dropped.
+    forged = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=torch.zeros(1, 3, 2, 2, 3, dtype=torch.float32),
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+
+    with pytest.raises(ValueError, match="already prepared"):
+        prepare_diffusion_media_for_transport(
+            forged,
+            od_config=_config(enabled=True),
+            sampling_params=_sampling(interpolate=True),
+        )
+
+
+def test_request_owned_tensor_rejects_a_detached_storage_view() -> None:
+    # ``tensor._base is None`` is not proof of ownership: a detached slice keeps a
+    # nonzero storage offset and the whole source storage. Serializing it would
+    # leak every other request's pixels, so it must be cloned before transport.
+    source = torch.randn(4, 3, 2, 4, 5)
+    detached_slice = source[1:2].detach()
+    assert detached_slice.is_contiguous()
+    assert detached_slice._base is None
+    assert detached_slice.storage_offset() != 0
+
+    owned = ensure_request_owned_tensor(detached_slice)
+
+    assert owned.shape == detached_slice.shape
+    assert owned.storage_offset() == 0
+    assert owned.untyped_storage().nbytes() == owned.numel() * owned.element_size()
+    assert owned.untyped_storage().data_ptr() != source.untyped_storage().data_ptr()
+    torch.testing.assert_close(owned, detached_slice)
+
+
+def test_prepared_media_rejects_non_compact_request_storage() -> None:
+    # A zero-offset view of a bigger buffer is contiguous but not compact, so
+    # prepared transport must refuse it just as it refuses a storage view.
+    source = torch.randn(2, 3, 2, 4, 5)
+    zero_offset_view = source[0:1]
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=zero_offset_view,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+
+    with pytest.raises(ValueError, match="request-local storage"):
+        media.validate()

tests/diffusion/test_device_reduction.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""GPU correctness/size tests for device-side video reduction.
+
+The reduction must be byte-identical to the current production output (SHM
+widens bf16->f32, then diffusers postprocess, then the API server's ``*255``
+rounding), otherwise enabling it would silently change pixels.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.postprocess.device_reduction import reduce_video_to_uint8_frames
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.diffusion,
+    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+]
+
+_GPU = torch.accelerator.is_available() if hasattr(torch, "accelerator") else torch.cuda.is_available()
+requires_gpu = pytest.mark.skipif(not _GPU, reason="device reduction is a GPU path")
+
+
+def _device() -> str:
+    return "cuda:0"
+
+
+def _reference_uint8(video: torch.Tensor, *, compute_dtype: torch.dtype | None = torch.float32) -> np.ndarray:
+    """The diffusers postprocess plus the API server's rounding, for the same input.
+
+    ``compute_dtype`` selects the precision the denormalize runs at.
+    ``torch.float32`` is the precision the device reduction uses; ``None`` keeps
+    the input's own dtype, which is what the engine's postprocess actually does.
+    """
+    from diffusers.video_processor import VideoProcessor
+
+    source = video if compute_dtype is None else video.to(compute_dtype)
+    processed = VideoProcessor(vae_scale_factor=8).postprocess_video(source, output_type="np")
+    scaled = np.clip(processed.astype(np.float32), 0.0, 1.0) * 255.0
+    np.rint(scaled, out=scaled)
+    return scaled.astype(np.uint8)
+
+
+@requires_gpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+def test_device_reduction_matches_a_float32_reference(dtype: torch.dtype) -> None:
+    """The reduction must be the float32 computation, exactly, for every input dtype."""
+    torch.manual_seed(0)
+    # Values slightly outside [-1, 1] so the clamp is actually exercised.
+    video = (torch.rand(1, 3, 8, 64, 96, device=_device()) * 2.4 - 1.2).to(dtype)
+
+    reference = _reference_uint8(video)
+    produced = reduce_video_to_uint8_frames(video).cpu().numpy()
+
+    assert produced.shape == reference.shape == (1, 8, 64, 96, 3)
+    assert produced.dtype == np.uint8
+    np.testing.assert_array_equal(produced, reference)
+
+
+@requires_gpu
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_device_reduction_deviates_from_a_narrow_postprocess_by_at_most_one(dtype: torch.dtype) -> None:
+    """Pin the cost of computing in float32: at most one 255th.
+
+    A postprocess that denormalizes in the VAE's own dtype (WAN, HunyuanVideo,
+    LTX-2, Cosmos3) does not match this reduction bit for bit; ~20% of pixels
+    differ by one on a real 480x832x81 WAN generation. Keep that a bound rather
+    than a surprise, so a regression widening it fails here.
+    """
+    torch.manual_seed(0)
+    video = (torch.rand(1, 3, 8, 64, 96, device=_device()) * 2.4 - 1.2).to(dtype)
+
+    narrow = _reference_uint8(video, compute_dtype=None)
+    produced = reduce_video_to_uint8_frames(video).cpu().numpy()
+
+    deviation = np.abs(produced.astype(np.int16) - narrow.astype(np.int16))
+    assert deviation.max() <= 1
+    # The narrow path really is different, so the bound above is not vacuous.
+    assert deviation.any()
+
+
+@requires_gpu
+def test_device_reduction_shrinks_hop1_payload_4x() -> None:
+    """The whole point: the D2H payload drops from float32 to uint8 (4x)."""
+    video = torch.rand(1, 3, 8, 704, 1280, device=_device(), dtype=torch.bfloat16) * 2 - 1
+
+    # Production transports bf16 widened to float32 over SHM.
+    widened_f32_bytes = video.numel() * 4
+    reduced = reduce_video_to_uint8_frames(video)
+    reduced_bytes = reduced.numel() * reduced.element_size()
+
+    assert reduced.dtype == torch.uint8
+    assert widened_f32_bytes / reduced_bytes == pytest.approx(4.0, rel=1e-6)
+
+
+@requires_gpu
+def test_device_reduction_stays_on_device() -> None:
+    """Reduction must happen before D2H, so the result is still on the GPU."""
+    video = torch.rand(1, 3, 4, 32, 48, device=_device(), dtype=torch.bfloat16) * 2 - 1
+    reduced = reduce_video_to_uint8_frames(video)
+    assert reduced.device.type == "cuda"
+    assert reduced.is_contiguous()
+
+
+@requires_gpu
+def test_device_reduction_without_denormalize_matches_already_unit_range() -> None:
+    """Pipelines with do_normalize=False emit [0,1]; skip the denorm step."""
+    video = torch.rand(1, 3, 4, 16, 16, device=_device(), dtype=torch.float32)  # already [0,1]
+
+    produced = reduce_video_to_uint8_frames(video, do_denormalize=False).cpu().numpy()
+
+    expected = video.permute(0, 2, 3, 4, 1).cpu().numpy()
+    expected = np.rint(np.clip(expected, 0, 1) * 255.0).astype(np.uint8)
+    np.testing.assert_array_equal(produced, expected)
+
+
+def test_device_reduction_rejects_non_5d() -> None:
+    """A wrong-rank tensor is a programming error, not a silent reshape."""
+    with pytest.raises(ValueError, match=r"\[B, C, F, H, W\]"):
+        reduce_video_to_uint8_frames(torch.rand(3, 8, 64, 64))

tests/diffusion/test_diffusion_model_runner.py

@@ -10,9 +10,19 @@ import torch
 
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, VideoOutputTransportConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    VideoMediaOutput,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTensorSpec,
+    VideoValueRange,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput, NewRequestData
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -144,6 +154,19 @@ def _make_request():
     )
 
 
+def _raw_video_media(batch_size: int = 1, tensor: torch.Tensor | None = None) -> DiffusionMediaOutput:
+    return DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=tensor if tensor is not None else torch.randn(batch_size, 3, 2, 4, 5),
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            ),
+        )
+    )
+
+
 def _make_request_with_params(req_id: str, sampling_params):
     return SimpleNamespace(
         request_id=req_id,
@@ -214,6 +237,7 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
         enable_cache_dit_summary=enable_cache_dit_summary,
         parallel_config=SimpleNamespace(use_hsdp=False),
         streaming_output=False,
+        video_output_transport=VideoOutputTransportConfig(),
     )
     runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend(
         vllm_config=runner.vllm_config,
@@ -546,6 +570,48 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
     assert fourth.get_request_output("req").finished is True
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_stepwise_prepares_typed_media(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.pipeline = _ChunkStepPipeline([DiffusionOutput(media=_raw_video_media())])
+    runner.od_config.streaming_output = True
+    runner.od_config.step_execution = True
+    runner.od_config.video_output_transport = VideoOutputTransportConfig(enable_device_postprocess=True)
+    req = OmniDiffusionRequest(
+        prompt="a prompt",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=4),
+        request_id="req",
+    )
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    scheduler_output = DiffusionSchedulerOutput(
+        step_id=0,
+        scheduled_new_reqs=[NewRequestData(request_id="req", req=req)],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        finished_req_ids=set(),
+        num_running_reqs=1,
+        num_waiting_reqs=0,
+    )
+
+    DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+    scheduler_output = DiffusionSchedulerOutput(
+        step_id=1,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData(request_ids=["req"]),
+        finished_req_ids=set(),
+        num_running_reqs=1,
+        num_waiting_reqs=0,
+    )
+    result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output).get_request_output("req").result
+
+    assert result.media.prepared_for_transport is True
+    assert result.media.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_execute_stepwise_streaming_decodes_final_only_pipeline(monkeypatch):
@@ -819,6 +885,7 @@ def _make_batch_runner(pipeline):
         cache_backend="none",
         enable_cache_dit_summary=False,
         parallel_config=SimpleNamespace(use_hsdp=False),
+        video_output_transport=VideoOutputTransportConfig(),
     )
     runner.kv_transfer_manager = SimpleNamespace(
         receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None,
@@ -938,6 +1005,21 @@ def test_execute_model_batch_uses_runner_output_helper(monkeypatch):
     assert helper_calls == [(["req-0", "req-1"], ["a", "b"])]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_batch_prepares_typed_media_before_return(monkeypatch):
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(media=_raw_video_media())]))
+    runner.od_config.video_output_transport = VideoOutputTransportConfig(enable_device_postprocess=True)
+
+    result = DiffusionModelRunner.execute_model_batch(runner, _make_scheduler_output(num_reqs=1), runner.od_config)
+
+    prepared = result.runner_outputs[0].result.media
+    assert prepared.prepared_for_transport is True
+    assert prepared.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_split_diffusion_output_by_request_slices_single_and_multi_request_outputs():
@@ -959,6 +1041,36 @@ def test_split_diffusion_output_by_request_slices_single_and_multi_request_outpu
     assert single[0].output == ["img-0a", "img-0b"]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_split_diffusion_output_by_request_preserves_typed_media():
+    reqs = [_make_request(), _make_request()]
+    batch = DiffusionRequestBatch(requests=reqs)
+    source = torch.randn(4, 3, 2, 4, 5)
+    result = DiffusionOutput(media=_raw_video_media(tensor=source))
+
+    outputs = split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=2)
+
+    assert [output.media.video.tensor.shape for output in outputs] == [
+        torch.Size([2, 3, 2, 4, 5]),
+        torch.Size([2, 3, 2, 4, 5]),
+    ]
+    assert all(output.media.prepared_for_transport is False for output in outputs)
+    torch.testing.assert_close(outputs[0].media.video.tensor, source[:2])
+    torch.testing.assert_close(outputs[1].media.video.tensor, source[2:])
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_split_diffusion_output_by_request_validates_media_batch_dimension():
+    reqs = [_make_request(), _make_request()]
+    batch = DiffusionRequestBatch(requests=reqs)
+    result = DiffusionOutput(media=_raw_video_media(batch_size=3))
+
+    with pytest.raises(ValueError, match=r"request_count \* num_outputs_per_prompt \(4\), got 3"):
+        split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=2)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_split_diffusion_output_by_request_slices_tuple_outputs():

tests/diffusion/test_ipc_async.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for IPC async D2H stream packing."""
 
+from unittest.mock import ANY
+
 import pytest
 import torch
 
@@ -43,7 +45,7 @@ class TestPackTensorIfLargeWithD2hStream:
         tensor = torch.zeros(1)
         d2h = mocker.MagicMock()
         _pack_value_if_large(tensor, d2h_stream=d2h)
-        mock_pack.assert_called_once_with(tensor, d2h_stream=d2h)
+        mock_pack.assert_called_once_with(tensor, d2h_stream=d2h, created=None)
 
     def test_pack_value_if_large_passes_stream_through_dict(self, mocker):
         """d2h_stream propagates through dict container packing."""
@@ -93,7 +95,7 @@ class TestPackDiffusionFieldsWithD2hStream:
         d2h = mocker.MagicMock()
         out = DiffusionOutput(output=torch.zeros(1))
         _pack_diffusion_fields(out, d2h_stream=d2h)
-        mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h)
+        mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h, created=None)
 
     def test_trajectory_fields_packed_with_stream(self, mocker):
         mock_pack_tensor = mocker.patch(
@@ -120,7 +122,12 @@ class TestPackDiffusionOutputShmWithD2hStream:
         d2h = mocker.MagicMock()
         out = DiffusionOutput()
         pack_diffusion_output_shm(out, d2h_stream=d2h)
-        mock_pack_fields.assert_called_once_with(out, d2h_stream=d2h)
+        mock_pack_fields.assert_called_once_with(
+            out,
+            d2h_stream=d2h,
+            created=ANY,
+            pending_updates=ANY,
+        )
 
     def test_rpc_envelope_path_with_stream(self, mocker):
         mock_pack_fields = mocker.patch(
@@ -130,11 +137,11 @@ class TestPackDiffusionOutputShmWithD2hStream:
         result = DiffusionOutput()
         envelope = {"type": "diffusion_rpc_result", "result": result}
         pack_diffusion_output_shm(envelope, d2h_stream=d2h)
-        # envelope["result"] is replaced in-place by pack_diffusion_output_shm;
-        # assert with the original reference, not the replaced one.
         mock_pack_fields.assert_called_once_with(
             result,
             d2h_stream=d2h,
+            created=ANY,
+            pending_updates=ANY,
         )
 
     @pytest.mark.parametrize("wrapped_in_rpc_envelope", [False, True])
@@ -147,7 +154,12 @@ class TestPackDiffusionOutputShmWithD2hStream:
         tagged = {"dp_rank": 1, "output": result}
         value = {"type": "diffusion_rpc_result", "result": tagged} if wrapped_in_rpc_envelope else tagged
         pack_diffusion_output_shm(value, d2h_stream=d2h)
-        mock_pack_fields.assert_called_once_with(result, d2h_stream=d2h)
+        mock_pack_fields.assert_called_once_with(
+            result,
+            d2h_stream=d2h,
+            created=ANY,
+            pending_updates=ANY,
+        )
 
     def test_runner_output_path_with_stream(self, mocker):
         mock_pack_fields = mocker.patch(
@@ -163,9 +175,9 @@ class TestPackDiffusionOutputShmWithD2hStream:
             result=result,
         )
         pack_diffusion_output_shm(runner_out, d2h_stream=d2h)
-        # runner_out.result is replaced in-place by pack_diffusion_output_shm;
-        # assert with the original reference, not the replaced one.
         mock_pack_fields.assert_called_once_with(
             result,
             d2h_stream=d2h,
+            created=ANY,
+            pending_updates=ANY,
         )

tests/diffusion/test_media_output_contract.py

@@ -0,0 +1,546 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
+import vllm_omni.diffusion.ipc as diffusion_ipc
+import vllm_omni.diffusion.postprocess.media as media_postprocess
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    FloatVideoConsumer,
+    VideoMediaOutput,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTensorSpec,
+    VideoTransportConstraints,
+    VideoValueRange,
+)
+from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport
+from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _video(
+    tensor: torch.Tensor,
+    *,
+    encoding: VideoTensorEncoding,
+    value_range: VideoValueRange,
+    consumers: frozenset[FloatVideoConsumer] = frozenset(),
+) -> VideoMediaOutput:
+    layout = VideoTensorLayout.BTHWC if encoding is VideoTensorEncoding.UINT8_FRAMES else VideoTensorLayout.BCTHW
+    return VideoMediaOutput(
+        tensor=tensor,
+        spec=VideoTensorSpec(
+            layout=layout,
+            encoding=encoding,
+            value_range=value_range,
+        ),
+        constraints=VideoTransportConstraints(pending_float_consumers=consumers),
+    )
+
+
+def _prepared_media(video: VideoMediaOutput) -> DiffusionMediaOutput:
+    return DiffusionMediaOutput(video=video, prepared_for_transport=True)
+
+
+def _od_config(*, enabled: bool = False) -> OmniDiffusionConfig:
+    config = OmniDiffusionConfig(
+        model=None,
+        video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled),
+    )
+    config.model_class_name = "WanPipeline"
+    return config
+
+
+def test_finalizer_rejects_media_that_skipped_runner_preparation() -> None:
+    media = DiffusionMediaOutput(
+        video=_video(
+            torch.randn(1, 3, 2, 4, 5),
+            encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+            value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+        )
+    )
+
+    with pytest.raises(ValueError, match="before transport preparation"):
+        finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np"))
+
+
+def test_uint8_media_finalization_preserves_frames() -> None:
+    frames = torch.randint(0, 256, (1, 4, 8, 8, 3), dtype=torch.uint8)
+
+    output = finalize_diffusion_media(
+        _prepared_media(
+            _video(
+                frames,
+                encoding=VideoTensorEncoding.UINT8_FRAMES,
+                value_range=VideoValueRange.ZERO_TO_255,
+            )
+        ),
+        sampling_params=OmniDiffusionSamplingParams(output_type="np"),
+    )
+
+    np.testing.assert_array_equal(output["payload"]["video"], frames.numpy())
+    assert output["metadata"] == {}
+
+
+def test_float_media_finalization_uses_the_declared_range() -> None:
+    video = torch.tensor([[[[[-1.0, 0.0, 1.0]]]]]).expand(1, 3, 1, 1, 3).contiguous()
+
+    output = finalize_diffusion_media(
+        _prepared_media(
+            _video(
+                video,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            )
+        ),
+        sampling_params=OmniDiffusionSamplingParams(output_type="np"),
+    )
+
+    expected = (video.permute(0, 2, 3, 4, 1).numpy() * 0.5 + 0.5).clip(0.0, 1.0)
+    np.testing.assert_array_equal(output["payload"]["video"], expected)
+
+
+def test_zero_to_one_float_media_is_not_denormalized() -> None:
+    video = torch.tensor([[[[[0.0, 0.5, 1.0]]]]]).expand(1, 3, 1, 1, 3).contiguous()
+    media = _prepared_media(
+        _video(
+            video,
+            encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+            value_range=VideoValueRange.ZERO_TO_ONE,
+        )
+    )
+
+    output = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np"))
+
+    expected = video.permute(0, 2, 3, 4, 1).numpy()
+    np.testing.assert_array_equal(output["payload"]["video"], expected)
+
+
+def test_presentation_is_resolved_outside_the_tensor_descriptor() -> None:
+    video = torch.zeros(1, 3, 2, 4, 5)
+    media = _prepared_media(
+        _video(
+            video,
+            encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+            value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+        )
+    )
+
+    output = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="pil"))
+
+    assert "presentation" not in media.video.__dataclass_fields__
+    assert len(output["payload"]["video"]) == 1
+    assert len(output["payload"]["video"][0]) == 2
+
+
+def test_frame_interpolation_consumes_its_constraint(monkeypatch: pytest.MonkeyPatch) -> None:
+    video = torch.randn(1, 3, 2, 4, 5)
+    interpolated = torch.randn(1, 3, 3, 4, 5)
+    monkeypatch.setattr(
+        media_postprocess,
+        "interpolate_video_tensor",
+        lambda *args, **kwargs: (interpolated, 2.0),
+    )
+
+    output = finalize_diffusion_media(
+        _prepared_media(
+            _video(
+                video,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+                consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}),
+            )
+        ),
+        sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True),
+    )
+
+    assert output["payload"]["video"].shape == (1, 3, 4, 5, 3)
+    assert output["metadata"] == {"video": {"video_fps_multiplier": 2.0}}
+
+
+def test_uint8_media_rejects_a_pending_float_consumer() -> None:
+    media = _prepared_media(
+        _video(
+            torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8),
+            encoding=VideoTensorEncoding.UINT8_FRAMES,
+            value_range=VideoValueRange.ZERO_TO_255,
+            consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}),
+        )
+    )
+
+    with pytest.raises(ValueError, match="cannot have pending float consumers"):
+        media.validate()
+
+
+def test_unconsumed_float_constraint_fails_before_public_formatting() -> None:
+    media = _prepared_media(
+        _video(
+            torch.randn(1, 3, 2, 4, 5),
+            encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+            value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            consumers=frozenset({FloatVideoConsumer.VIDEO_GUARDRAILS}),
+        )
+    )
+
+    with pytest.raises(ValueError, match="pending float consumers"):
+        finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np"))
+
+
+def test_engine_routes_typed_media_exclusively(monkeypatch: pytest.MonkeyPatch) -> None:
+    media = _prepared_media(
+        _video(
+            torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8),
+            encoding=VideoTensorEncoding.UINT8_FRAMES,
+            value_range=VideoValueRange.ZERO_TO_255,
+        )
+    )
+    request = OmniDiffusionRequest(
+        prompt="test",
+        sampling_params=OmniDiffusionSamplingParams(output_type="np"),
+        request_id="request-0",
+    )
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = _od_config()
+    engine.post_process_func = lambda _: pytest.fail("legacy postprocessor was called")
+    engine._post_process_accepts_sampling_params = False
+
+    finalized = {"payload": {"video": np.zeros((1, 1, 2, 2, 3), dtype=np.uint8)}, "metadata": {}}
+
+    def _format_outputs(
+        *,
+        request: object,
+        od_config: object,
+        diffusion_output: object,
+        output_data: object,
+        postprocess_output: object,
+    ) -> list[object]:
+        del request, od_config, diffusion_output, output_data
+        return [postprocess_output]
+
+    monkeypatch.setattr(diffusion_engine_module, "finalize_diffusion_media", lambda *args, **kwargs: finalized)
+    monkeypatch.setattr(diffusion_engine_module, "normalize_diffusion_postprocess_output", lambda value: value)
+    monkeypatch.setattr(diffusion_engine_module, "format_diffusion_outputs", _format_outputs)
+
+    result = engine.postprocess_output(request, DiffusionOutput(media=media))
+
+    assert result == [finalized]
+
+
+def test_typed_media_crosses_runner_ipc_and_engine_without_legacy_postprocess(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    sampling_params = OmniDiffusionSamplingParams(output_type="np", num_outputs_per_prompt=1)
+    raw_media = DiffusionMediaOutput(
+        video=_video(
+            torch.linspace(-1.0, 1.0, 1 * 3 * 2 * 4 * 5).reshape(1, 3, 2, 4, 5),
+            encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+            value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+        )
+    )
+    prepared = prepare_diffusion_media_for_transport(
+        raw_media,
+        od_config=_od_config(enabled=True),
+        sampling_params=sampling_params,
+    )
+    diffusion_output = DiffusionOutput(media=prepared)
+    monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1)
+    pack_diffusion_output_shm(diffusion_output)
+    unpack_diffusion_output_shm(diffusion_output)
+
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = _od_config()
+    engine.post_process_func = lambda _: pytest.fail("legacy postprocessor was called")
+    engine._post_process_accepts_sampling_params = False
+    request = OmniDiffusionRequest(prompt="test", sampling_params=sampling_params, request_id="request-0")
+
+    outputs = engine.postprocess_output(request, diffusion_output)
+
+    assert len(outputs) == 1
+    assert outputs[0].images[0].dtype == np.uint8
+    assert outputs[0].images[0].shape == (1, 2, 4, 5, 3)
+
+
+def test_diffusion_output_rejects_typed_and_legacy_output_together() -> None:
+    media = _prepared_media(
+        _video(
+            torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8),
+            encoding=VideoTensorEncoding.UINT8_FRAMES,
+            value_range=VideoValueRange.ZERO_TO_255,
+        )
+    )
+    with pytest.raises(ValueError, match="both media and legacy output"):
+        DiffusionOutput(output=torch.zeros(1), media=media)
+
+
+def test_diffusion_output_rejects_postprocessor_on_typed_media() -> None:
+    media = _prepared_media(
+        _video(
+            torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8),
+            encoding=VideoTensorEncoding.UINT8_FRAMES,
+            value_range=VideoValueRange.ZERO_TO_255,
+        )
+    )
+
+    with pytest.raises(ValueError, match="model-specific post_process_func"):
+        DiffusionOutput(media=media, post_process_func=lambda value: value)
+
+
+def test_diffusion_output_preserves_the_legacy_positional_constructor() -> None:
+    # ``media`` was originally declared right after ``output``, which broke
+    # ``DiffusionOutput(output, trajectory_timesteps)`` by binding the second
+    # argument to media. It must stay last so out-of-tree positional callers keep
+    # the pre-existing field order.
+    output_tensor = torch.randn(1)
+    timesteps = torch.tensor([1, 2, 3])
+
+    output = DiffusionOutput(output_tensor, timesteps)
+
+    assert output.output is output_tensor
+    assert output.trajectory_timesteps is timesteps
+    assert output.media is None
+
+
+def test_frame_interpolation_uses_the_declared_range_negative_one_to_one(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # RIFE infers [0,1] vs [-1,1] from the tensor min/max, so an all-nonnegative
+    # [-1,1] clip is misread as [0,1] and interpolated in the wrong space. The
+    # interpolator must receive the unit-range form per the declared spec.
+    video = torch.full((1, 3, 2, 3, 3), 0.25)
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=video,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            ),
+            constraints=VideoTransportConstraints(
+                pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION})
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+
+    seen: list[torch.Tensor] = []
+
+    def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]:
+        seen.append(tensor)
+        return tensor * tensor, 2.0
+
+    monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator)
+
+    output = finalize_diffusion_media(
+        media,
+        sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True),
+    )
+
+    # The interpolator is fed the unit-range form 0.25→0.625, not raw 0.25.
+    assert len(seen) == 1
+    assert seen[0].shape == video.shape
+    assert float(seen[0][0, 0, 0, 0, 0]) == pytest.approx(0.625, abs=1e-6)
+    # Squaring gives 0.390625; restored to [-1,1] (-0.21875) by the fix, then the
+    # existing denormalize maps it back to [0,1] as 0.390625 in the public frames.
+    frames = output["payload"]["video"]
+    assert np.allclose(np.asarray(frames), 0.390625, atol=1e-6)
+
+
+def test_frame_interpolation_uses_the_declared_range_zero_to_one(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # With a declared [0,1] range, the all-0.25 clip is already in unit space, so
+    # squaring gives 0.390625 and there is no restore step.
+    video = torch.full((1, 3, 2, 3, 3), 0.25)
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=video,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.ZERO_TO_ONE,
+            ),
+            constraints=VideoTransportConstraints(
+                pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION})
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+
+    def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]:
+        return tensor * tensor, 2.0
+
+    monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator)
+
+    output = finalize_diffusion_media(
+        media,
+        sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True),
+    )
+
+    # 0.25 is already [0,1]; squaring gives 0.0625, and no denormalize is applied.
+    frames = output["payload"]["video"]
+    assert np.allclose(np.asarray(frames), 0.0625, atol=1e-6)
+
+
+def test_frame_interpolation_clamps_zero_to_one_overshoot(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # A declared [0,1] clip whose sample is almost entirely inside [0,1] — but
+    # with one bf16 value overshooting 1 — is still in unit space, so it must be
+    # clamped to [0,1] before RIFE and interpolated as [0,1]. Without the clamp
+    # the interpolator classifies it as [-1,1] (amax > 1) and returns the wrong
+    # space, which is then written back with no *2-1 restore.
+    video = torch.full((1, 3, 2, 3, 3), 0.25, dtype=torch.bfloat16)
+    video[0, 0, 0, 0, 0] = 1.25  # single overshoot > 1 -> amax > 1
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=video,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.ZERO_TO_ONE,
+            ),
+            constraints=VideoTransportConstraints(
+                pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION})
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+
+    seen: list[torch.Tensor] = []
+
+    def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]:
+        seen.append(tensor)
+        return tensor.float().square(), 2.0
+
+    monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator)
+
+    output = finalize_diffusion_media(
+        media,
+        sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True),
+    )
+
+    # The interpolator is fed the clamped bf16 unit-range tensor (max now 1.0),
+    # not the raw 1.25 overshoot that would be misread as [-1,1].
+    assert len(seen) == 1
+    assert seen[0].dtype is torch.bfloat16
+    assert float(seen[0].min()) == pytest.approx(0.25, abs=1e-6)
+    assert float(seen[0].max()) == pytest.approx(1.0, abs=1e-6)
+    # Squaring the clamped values gives 0.0625 at the 0.25 samples and 1.0 at the
+    # overshoot; the overall shape is unchanged (no *2-1 for ZERO_TO_ONE).
+    frames = output["payload"]["video"]
+    assert np.allclose(np.asarray(frames).reshape(-1)[1:], 0.0625, atol=1e-6)
+    assert np.isclose(float(np.asarray(frames).reshape(-1)[0]), 1.0, atol=1e-6)
+
+
+def test_packing_failure_unlinks_created_segments_and_leaves_output_untouched(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # Packing a late field can fail after an SHM segment is already created;
+    # the segment must be unlinked and ``output`` left unmutated so a partial
+    # object is never enqueued.
+    prepared = prepare_diffusion_media_for_transport(
+        DiffusionMediaOutput(
+            video=_video(
+                torch.linspace(-1.0, 1.0, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5),
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            )
+        ),
+        od_config=_od_config(enabled=False),
+        sampling_params=OmniDiffusionSamplingParams(output_type="np"),
+    )
+    output = DiffusionOutput(media=prepared, trajectory_latents=torch.randn(30))
+    original_media = prepared
+    original_latents = output.trajectory_latents
+
+    real_array_to_shm = diffusion_ipc._array_to_shm
+    created_handles: list[dict[str, object]] = []
+
+    def _fail_on_second(array: np.ndarray) -> dict[str, object]:
+        if created_handles:
+            raise RuntimeError("injected packing failure")
+        handle = real_array_to_shm(array)
+        created_handles.append(handle)
+        return handle
+
+    monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1)
+    monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second)
+
+    with pytest.raises(RuntimeError, match="injected packing failure"):
+        pack_diffusion_output_shm(output)
+
+    assert created_handles, "no SHM segment was created before the injected failure"
+    # Output was not mutated despite the media segment being created first.
+    assert output.media is original_media
+    assert isinstance(output.media, DiffusionMediaOutput)
+    assert output.trajectory_latents is original_latents
+
+    from multiprocessing import shared_memory
+
+    with pytest.raises(FileNotFoundError):
+        shared_memory.SharedMemory(name=str(created_handles[0]["name"]))
+
+
+def test_second_batch_entry_failure_leaves_the_whole_batch_untouched(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from multiprocessing import shared_memory
+
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    first_media = _prepared_media(
+        _video(
+            torch.zeros(1, 3, 2, 4, 5),
+            encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+            value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+        )
+    )
+    second_media = _prepared_media(
+        _video(
+            torch.ones(1, 3, 2, 4, 5),
+            encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+            value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+        )
+    )
+    first = DiffusionOutput(media=first_media)
+    second = DiffusionOutput(media=second_media)
+    batch = BatchRunnerOutput(
+        runner_outputs=[
+            RunnerOutput(request_id="first", result=first),
+            RunnerOutput(request_id="second", result=second),
+        ]
+    )
+
+    real_array_to_shm = diffusion_ipc._array_to_shm
+    created_handles: list[dict[str, object]] = []
+
+    def _fail_on_second(array: np.ndarray) -> dict[str, object]:
+        if created_handles:
+            raise RuntimeError("second batch entry failed")
+        handle = real_array_to_shm(array)
+        created_handles.append(handle)
+        return handle
+
+    monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1)
+    monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second)
+
+    with pytest.raises(RuntimeError, match="second batch entry failed"):
+        pack_diffusion_output_shm(batch)
+
+    # No per-entry assignment is committed until every entry has packed.
+    assert first.media is first_media
+    assert second.media is second_media
+    with pytest.raises(FileNotFoundError):
+        shared_memory.SharedMemory(name=str(created_handles[0]["name"]))

tests/diffusion/test_video_transport_defaults_are_inert.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Pin the default behavior and the WAN reference integration boundary."""
+
+from __future__ import annotations
+
+import inspect
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig, VideoOutputTransportConfig
+from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_transport_defaults() -> None:
+    transport = VideoOutputTransportConfig()
+
+    assert transport.enable_device_postprocess is False
+
+
+def test_a_default_diffusion_config_carries_a_default_transport() -> None:
+    config = OmniDiffusionConfig(model=None)
+
+    assert isinstance(config.video_output_transport, VideoOutputTransportConfig)
+    assert config.video_output_transport.enable_device_postprocess is False
+
+
+def test_wan_uses_the_typed_media_contract() -> None:
+    import vllm_omni.diffusion.models as models_pkg
+
+    source = (Path(inspect.getfile(models_pkg)).parent / "wan2_2/pipeline_wan2_2.py").read_text()
+    assert "DiffusionMediaOutput(" in source
+    assert "reduce_video_to_uint8_frames(" not in source
+    # The migrated built-in WAN pipeline emits typed media, so it never invokes
+    # the legacy postprocess hook. The mappings are intentionally retained so
+    # legacy/out-of-tree WAN replacements (registered with
+    # post_process_func_name=None) keep the built-in hook instead of exposing raw
+    # BCTHW tensors.
+    assert _DIFFUSION_POST_PROCESS_FUNCS["WanPipeline"] == "get_wan22_post_process_func"
+    assert _DIFFUSION_POST_PROCESS_FUNCS["WanDMDPipeline"] == "get_wan22_post_process_func"
+    assert _DIFFUSION_POST_PROCESS_FUNCS["WanT2VDMD2Pipeline"] == "get_wan22_post_process_func"
+
+
+def test_non_reference_models_do_not_use_the_runtime_reducer() -> None:
+    import vllm_omni.diffusion.models as models_pkg
+
+    root = Path(inspect.getfile(models_pkg)).parent
+    paths = (
+        "cosmos3/pipeline_cosmos3.py",
+        "hunyuan_video/pipeline_hunyuan_video_1_5.py",
+        "lingbot_video/pipeline_lingbot_video.py",
+        "ltx2/ltx2_runtime.py",
+        "minimax_h3/pipeline_minimax_h3.py",
+    )
+    for relative in paths:
+        source = (root / relative).read_text()
+        assert "prepare_diffusion_media_for_transport(" not in source
+        assert "reduce_video_to_uint8_frames(" not in source

tests/diffusion/test_wan_device_reduction.py

@@ -0,0 +1,348 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""WAN reference coverage for the typed pre-D2H media contract."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig
+from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    VideoMediaOutput,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTensorSpec,
+    VideoValueRange,
+)
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import get_wan22_post_process_func
+from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport
+from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media
+from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_uint8_frames
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.diffusion,
+    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+]
+
+_GPU = torch.accelerator.is_available() if hasattr(torch, "accelerator") else torch.cuda.is_available()
+requires_gpu = pytest.mark.skipif(not _GPU, reason="device reduction is a GPU path")
+
+
+def _decoded_wan_video(tensor: torch.Tensor) -> DiffusionMediaOutput:
+    return DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=tensor,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            ),
+        )
+    )
+
+
+def _config(*, enabled: bool) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model=None,
+        video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled),
+    )
+
+
+def _encoder_frames(video_payload: np.ndarray) -> list[np.ndarray]:
+    return [_coerce_video_to_uint8_frames(video_payload[i]) for i in range(video_payload.shape[0])]
+
+
+@requires_gpu
+def test_wan_media_reduction_matches_float32_reference_and_bounds_native_path() -> None:
+    torch.manual_seed(0)
+    vae_out = torch.rand(2, 3, 8, 64, 96, device="cuda:0", dtype=torch.bfloat16) * 2.4 - 1.2
+    sampling = OmniDiffusionSamplingParams(output_type=None, enable_frame_interpolation=False)
+
+    legacy_postprocess = get_wan22_post_process_func(_config(enabled=False))
+    widened_out = legacy_postprocess(vae_out.float(), output_type="np", sampling_params=sampling)
+    widened_frames = _encoder_frames(widened_out["payload"]["video"])
+    native_out = legacy_postprocess(vae_out, output_type="np", sampling_params=sampling)
+    native_frames = _encoder_frames(native_out["payload"]["video"])
+
+    prepared = prepare_diffusion_media_for_transport(
+        _decoded_wan_video(vae_out),
+        od_config=_config(enabled=True),
+        sampling_params=sampling,
+    )
+    reduced_out = finalize_diffusion_media(prepared, sampling_params=sampling)
+    assert prepared.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES
+    assert reduced_out["payload"]["video"].dtype == np.uint8
+    reduced_frames = _encoder_frames(reduced_out["payload"]["video"])
+
+    assert len(widened_frames) == len(reduced_frames) == 2
+    for expected, native, produced in zip(widened_frames, native_frames, reduced_frames, strict=True):
+        assert produced.shape == expected.shape == (8, 64, 96, 3)
+        np.testing.assert_array_equal(produced, expected)
+        deviation = np.abs(produced.astype(np.int16) - native.astype(np.int16))
+        assert deviation.max() <= 1
+
+
+@requires_gpu
+def test_wan_uint8_media_finalization_does_not_rescale() -> None:
+    frames = torch.randint(0, 256, (1, 4, 16, 16, 3), dtype=torch.uint8, device="cuda:0")
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=frames,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BTHWC,
+                encoding=VideoTensorEncoding.UINT8_FRAMES,
+                value_range=VideoValueRange.ZERO_TO_255,
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+
+    out = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np"))
+
+    np.testing.assert_array_equal(out["payload"]["video"], frames.cpu().numpy())
+
+
+@requires_gpu
+def test_diffusion_output_to_cpu_moves_wan_media_off_device() -> None:
+    prepared = prepare_diffusion_media_for_transport(
+        _decoded_wan_video(torch.randn(1, 3, 2, 4, 5, device="cuda:0")),
+        od_config=_config(enabled=False),
+        sampling_params=OmniDiffusionSamplingParams(output_type="np"),
+    )
+    output = DiffusionOutput(media=prepared, to_cpu=True)
+
+    assert output.media is not None
+    assert output.media.video.tensor.device.type == "cpu"
+
+
+@requires_gpu
+def test_small_typed_media_is_on_cpu_before_ipc() -> None:
+    sampling = OmniDiffusionSamplingParams(output_type="np")
+    prepared = prepare_diffusion_media_for_transport(
+        _decoded_wan_video(torch.randn(1, 3, 2, 4, 5, device="cuda:0")),
+        od_config=_config(enabled=True),
+        sampling_params=sampling,
+    )
+    output = DiffusionOutput(media=prepared)
+
+    pack_diffusion_output_shm(output)
+
+    assert isinstance(output.media, dict)
+    packed_tensor = output.media["video"]["tensor"]
+    assert isinstance(packed_tensor, torch.Tensor)
+    assert packed_tensor.device.type == "cpu"
+    unpack_diffusion_output_shm(output)
+
+
+@requires_gpu
+def test_wan_float_media_finalization_matches_the_legacy_path() -> None:
+    video = torch.rand(1, 3, 4, 16, 16, device="cuda:0", dtype=torch.bfloat16) * 2 - 1
+    sampling = OmniDiffusionSamplingParams(output_type="np")
+    prepared = prepare_diffusion_media_for_transport(
+        _decoded_wan_video(video),
+        od_config=_config(enabled=False),
+        sampling_params=sampling,
+    )
+
+    produced = finalize_diffusion_media(prepared, sampling_params=sampling)
+    expected = get_wan22_post_process_func(_config(enabled=False))(
+        video,
+        output_type="np",
+        sampling_params=sampling,
+    )
+
+    assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+    np.testing.assert_array_equal(produced["payload"]["video"], expected["payload"]["video"])
+
+
+@requires_gpu
+def test_oom_fallback_media_packs_without_device_widening(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Regression: when the fp32 reduction OOMs, the fallback keeps the bf16 float
+    # payload. Async IPC must copy it to the host in bf16 and widen on the CPU;
+    # a second device-side fp32 conversion would OOM again and drop the payload.
+    from vllm_omni.diffusion.postprocess import device_reduction
+
+    # Large enough (> 1 MB in bf16) that packing routes the tensor through the
+    # shared-memory D2H path rather than the small-tensor CPU copy.
+    video = torch.rand(1, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16) * 2 - 1
+
+    def _raise_oom(*_args: object, **_kwargs: object) -> torch.Tensor:
+        raise torch.OutOfMemoryError("forced fp32 reduction OOM")
+
+    monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", _raise_oom)
+
+    prepared = prepare_diffusion_media_for_transport(
+        _decoded_wan_video(video),
+        od_config=_config(enabled=True),
+        sampling_params=OmniDiffusionSamplingParams(output_type="np"),
+    )
+    assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+    assert prepared.video.tensor.dtype == torch.bfloat16
+
+    output = DiffusionOutput(media=prepared)
+
+    pinned_dtypes: list[torch.dtype] = []
+    real_empty = torch.empty
+
+    def _spy_empty(*args: object, **kwargs: object) -> torch.Tensor:
+        tensor = real_empty(*args, **kwargs)
+        if kwargs.get("pin_memory"):
+            pinned_dtypes.append(tensor.dtype)
+        return tensor
+
+    monkeypatch.setattr(torch, "empty", _spy_empty)
+
+    d2h_stream = torch.Stream(device=torch.device("cuda", 0))
+    pack_diffusion_output_shm(output, d2h_stream=d2h_stream)
+    d2h_stream.synchronize()
+
+    assert pinned_dtypes == [torch.bfloat16], (
+        "async packing must D2H in the original dtype and widen on the CPU, "
+        "never widen bf16 to fp32 on the accelerator before the copy"
+    )
+
+    unpack_diffusion_output_shm(output)
+    assert isinstance(output.media, DiffusionMediaOutput)
+    assert output.media.video.tensor.dtype == torch.bfloat16
+    torch.testing.assert_close(output.media.video.tensor, video.cpu())
+
+
+@requires_gpu
+def test_oom_fallback_recovers_the_allocator(monkeypatch: pytest.MonkeyPatch) -> None:
+    # A real float32 OOM leaves the CUDA caching allocator in a failed state;
+    # the fallback must empty the cache and drain the device before marking the
+    # float representation transport-ready, so the worker D2H of the same tensor
+    # can actually run. Use a genuine allocation failure (per-process cap) rather
+    # than injecting the exception, which never poisons the allocator.
+    from vllm_omni.platforms import current_omni_platform
+
+    video = torch.empty(1, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16)
+    video.uniform_(-1.0, 1.0)
+    expected = video.cpu()
+    d2h_stream = torch.Stream(device=video.device)
+    base_bytes = video.numel() * video.element_size()
+    _free, total = current_omni_platform.get_device_memory()
+
+    # Cap the process so the fp32 conversion (2x base bytes) cannot fit, while
+    # retaining the original bf16 tensor. This forces a real allocator OOM at
+    # the .to(fp32), rather than injecting an exception from a monkeypatch.
+    torch.cuda.set_per_process_memory_fraction(min(0.99, (1.2 * base_bytes) / total))
+    real_empty_cache = current_omni_platform.empty_cache
+    real_synchronize = current_omni_platform.synchronize
+    try:
+        empty_calls: list[int] = []
+        sync_calls: list[bool] = []
+
+        def _empty_cache() -> None:
+            real_empty_cache()
+            empty_calls.append(torch.accelerator.current_device_index())
+
+        def _synchronize() -> None:
+            real_synchronize()
+            sync_calls.append(True)
+
+        monkeypatch.setattr(current_omni_platform, "empty_cache", _empty_cache)
+        monkeypatch.setattr(current_omni_platform, "synchronize", _synchronize)
+
+        prepared = prepare_diffusion_media_for_transport(
+            _decoded_wan_video(video),
+            od_config=_config(enabled=True),
+            sampling_params=OmniDiffusionSamplingParams(output_type="np"),
+        )
+
+        # The fallback ran past a genuine OOM and kept the bf16 float payload.
+        assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+        assert prepared.video.tensor.dtype == torch.bfloat16
+        assert empty_calls, "fallback must empty the CUDA allocator cache after a real OOM"
+        assert sync_calls, "fallback must drain the device after a real OOM"
+
+        # Exercise the next production step: async D2H/SHM packing of the same
+        # tensor must succeed after allocator recovery and preserve its values.
+        output = DiffusionOutput(media=prepared)
+        pack_diffusion_output_shm(output, d2h_stream=d2h_stream)
+        d2h_stream.synchronize()
+        unpack_diffusion_output_shm(output)
+        assert isinstance(output.media, DiffusionMediaOutput)
+        torch.testing.assert_close(output.media.video.tensor, expected)
+    finally:
+        torch.cuda.set_per_process_memory_fraction(1.0)
+        real_empty_cache()
+        real_synchronize()
+
+
+@requires_gpu
+def test_oom_fallback_releases_late_uint8_traceback_before_clone(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A final uint8 allocation OOM must release float intermediates first."""
+    from vllm_omni.diffusion.postprocess import device_reduction
+    from vllm_omni.platforms import current_omni_platform
+
+    # A split request is a non-compact view, so float fallback must allocate a
+    # request-owned clone after reduction fails.
+    backing = torch.empty(2, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16)
+    backing.uniform_(-1.0, 1.0)
+    video = backing[1:2].detach()
+    expected = video.cpu()
+    _, total = current_omni_platform.get_device_memory()
+    real_to = torch.Tensor.to
+    real_empty_cache = current_omni_platform.empty_cache
+    real_synchronize = current_omni_platform.synchronize
+    real_ensure_request_owned = device_reduction.ensure_request_owned_tensor
+    uint8_attempts: list[int] = []
+    fallback_allocations: list[int] = []
+
+    def _record_fallback_allocation(tensor: torch.Tensor) -> torch.Tensor:
+        fallback_allocations.append(torch.accelerator.memory_allocated(tensor.device))
+        return real_ensure_request_owned(tensor)
+
+    def _fail_real_uint8_allocation(
+        tensor: torch.Tensor,
+        *args: object,
+        **kwargs: object,
+    ) -> torch.Tensor:
+        dtype = kwargs.get("dtype")
+        if dtype is None and args and isinstance(args[0], torch.dtype):
+            dtype = args[0]
+        if tensor.is_cuda and dtype is torch.uint8:
+            # Earlier float32 work runs without a cap. Immediately before the
+            # real final uint8 allocation, cap the allocator below that request
+            # but above the post-unwind request-owned bf16 clone requirement.
+            real_empty_cache()
+            allocated = torch.accelerator.memory_allocated(tensor.device)
+            uint8_bytes = tensor.numel()
+            limit_bytes = allocated + max(1, uint8_bytes // 2)
+            torch.cuda.set_per_process_memory_fraction(min(0.99, limit_bytes / total))
+            uint8_attempts.append(allocated)
+        return real_to(tensor, *args, **kwargs)
+
+    monkeypatch.setattr(device_reduction, "ensure_request_owned_tensor", _record_fallback_allocation)
+    monkeypatch.setattr(torch.Tensor, "to", _fail_real_uint8_allocation)
+    try:
+        prepared = prepare_diffusion_media_for_transport(
+            _decoded_wan_video(video),
+            od_config=_config(enabled=True),
+            sampling_params=OmniDiffusionSamplingParams(output_type="np"),
+        )
+
+        assert uint8_attempts, "the real final .to(uint8) allocation was not attempted"
+        assert fallback_allocations, "float fallback did not request compact storage"
+        assert fallback_allocations[0] < uint8_attempts[0], (
+            "the failed reduction's float intermediates remained live during fallback"
+        )
+        assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
+        assert prepared.video.tensor.dtype is torch.bfloat16
+        assert prepared.video.tensor.storage_offset() == 0
+        assert prepared.video.tensor.untyped_storage().data_ptr() != video.untyped_storage().data_ptr()
+        torch.testing.assert_close(prepared.video.tensor.cpu(), expected)
+    finally:
+        torch.cuda.set_per_process_memory_fraction(1.0)
+        real_empty_cache()
+        real_synchronize()

tests/docs/test_generate_argparse.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+ROOT_DIR = Path(__file__).parents[2]
+HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_argparse.py"
+SPEC = importlib.util.spec_from_file_location("generate_argparse", HOOK_PATH)
+assert SPEC and SPEC.loader
+generate_argparse = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(generate_argparse)
+
+
+def test_static_serve_parser_supports_video_output_transport() -> None:
+    parser = generate_argparse.create_parser_subparser_init(generate_argparse.OmniServeCommand)
+
+    args = parser.parse_args(["--omni", "--video-output-transport", '{"enable_device_postprocess": true}'])
+
+    assert args.video_output_transport == {"enable_device_postprocess": True}

tests/e2e/offline_inference/test_device_postprocess_equivalence.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""End-to-end WAN check for the typed pre-D2H media contract.
+
+Runs the whole engine twice with the same seed, once per flag state, and compares
+the frames a client receives. WAN denormalizes historical bfloat16 output in the
+input dtype while device preparation computes in float32, so the paths may differ
+by one 255th.
+
+Run with a WAN2.2 checkpoint::
+
+    VLLM_OMNI_DEVICE_POSTPROCESS_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+        pytest -s tests/e2e/offline_inference/test_device_postprocess_equivalence.py
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_MODEL")
+
+pytestmark = [
+    pytest.mark.full_model,
+    pytest.mark.diffusion,
+    pytest.mark.gpu,
+    pytest.mark.skipif(
+        not MODEL,
+        reason="Set VLLM_OMNI_DEVICE_POSTPROCESS_MODEL to a WAN2.2 checkpoint to run this.",
+    ),
+]
+
+_SEED = 1234
+_STEPS = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_STEPS", "2"))
+_SIZE = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_SIZE", "256"))
+_FRAMES = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_FRAMES", "17"))
+
+
+def _generate(enable_device_postprocess: bool) -> np.ndarray:
+    from vllm_omni.entrypoints.omni import Omni
+
+    engine = Omni(
+        model=MODEL,
+        num_gpus=1,
+        video_output_transport={"enable_device_postprocess": enable_device_postprocess},
+    )
+    try:
+        # output_type must travel with the request: the pipelines read it from the
+        # sampling params, and the reduction only covers the "np" path.
+        sampling_params = OmniDiffusionSamplingParams(
+            output_type="np",
+            seed=_SEED,
+            num_inference_steps=_STEPS,
+            height=_SIZE,
+            width=_SIZE,
+            num_frames=_FRAMES,
+        )
+        outputs = engine.generate({"prompt": "a robot waving"}, sampling_params)
+    finally:
+        engine.close()
+
+    video = outputs[0].images[0]
+    return video.numpy() if hasattr(video, "numpy") else np.asarray(video)
+
+
+def _as_uint8_frames(video: np.ndarray) -> np.ndarray:
+    """Reproduce what the API server does before encoding.
+
+    A float payload is scaled and rounded there, so that is the only fair place to
+    compare it against a payload the worker already reduced.
+    """
+    if video.dtype == np.uint8:
+        return video
+    return np.rint(np.clip(video, 0.0, 1.0) * 255.0).astype(np.uint8)
+
+
+# A dtype difference moves ~20% of pixels by one; dropping the round moves ~50%.
+_MAX_DIFFERING_FRACTION = 0.35
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_device_postprocess_produces_equivalent_frames() -> None:
+    float_path = _generate(False)
+    device_path = _generate(True)
+
+    # Guard against the run silently taking the float path twice, which would make
+    # the comparison meaningless.
+    assert float_path.dtype != np.uint8, "the float path returned uint8; the gate may be stuck open"
+    assert device_path.dtype == np.uint8, "the WAN typed-media path did not prepare uint8 frames"
+    assert float_path.shape == device_path.shape
+
+    expected = _as_uint8_frames(float_path)
+    deviation = np.abs(device_path.astype(np.int16) - expected.astype(np.int16))
+    differing = float((deviation > 0).mean())
+    print(f"maxdiff={deviation.max()} differing={differing:.4f} ({int((deviation > 0).sum())}/{deviation.size})")
+
+    assert deviation.max() <= 1, "the reduction moved a pixel by more than one 255th"
+    assert differing < _MAX_DIFFERING_FRACTION, (
+        f"{differing:.1%} of pixels differ, more than a denormalize-precision difference explains"
+    )

tests/entrypoints/test_serve.py

@@ -69,6 +69,30 @@ def test_serve_parser_accepts_deploy_config() -> None:
     assert args.get_explicit_kwargs_dict()["deploy_config"] == "/tmp/deploy.yaml"
 
 
+def test_serve_parser_accepts_video_output_transport() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        ["serve", "fake-model", "--omni", "--video-output-transport", '{"enable_device_postprocess": true}']
+    )
+
+    expected = {"enable_device_postprocess": True}
+    assert args.video_output_transport == expected
+    assert args.get_explicit_kwargs_dict()["video_output_transport"] == expected
+
+
+@pytest.mark.parametrize("value", ["{not json", "[]"])
+def test_serve_parser_rejects_invalid_video_output_transport(value: str) -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    with pytest.raises(SystemExit):
+        parser.parse_args(["serve", "fake-model", "--omni", "--video-output-transport", value])
+
+
 def test_serve_parser_rejects_stage_configs_path() -> None:
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="subcommand")

vllm_omni/config/omni_config.py

@@ -718,6 +718,7 @@ class _DiffusionConfigProjection:
     cache_strategy: str = "none"
     cache_backend: str = "none"
     cache_config: Any = field(default_factory=dict)
+    video_output_transport: object = field(default_factory=dict)
     enable_cache_dit_summary: bool = False
     diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
     diffusion_kv_max_rows_per_request: int | None = Field(default=None, ge=1, strict=True)
@@ -812,6 +813,7 @@ class _DiffusionConfigProjection:
             AttentionConfig,
             DiffusionCacheConfig,
             TransformerConfig,
+            VideoOutputTransportConfig,
             build_attention_config,
             parse_kv_cache_skip_selector,
             validate_dlo_host_registration_options,
@@ -854,6 +856,13 @@ class _DiffusionConfigProjection:
         elif not isinstance(self.cache_config, DiffusionCacheConfig):
             self.cache_config = DiffusionCacheConfig()
 
+        if self.video_output_transport is None:
+            self.video_output_transport = VideoOutputTransportConfig()
+        elif isinstance(self.video_output_transport, Mapping):
+            self.video_output_transport = VideoOutputTransportConfig(**dict(self.video_output_transport))
+        elif not isinstance(self.video_output_transport, VideoOutputTransportConfig):
+            raise TypeError("video_output_transport must be a VideoOutputTransportConfig or mapping")
+
         self._propagate_quantization_from_tf_config(self.tf_model_config)
         if self.quantization_config is not None:
             if isinstance(self.quantization_config, QuantizationConfig):

vllm_omni/config/stage_config.py

@@ -458,6 +458,7 @@ class StageDeployConfig:
     fa_deterministic: bool | None = None
     cache_backend: str | None = None
     cache_config: dict[str, Any] | None = None
+    video_output_transport: dict[str, Any] | None = None
     enable_cache_dit_summary: bool | None = None
     step_execution: bool | None = None
     vae_use_slicing: bool | None = None

vllm_omni/diffusion/data.py

@@ -29,6 +29,7 @@ from vllm_omni.diffusion.diffusion_kv.config import (
     parse_diffusion_kv_cache_mode,
 )
 from vllm_omni.diffusion.lora.manager import LoRABackend
+from vllm_omni.diffusion.media import DiffusionMediaOutput
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 from vllm_omni.errors import client_error_metadata
@@ -758,6 +759,15 @@ def uses_diffusers_adapter(od_config: object) -> bool:
     )
 
 
+@dataclass
+class VideoOutputTransportConfig:
+    enable_device_postprocess: bool = False
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.enable_device_postprocess, bool):
+            raise TypeError("enable_device_postprocess must be a bool")
+
+
 @dataclass
 class OmniDiffusionConfig:
     # Model and path configuration (for convenience)
@@ -793,6 +803,7 @@ class OmniDiffusionConfig:
     # Cache backend configuration (NEW)
     cache_backend: str = "none"  # "tea_cache", "deep_cache", etc.
     cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
+    video_output_transport: VideoOutputTransportConfig = field(default_factory=VideoOutputTransportConfig)
     enable_cache_dit_summary: bool = False
 
     # Prompt-embedding cache. When enabled, ``DiffusionModelRunner`` wraps the
@@ -1258,6 +1269,13 @@ class OmniDiffusionConfig:
             # If it's neither dict nor DiffusionCacheConfig, convert to empty config
             self.cache_config = DiffusionCacheConfig()
 
+        if self.video_output_transport is None:
+            self.video_output_transport = VideoOutputTransportConfig()
+        elif isinstance(self.video_output_transport, Mapping):
+            self.video_output_transport = VideoOutputTransportConfig(**dict(self.video_output_transport))
+        elif not isinstance(self.video_output_transport, VideoOutputTransportConfig):
+            raise TypeError("video_output_transport must be a VideoOutputTransportConfig or mapping")
+
         # Auto-detect quantization from TransformerConfig if not explicitly set.
         # This covers the case where tf_model_config is passed at construction
         # time. For late (post-construction) assignment, callers should use
@@ -1618,8 +1636,8 @@ class DiffusionOutput:
 
     output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None
 
-    # Legacy compatibility fields. New pipeline-specific payloads should be
-    # carried by output["payload"] instead.
+    # Legacy compatibility fields. Decoded video uses ``media``; other new
+    # payloads should be carried by output["payload"] instead.
     trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None
     trajectory_latents: torch.Tensor | dict[str, Any] | None = None
     trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None
@@ -1658,7 +1676,18 @@ class DiffusionOutput:
     # mode) and the receiving side must not initialise a stray CUDA context.
     to_cpu: bool = False
 
+    # Typed video-media contract. Declared last so the pre-existing positional
+    # constructor order (output, trajectory_timesteps, ...) that out-of-tree
+    # pipelines rely on is preserved. Mutually exclusive with ``output``.
+    media: DiffusionMediaOutput | None = None
+
     def __post_init__(self) -> None:
+        if self.media is not None and not isinstance(self.media, DiffusionMediaOutput):
+            raise TypeError(f"media must be DiffusionMediaOutput, got {type(self.media).__name__}")
+        if self.media is not None and self.output is not None:
+            raise ValueError("DiffusionOutput cannot contain both media and legacy output")
+        if self.media is not None and self.post_process_func is not None:
+            raise ValueError("Typed diffusion media cannot carry a model-specific post_process_func")
         if not self.to_cpu:
             return
 
@@ -1674,6 +1703,10 @@ class DiffusionOutput:
             return value
 
         self.output = _maybe_to_cpu(self.output)
+        if self.media is not None:
+            if not self.media.prepared_for_transport:
+                raise ValueError("Diffusion media must be prepared before to_cpu=True")
+            self.media = self.media.to_cpu()
         self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps)
         self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents)
         self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs)

vllm_omni/diffusion/diffusion_engine.py

@@ -47,6 +47,7 @@ from vllm_omni.diffusion.output_formatter import (
     format_empty_diffusion_outputs,
     normalize_diffusion_postprocess_output,
 )
+from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
     get_diffusion_post_process_func,
@@ -534,6 +535,13 @@ class DiffusionEngine:
             raise RuntimeError(output.error)
         logger.debug("Generation completed successfully.")
 
+        if output.media is not None:
+            if output.output is not None:
+                raise ValueError("DiffusionOutput cannot contain both media and legacy output")
+            media = output.media.to_cpu() if self.od_config.enable_cpu_offload else output.media
+            output_data = media.video.tensor
+            outputs = finalize_diffusion_media(media, sampling_params=request.sampling_params)
+        else:
             if output.output is None:
                 logger.warning("Output is None, returning empty OmniRequestOutput")
                 return format_empty_diffusion_outputs(request, finished=output.finished)

vllm_omni/diffusion/ipc.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """IPC utilities for transferring large tensors via POSIX shared memory.
 
@@ -11,15 +11,32 @@ serialised through the queue.
 
 from __future__ import annotations
 
+from collections.abc import Iterator
 from typing import Any
 
 import numpy as np
 import torch
 
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    FloatVideoConsumer,
+    VideoColorModel,
+    VideoMediaOutput,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTensorSpec,
+    VideoTransportConstraints,
+    VideoValueRange,
+)
 
 _SHM_TENSOR_THRESHOLD = 1_000_000  # 1 MB
 DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"
+_DIFFUSION_MEDIA_WIRE_TYPE = "diffusion_media_v1"
+
+# Sentinel so compute-then-assign packing can tell "field not packed" apart from
+# "field packed to None".
+_UNSET = object()
 
 
 def _array_to_shm(array: np.ndarray) -> dict[str, Any]:
@@ -61,9 +78,39 @@ def _array_from_shm(handle: dict[str, Any]) -> np.ndarray:
     return array
 
 
+def _unlink_shm_handle(handle: dict[str, Any]) -> None:
+    """Best-effort removal of a shared-memory segment created during packing."""
+    from multiprocessing import shared_memory
+
+    name = handle.get("name")
+    if not name:
+        return
+    try:
+        shm = shared_memory.SharedMemory(name=name)
+    except FileNotFoundError:
+        return
+    try:
+        shm.close()
+    finally:
+        try:
+            shm.unlink()
+        except FileNotFoundError:
+            pass
+
+
+def _unlink_shm_handles(handles: list[dict[str, Any]]) -> None:
+    for handle in handles:
+        try:
+            _unlink_shm_handle(handle)
+        except Exception:
+            # Cleanup is best-effort; never mask the original packing failure.
+            pass
+
+
 def _tensor_to_shm(
     tensor: torch.Tensor,
     d2h_stream: torch.Stream | None = None,
+    created: list[dict[str, Any]] | None = None,
 ) -> dict[str, Any]:
     """Copy a tensor into POSIX shared memory and return a metadata handle.
 
@@ -82,14 +129,17 @@ def _tensor_to_shm(
         torch.accelerator.set_stream(d2h_stream)
         try:
             t = tensor.detach()
-            if original_dtype == torch.bfloat16:
-                t = t.to(torch.float32)
             cpu = torch.empty(t.shape, dtype=t.dtype, pin_memory=True)
             cpu.copy_(t, non_blocking=True)
         finally:
             torch.accelerator.set_stream(old_stream)
         d2h_stream.synchronize()
         tensor = cpu
+        # numpy has no bfloat16, so widen on the host after D2H. Never widen on
+        # the accelerator before the copy: a payload that already fell back after
+        # an fp32 OOM would OOM again on a second device-side fp32 allocation.
+        if original_dtype == torch.bfloat16:
+            tensor = tensor.to(torch.float32)
     else:
         tensor = tensor.detach().cpu().contiguous()
         if original_dtype == torch.bfloat16:
@@ -101,6 +151,8 @@ def _tensor_to_shm(
             "torch_dtype": str(original_dtype),
         }
     )
+    if created is not None:
+        created.append(handle)
     return handle
 
 
@@ -120,6 +172,7 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
 def _pack_tensor_if_large(
     val: torch.Tensor,
     d2h_stream: torch.Stream | None = None,
+    created: list[dict[str, Any]] | None = None,
 ) -> torch.Tensor | dict:
     """Replace a tensor with an SHM handle if it exceeds the threshold.
 
@@ -134,14 +187,16 @@ def _pack_tensor_if_large(
     except Exception:
         storage_bytes = view_bytes
     if max(view_bytes, storage_bytes) > _SHM_TENSOR_THRESHOLD:
-        return _tensor_to_shm(val, d2h_stream=d2h_stream)
+        return _tensor_to_shm(val, d2h_stream=d2h_stream, created=created)
     return val
 
 
-def _ndarray_to_shm(array: np.ndarray) -> dict[str, Any]:
+def _ndarray_to_shm(array: np.ndarray, created: list[dict[str, Any]] | None = None) -> dict[str, Any]:
     """Copy a contiguous NumPy array into POSIX shared memory."""
     handle = _array_to_shm(array)
     handle["__ndarray_shm__"] = True
+    if created is not None:
+        created.append(handle)
     return handle
 
 
@@ -153,6 +208,7 @@ def _ndarray_from_shm(handle: dict[str, Any]) -> np.ndarray:
 def _pack_value_if_large(
     val: object,
     d2h_stream: torch.Stream | None = None,
+    created: list[dict[str, Any]] | None = None,
 ) -> object:
     """Recursively replace large tensors with SHM handles.
 
@@ -163,15 +219,19 @@ def _pack_value_if_large(
     the two in sync.
     """
     if isinstance(val, torch.Tensor):
-        return _pack_tensor_if_large(val, d2h_stream=d2h_stream)
+        return _pack_tensor_if_large(val, d2h_stream=d2h_stream, created=created)
     if isinstance(val, np.ndarray):
-        return _ndarray_to_shm(val) if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD else val
+        return (
+            _ndarray_to_shm(val, created=created)
+            if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD
+            else val
+        )
     if isinstance(val, dict):
-        return {key: _pack_value_if_large(value, d2h_stream=d2h_stream) for key, value in val.items()}
+        return {key: _pack_value_if_large(value, d2h_stream=d2h_stream, created=created) for key, value in val.items()}
     if isinstance(val, list):
-        return [_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val]
+        return [_pack_value_if_large(item, d2h_stream=d2h_stream, created=created) for item in val]
     if isinstance(val, tuple):
-        return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val)
+        return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream, created=created) for item in val)
     return val
 
 
@@ -190,18 +250,116 @@ def _unpack_if_shm_handle(val: object) -> object:
     return val
 
 
+def _pack_diffusion_media(
+    media: DiffusionMediaOutput,
+    *,
+    d2h_stream: torch.Stream | None,
+    created: list[dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    media.validate()
+    if not media.prepared_for_transport:
+        raise ValueError("Diffusion media must be prepared before IPC packing")
+    video = media.video
+    packed_tensor = _pack_tensor_if_large(video.tensor, d2h_stream=d2h_stream, created=created)
+    if isinstance(packed_tensor, torch.Tensor) and packed_tensor.device.type != "cpu":
+        packed_tensor = packed_tensor.detach().cpu().contiguous()
+    return {
+        "__type__": _DIFFUSION_MEDIA_WIRE_TYPE,
+        "schema_version": video.schema_version,
+        "prepared_for_transport": True,
+        "video": {
+            "tensor": packed_tensor,
+            "layout": video.spec.layout.value,
+            "encoding": video.spec.encoding.value,
+            "value_range": video.spec.value_range.value,
+            "color_model": video.spec.color_model.value,
+            "pending_float_consumers": sorted(consumer.value for consumer in video.constraints.pending_float_consumers),
+        },
+    }
+
+
+def _unpack_diffusion_media(packed: dict[str, Any]) -> DiffusionMediaOutput:
+    if packed.get("__type__") != _DIFFUSION_MEDIA_WIRE_TYPE:
+        raise ValueError("Invalid diffusion media wire payload")
+    video_payload = packed.get("video")
+    if not isinstance(video_payload, dict):
+        raise ValueError("Diffusion media wire payload is missing video metadata")
+    tensor = _unpack_if_shm_handle(video_payload.get("tensor"))
+    if not isinstance(tensor, torch.Tensor):
+        raise ValueError("Diffusion media wire payload contains an invalid video tensor")
+    schema_version = packed.get("schema_version")
+    if schema_version != 1:
+        raise ValueError(f"Unsupported video media schema version: {schema_version}")
+    if packed.get("prepared_for_transport") is not True:
+        raise ValueError("Diffusion media wire payload was not prepared for transport")
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=tensor,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout(video_payload["layout"]),
+                encoding=VideoTensorEncoding(video_payload["encoding"]),
+                value_range=VideoValueRange(video_payload["value_range"]),
+                color_model=VideoColorModel(video_payload["color_model"]),
+            ),
+            constraints=VideoTransportConstraints(
+                pending_float_consumers=frozenset(
+                    FloatVideoConsumer(value) for value in video_payload.get("pending_float_consumers", [])
+                )
+            ),
+            schema_version=1,
+        ),
+        prepared_for_transport=True,
+    )
+    media.validate()
+    return media
+
+
 def _pack_diffusion_fields(
     output: DiffusionOutput,
     d2h_stream: torch.Stream | None = None,
+    created: list[dict[str, Any]] | None = None,
+    pending_updates: list[tuple[DiffusionOutput, str, object]] | None = None,
 ) -> DiffusionOutput:
+    if output.media is not None and output.output is not None:
+        raise ValueError("DiffusionOutput cannot contain both media and legacy output")
+    if output.media is not None and not isinstance(output.media, DiffusionMediaOutput):
+        raise TypeError(f"DiffusionOutput.media must be DiffusionMediaOutput, got {type(output.media).__name__}")
+
+    # Pack into locals first so a mid-payload failure leaves ``output`` unmutated
+    # (never enqueue a half-packed object); the caller unlinks any SHM segments
+    # accumulated in ``created`` on failure.
+    packed_output = _UNSET
     if output.output is not None:
-        output.output = _pack_value_if_large(output.output, d2h_stream=d2h_stream)
-    if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor):
-        output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream)
-    if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor):
-        output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream)
-    if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor):
-        output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream)
+        packed_output = _pack_value_if_large(output.output, d2h_stream=d2h_stream, created=created)
+    packed_media = _UNSET
+    if output.media is not None:
+        packed_media = _pack_diffusion_media(output.media, d2h_stream=d2h_stream, created=created)
+    packed_latents = _UNSET
+    if isinstance(output.trajectory_latents, torch.Tensor):
+        packed_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream, created=created)
+    packed_timesteps = _UNSET
+    if isinstance(output.trajectory_timesteps, torch.Tensor):
+        packed_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream, created=created)
+    packed_log_probs = _UNSET
+    if isinstance(output.trajectory_log_probs, torch.Tensor):
+        packed_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream, created=created)
+
+    updates = [
+        (output, field_name, value)
+        for field_name, value in (
+            ("output", packed_output),
+            ("media", packed_media),
+            ("trajectory_latents", packed_latents),
+            ("trajectory_timesteps", packed_timesteps),
+            ("trajectory_log_probs", packed_log_probs),
+        )
+        if value is not _UNSET
+    ]
+    if pending_updates is None:
+        for target, field_name, value in updates:
+            setattr(target, field_name, value)
+    else:
+        pending_updates.extend(updates)
     return output
 
 
@@ -223,35 +381,84 @@ def pack_diffusion_output_shm(
 
     If *d2h_stream* is provided, D2H copies use that stream (non-blocking on
     the default stream).  The caller must synchronize *d2h_stream* afterward.
+
+    Packing is failure-atomic: every SHM segment created during the call is
+    tracked and unlinked if any part of the payload fails to pack, and field
+    mutations are committed only after the complete payload succeeds.
     """
+    created: list[dict[str, Any]] = []
+    pending_updates: list[tuple[DiffusionOutput, str, object]] = []
+    try:
+        packed = _pack_output_shm(
+            output,
+            d2h_stream=d2h_stream,
+            created=created,
+            pending_updates=pending_updates,
+        )
+    except BaseException:
+        _unlink_shm_handles(created)
+        raise
+
+    # Commit mutations only after every field and batch entry packed. Until
+    # this point, a failure leaves the complete caller-owned payload untouched.
+    for target, field_name, value in pending_updates:
+        setattr(target, field_name, value)
+    return packed
+
+
+def _iter_diffusion_outputs(output: object) -> Iterator[DiffusionOutput]:
+    """Yield every packable ``DiffusionOutput`` leaf in a worker payload."""
     if isinstance(output, DiffusionOutput):
-        return _pack_diffusion_fields(output, d2h_stream=d2h_stream)
+        yield output
+        return
 
     # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput}
     if isinstance(output, dict) and "dp_rank" in output and "output" in output:
         inner = output["output"]
         if isinstance(inner, DiffusionOutput):
-            output["output"] = _pack_diffusion_fields(inner, d2h_stream=d2h_stream)
-        return output
+            yield inner
+        return
 
     if _is_rpc_result_envelope(output):
-        result = output.get("result")
-        output["result"] = pack_diffusion_output_shm(result, d2h_stream=d2h_stream)
-        return output
+        yield from _iter_diffusion_outputs(output.get("result"))
+        return
 
     result = getattr(output, "result", None)
     if isinstance(result, DiffusionOutput):
-        output.result = _pack_diffusion_fields(result, d2h_stream=d2h_stream)
+        yield result
 
     runner_outputs = getattr(output, "runner_outputs", None)
     if isinstance(runner_outputs, list):
         for runner_output in runner_outputs:
-            pack_diffusion_output_shm(runner_output, d2h_stream=d2h_stream)
+            yield from _iter_diffusion_outputs(runner_output)
+
+
+def payload_carries_typed_media(output: object) -> bool:
+    """True if *output* carries any typed ``DiffusionOutput.media`` payload."""
+    return any(item.media is not None for item in _iter_diffusion_outputs(output))
+
+
+def _pack_output_shm(
+    output: object,
+    *,
+    d2h_stream: torch.Stream | None,
+    created: list[dict[str, Any]],
+    pending_updates: list[tuple[DiffusionOutput, str, object]],
+) -> object:
+    for item in _iter_diffusion_outputs(output):
+        _pack_diffusion_fields(
+            item,
+            d2h_stream=d2h_stream,
+            created=created,
+            pending_updates=pending_updates,
+        )
     return output
 
 
 def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
     output.output = _unpack_if_shm_handle(output.output)
+    if isinstance(output.media, dict):
+        output.media = _unpack_diffusion_media(output.media)
     output.trajectory_latents = _unpack_if_shm_handle(output.trajectory_latents)
     output.trajectory_timesteps = _unpack_if_shm_handle(output.trajectory_timesteps)
     output.trajectory_log_probs = _unpack_if_shm_handle(output.trajectory_log_probs)

vllm_omni/diffusion/media.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field, replace
+from enum import Enum
+from typing import Literal
+
+import torch
+
+
+class VideoTensorLayout(str, Enum):
+    BCTHW = "bcthw"
+    BTHWC = "bthwc"
+
+
+class VideoTensorEncoding(str, Enum):
+    NORMALIZED_FLOAT = "normalized_float"
+    UINT8_FRAMES = "uint8_frames"
+
+
+class VideoValueRange(str, Enum):
+    NEGATIVE_ONE_TO_ONE = "negative_one_to_one"
+    ZERO_TO_ONE = "zero_to_one"
+    ZERO_TO_255 = "zero_to_255"
+
+
+class VideoColorModel(str, Enum):
+    RGB = "rgb"
+
+
+class FloatVideoConsumer(str, Enum):
+    FRAME_INTERPOLATION = "frame_interpolation"
+    VIDEO_GUARDRAILS = "video_guardrails"
+
+
+@dataclass(frozen=True, slots=True)
+class VideoTensorSpec:
+    layout: VideoTensorLayout
+    encoding: VideoTensorEncoding
+    value_range: VideoValueRange
+    color_model: VideoColorModel = VideoColorModel.RGB
+
+
+@dataclass(frozen=True, slots=True)
+class VideoTransportConstraints:
+    pending_float_consumers: frozenset[FloatVideoConsumer] = frozenset()
+
+
+@dataclass(frozen=True, slots=True)
+class VideoMediaOutput:
+    tensor: torch.Tensor
+    spec: VideoTensorSpec
+    constraints: VideoTransportConstraints = field(default_factory=VideoTransportConstraints)
+    schema_version: Literal[1] = 1
+
+    def validate(self) -> None:
+        if self.schema_version != 1:
+            raise ValueError(f"Unsupported video media schema version: {self.schema_version}")
+        if not isinstance(self.spec, VideoTensorSpec):
+            raise TypeError(f"Video media spec must be a VideoTensorSpec, got {type(self.spec).__name__}")
+        if not isinstance(self.constraints, VideoTransportConstraints):
+            raise TypeError(
+                f"Video media constraints must be VideoTransportConstraints, got {type(self.constraints).__name__}"
+            )
+        if not isinstance(self.constraints.pending_float_consumers, frozenset):
+            raise TypeError("pending_float_consumers must be a frozenset")
+        if not isinstance(self.tensor, torch.Tensor):
+            raise TypeError(f"Video media tensor must be a torch.Tensor, got {type(self.tensor).__name__}")
+        if self.tensor.ndim != 5:
+            raise ValueError(f"Video media tensor must be rank 5, got shape {tuple(self.tensor.shape)}")
+        if any(size <= 0 for size in self.tensor.shape):
+            raise ValueError(f"Video media dimensions must be positive, got shape {tuple(self.tensor.shape)}")
+        if self.spec.color_model is not VideoColorModel.RGB:
+            raise ValueError(f"Unsupported video color model: {self.spec.color_model}")
+        if any(not isinstance(consumer, FloatVideoConsumer) for consumer in self.constraints.pending_float_consumers):
+            raise ValueError("pending_float_consumers must contain only FloatVideoConsumer values")
+
+        if self.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT:
+            self._validate_normalized_float()
+            return
+        if self.spec.encoding is VideoTensorEncoding.UINT8_FRAMES:
+            self._validate_uint8_frames()
+            return
+        raise ValueError(f"Unsupported video tensor encoding: {self.spec.encoding}")
+
+    def _validate_normalized_float(self) -> None:
+        if self.spec.layout is not VideoTensorLayout.BCTHW:
+            raise ValueError("NORMALIZED_FLOAT video must use BCTHW layout")
+        if self.tensor.shape[1] != 3:
+            raise ValueError(f"RGB BCTHW video must have 3 channels, got {self.tensor.shape[1]}")
+        if self.tensor.dtype not in {torch.float16, torch.bfloat16, torch.float32}:
+            raise ValueError(f"NORMALIZED_FLOAT video must use fp16, bf16, or fp32, got {self.tensor.dtype}")
+        if self.spec.value_range not in {
+            VideoValueRange.NEGATIVE_ONE_TO_ONE,
+            VideoValueRange.ZERO_TO_ONE,
+        }:
+            raise ValueError(f"NORMALIZED_FLOAT video has incompatible value range: {self.spec.value_range}")
+
+    def _validate_uint8_frames(self) -> None:
+        if self.spec.layout is not VideoTensorLayout.BTHWC:
+            raise ValueError("UINT8_FRAMES video must use BTHWC layout")
+        if self.tensor.shape[-1] != 3:
+            raise ValueError(f"RGB BTHWC video must have 3 channels, got {self.tensor.shape[-1]}")
+        if self.tensor.dtype is not torch.uint8:
+            raise ValueError(f"UINT8_FRAMES video must use torch.uint8, got {self.tensor.dtype}")
+        if self.spec.value_range is not VideoValueRange.ZERO_TO_255:
+            raise ValueError(f"UINT8_FRAMES video has incompatible value range: {self.spec.value_range}")
+        if self.constraints.pending_float_consumers:
+            raise ValueError("UINT8_FRAMES video cannot have pending float consumers")
+        if not self.tensor.is_contiguous():
+            raise ValueError("UINT8_FRAMES video tensor must be contiguous")
+
+    def with_tensor(self, tensor: torch.Tensor) -> VideoMediaOutput:
+        return replace(self, tensor=tensor)
+
+    def to_cpu(self) -> VideoMediaOutput:
+        return self.with_tensor(self.tensor.detach().cpu())
+
+
+@dataclass(frozen=True, slots=True)
+class DiffusionMediaOutput:
+    video: VideoMediaOutput
+    # Only the model runner may set this before worker output packing.
+    prepared_for_transport: bool = False
+
+    def validate(self) -> None:
+        if not isinstance(self.video, VideoMediaOutput):
+            raise TypeError(f"Diffusion media video must be a VideoMediaOutput, got {type(self.video).__name__}")
+        if not isinstance(self.prepared_for_transport, bool):
+            raise TypeError("prepared_for_transport must be a bool")
+        self.video.validate()
+        if not self.prepared_for_transport:
+            return
+        if not self.video.tensor.is_contiguous():
+            raise ValueError("Prepared video media tensor must be contiguous")
+        if not _owns_compact_storage(self.video.tensor):
+            raise ValueError("Prepared video media tensor must own request-local storage")
+
+    def with_video(self, video: VideoMediaOutput) -> DiffusionMediaOutput:
+        return replace(self, video=video)
+
+    def to_cpu(self) -> DiffusionMediaOutput:
+        self.validate()
+        moved = self.with_video(self.video.to_cpu())
+        moved.validate()
+        return moved
+
+
+def _owns_compact_storage(tensor: torch.Tensor) -> bool:
+    """True if *tensor* is contiguous and backed by exactly its own elements.
+
+    ``tensor._base is None`` is not sufficient: ``batch[i : i + 1].detach()`` is
+    contiguous with ``_base=None`` yet keeps a nonzero storage offset and the
+    whole batch storage, so serializing it would leak every other request's
+    pixels. Require a zero offset and a storage sized to the logical elements.
+    """
+    if not tensor.is_contiguous():
+        return False
+    if tensor.storage_offset() != 0:
+        return False
+    try:
+        storage_bytes = tensor.untyped_storage().nbytes()
+    except Exception:
+        return False
+    return storage_bytes == tensor.numel() * tensor.element_size()
+
+
+def ensure_request_owned_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    if _owns_compact_storage(tensor):
+        return tensor
+    # A contiguous view may still retain another request's full batch storage.
+    return tensor.clone(memory_format=torch.contiguous_format)
+
+
+def slice_diffusion_media_output(
+    media: DiffusionMediaOutput,
+    start: int,
+    stop: int,
+) -> DiffusionMediaOutput:
+    media.validate()
+    source = media.video.tensor
+    tensor = source if start == 0 and stop == source.shape[0] else source[start:stop]
+    sliced = replace(media, video=media.video.with_tensor(tensor), prepared_for_transport=False)
+    sliced.validate()
+    return sliced

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -26,6 +26,14 @@ from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, Pipe
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
 from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    VideoMediaOutput,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTensorSpec,
+    VideoValueRange,
+)
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -907,6 +915,7 @@ class Wan22Pipeline(
 
         if DEBUG_PERF:
             _t_decode_start = time.perf_counter()
+        media = None
         if output_type == "latent":
             output = latents
         else:
@@ -920,7 +929,27 @@ class Wan22Pipeline(
                 latents.device, latents.dtype
             )
             latents = latents / latents_std + latents_mean
-            output = self.vae.decode(latents, return_dict=False)[0]
+            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()
@@ -947,6 +976,7 @@ class Wan22Pipeline(
         return split_diffusion_output_by_request(
             DiffusionOutput(
                 output=output,
+                media=media,
                 stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
             ),
             req,

vllm_omni/diffusion/postprocess/device_reduction.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Device-side reduction of decoded video tensors to uint8 frames."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    FloatVideoConsumer,
+    VideoMediaOutput,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTensorSpec,
+    VideoTransportConstraints,
+    VideoValueRange,
+    ensure_request_owned_tensor,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+# VaeImageProcessor.denormalize: (x * 0.5 + 0.5).clamp(0, 1)
+_DENORM_SCALE = 0.5
+_DENORM_SHIFT = 0.5
+
+logger = init_logger(__name__)
+
+
+def _request_float_consumers(
+    sampling_params: OmniDiffusionSamplingParams | None,
+) -> frozenset[FloatVideoConsumer]:
+    if sampling_params is not None and sampling_params.enable_frame_interpolation:
+        return frozenset({FloatVideoConsumer.FRAME_INTERPOLATION})
+    return frozenset()
+
+
+def _request_supports_uint8_frames(sampling_params: OmniDiffusionSamplingParams | None) -> bool:
+    return sampling_params is None or (sampling_params.output_type or "np") == "np"
+
+
+def _prepare_float_media_for_transport(
+    media: DiffusionMediaOutput,
+    video: VideoMediaOutput,
+) -> DiffusionMediaOutput:
+    prepared = replace(
+        media,
+        video=video.with_tensor(ensure_request_owned_tensor(video.tensor)),
+        prepared_for_transport=True,
+    )
+    prepared.validate()
+    return prepared
+
+
+def prepare_diffusion_media_for_transport(
+    media: DiffusionMediaOutput,
+    *,
+    od_config: OmniDiffusionConfig,
+    sampling_params: OmniDiffusionSamplingParams | None = None,
+) -> DiffusionMediaOutput:
+    """Validate and prepare one request-local media output before D2H."""
+    media.validate()
+    if media.prepared_for_transport:
+        # Only the model runner prepares media, exactly once, on media the
+        # pipeline emits unprepared. Accepting pre-prepared media here would let
+        # a pipeline forge runner-owned state and skip request policy (e.g. a
+        # forged uint8 payload silently bypasses frame interpolation).
+        raise ValueError(
+            "Media reached transport preparation already prepared; pipelines must "
+            "emit unprepared media so request policy is applied by the runner"
+        )
+    if media.video.spec.encoding is not VideoTensorEncoding.NORMALIZED_FLOAT:
+        raise ValueError("Pipelines must emit unprepared media as NORMALIZED_FLOAT video")
+
+    video = media.video
+    constraints = VideoTransportConstraints(
+        pending_float_consumers=(video.constraints.pending_float_consumers | _request_float_consumers(sampling_params))
+    )
+    constrained_video = replace(video, constraints=constraints)
+
+    enabled = od_config.video_output_transport.enable_device_postprocess
+    output_supported = _request_supports_uint8_frames(sampling_params)
+    if not enabled or constraints.pending_float_consumers or not output_supported:
+        if not enabled:
+            reason = "disabled"
+        elif constraints.pending_float_consumers:
+            reason = ",".join(sorted(consumer.value for consumer in constraints.pending_float_consumers))
+        else:
+            reason = "unsupported_presentation"
+        logger.debug("Device video preparation kept normalized float: reason=%s", reason)
+        return _prepare_float_media_for_transport(media, constrained_video)
+
+    try:
+        frames = reduce_video_to_uint8_frames(
+            constrained_video.tensor,
+            do_denormalize=constrained_video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE,
+        )
+    except torch.OutOfMemoryError:
+        logger.warning("Device video preparation ran out of memory; using normalized float transport")
+    else:
+        prepared_video = replace(
+            constrained_video,
+            tensor=frames,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BTHWC,
+                encoding=VideoTensorEncoding.UINT8_FRAMES,
+                value_range=VideoValueRange.ZERO_TO_255,
+            ),
+            constraints=VideoTransportConstraints(),
+        )
+        prepared = replace(media, video=prepared_video, prepared_for_transport=True)
+        prepared.validate()
+        logger.debug(
+            "Device video preparation converted normalized float to uint8: shape=%s",
+            tuple(frames.shape),
+        )
+        return prepared
+
+    # Leave the exception handler before allocator cleanup and fallback
+    # construction. A late conversion OOM otherwise keeps the helper traceback
+    # (and its live float32 intermediates) reachable while the fallback may need
+    # to clone a non-compact request view.
+    current_omni_platform.empty_cache()
+    current_omni_platform.synchronize()
+    return _prepare_float_media_for_transport(media, constrained_video)
+
+
+def reduce_video_to_uint8_frames(video: torch.Tensor, *, do_denormalize: bool = True) -> torch.Tensor:
+    """Reduce a decoded ``[B, C, F, H, W]`` video to uint8 ``[B, F, H, W, C]`` frames.
+
+    Runs denormalize/clamp/permute/round on the input's device so the following
+    D2H copy carries uint8 instead of float. The result matches
+    ``VideoProcessor.postprocess_video(output_type="np")`` then the ``*255``
+    rounding done in the API server. Pass ``do_denormalize=False`` for VAEs that
+    already emit ``[0, 1]``.
+    """
+    if video.dim() != 5:
+        raise ValueError(f"expected a [B, C, F, H, W] video tensor, got shape {tuple(video.shape)}")
+
+    # Match the numpy path, which promotes to float before scaling.
+    frames = video.to(torch.float32)
+    if do_denormalize:
+        frames = frames.mul(_DENORM_SCALE).add(_DENORM_SHIFT).clamp_(0.0, 1.0)
+    else:
+        frames = frames.clamp(0.0, 1.0)
+    frames = frames.permute(0, 2, 3, 4, 1)
+    frames = frames.mul_(255.0).round_().clamp_(0.0, 255.0).to(torch.uint8)
+    return frames.contiguous()

vllm_omni/diffusion/postprocess/media.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import replace
+
+from vllm_omni.diffusion.media import (
+    DiffusionMediaOutput,
+    FloatVideoConsumer,
+    VideoTensorEncoding,
+    VideoTensorLayout,
+    VideoTransportConstraints,
+    VideoValueRange,
+)
+from vllm_omni.diffusion.postprocess.rife_interpolator import interpolate_video_tensor
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def finalize_diffusion_media(
+    media: DiffusionMediaOutput,
+    *,
+    sampling_params: OmniDiffusionSamplingParams,
+) -> dict[str, object]:
+    """Finalize typed video media without invoking a model postprocessor."""
+    media.validate()
+    if not media.prepared_for_transport:
+        raise ValueError("Diffusion media reached the engine before transport preparation")
+
+    video = media.video
+    metadata: dict[str, object] = {}
+    consumers = video.constraints.pending_float_consumers
+    if FloatVideoConsumer.FRAME_INTERPOLATION in consumers:
+        if video.spec.encoding is not VideoTensorEncoding.NORMALIZED_FLOAT:
+            raise ValueError("Frame interpolation requires NORMALIZED_FLOAT video")
+        # RIFE guesses [0,1] vs [-1,1] from the tensor's min/max, so relying on
+        # the sample content misclassifies both an all-nonnegative [-1,1] clip
+        # (read as [0,1]) and a [0,1] clip whose single bf16 value overshoots 1
+        # (read as [-1,1]). Always map to unit range per the declared spec and
+        # clamp to [0,1] before RIFE, then restore the declared range so
+        # downstream denormalization is correct regardless of the sample.
+        to_negative_one = video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE
+        if to_negative_one:
+            interp_input = video.tensor.mul(0.5).add(0.5).clamp_(0.0, 1.0)
+        else:
+            interp_input = video.tensor.clamp(0.0, 1.0)
+        interpolated, multiplier = interpolate_video_tensor(
+            interp_input,
+            exp=sampling_params.frame_interpolation_exp,
+            scale=sampling_params.frame_interpolation_scale,
+            model_path=sampling_params.frame_interpolation_model_path,
+        )
+        if to_negative_one:
+            interpolated = interpolated.mul(2.0).sub(1.0)
+        consumers = consumers - {FloatVideoConsumer.FRAME_INTERPOLATION}
+        video = replace(
+            video,
+            tensor=interpolated,
+            constraints=VideoTransportConstraints(pending_float_consumers=frozenset(consumers)),
+        )
+        metadata["video"] = {"video_fps_multiplier": multiplier}
+
+    if consumers:
+        names = sorted(consumer.value for consumer in consumers)
+        raise ValueError(f"Video has pending float consumers at finalization: {names}")
+
+    output_type = getattr(sampling_params, "output_type", None) or "np"
+    if output_type == "latent":
+        raise ValueError("Latent output cannot be finalized from VideoMediaOutput")
+
+    if video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES:
+        if video.spec.layout is not VideoTensorLayout.BTHWC:
+            raise ValueError("UINT8_FRAMES video must use BTHWC layout")
+        if output_type != "np":
+            raise ValueError(f"UINT8_FRAMES finalization does not support output_type={output_type!r}")
+        frames = video.tensor.detach().cpu().numpy()
+    else:
+        if video.spec.layout is not VideoTensorLayout.BCTHW:
+            raise ValueError("NORMALIZED_FLOAT video must use BCTHW layout")
+        from diffusers.video_processor import VideoProcessor
+
+        do_denormalize = video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE
+        frames = VideoProcessor(vae_scale_factor=8).postprocess_video(
+            video.tensor,
+            output_type=output_type,
+            do_denormalize=[do_denormalize] * video.tensor.shape[2],
+        )
+
+    return {
+        "payload": {"video": frames},
+        "metadata": metadata,
+    }

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -53,6 +53,7 @@ from vllm_omni.diffusion.models.interface import (
 )
 from vllm_omni.diffusion.offloader import enable_offload_backend
 from vllm_omni.diffusion.offloader.config import TEXT_ENCODER_COMPONENT, resolve_offload
+from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
@@ -680,6 +681,21 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 od_config.cache_backend,
             )
 
+    def _prepare_output_for_transport(
+        self,
+        output: DiffusionOutput,
+        sampling_params: OmniDiffusionSamplingParams,
+    ) -> DiffusionOutput:
+        if output.media is not None:
+            if output.output is not None:
+                raise ValueError("DiffusionOutput cannot contain both media and legacy output")
+            output.media = prepare_diffusion_media_for_transport(
+                output.media,
+                od_config=self.od_config,
+                sampling_params=sampling_params,
+            )
+        return output
+
     def _runner_output_from_outputs(
         self,
         reqs: list[OmniDiffusionRequest],
@@ -774,6 +790,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                         allow_single_output=allow_single_output,
                         pipeline_name=type(self.pipeline).__name__,
                     )
+                with record_function("prepare_output_for_transport"):
+                    outputs = [
+                        self._prepare_output_for_transport(output, req.sampling_params)
+                        for req, output in zip(reqs, outputs, strict=True)
+                    ]
 
             if is_primary and outputs and record_output_peak_memory:
                 batch_peak_memory_mb = self._sample_peak_memory_mb()
@@ -1289,6 +1310,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                 clear_pipeline_stage_durations(self.pipeline)
                                 result = self.pipeline.post_decode(req)
                                 if result is not None:
+                                    result = self._prepare_output_for_transport(result, req.sampling)
                                     self._attach_stepwise_metadata(
                                         req,
                                         result,

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -64,7 +64,11 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     model_parallel_is_initialized,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
+from vllm_omni.diffusion.ipc import (
+    DIFFUSION_RPC_RESULT_ENVELOPE,
+    pack_diffusion_output_shm,
+    payload_carries_typed_media,
+)
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend
 from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -1131,7 +1135,15 @@ class WorkerProc:
         # Sync path (original, or async fallback).
         try:
             pack_diffusion_output_shm(output)
+        except (TypeError, ValueError):
+            raise
         except Exception as e:
+            # Typed media that fails to pack is left unprepared/on device (the
+            # pack is failure-atomic and does not mutate it), so enqueueing it
+            # would ship a broken payload. Re-raise memory/packing failures here
+            # instead of swallowing them for typed media.
+            if payload_carries_typed_media(output):
+                raise
             if hasattr(output, "output"):
                 logger.warning("SHM pack failed for model output: %s", e)
         self._enqueue_result(output)

vllm_omni/diffusion/worker/request_batch.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Request-level batch abstraction for diffusion runner."""
 
 from __future__ import annotations
@@ -12,6 +12,7 @@ import numpy as np
 import torch
 
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.media import DiffusionMediaOutput, slice_diffusion_media_output
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
@@ -35,6 +36,18 @@ def split_diffusion_output_by_request(
     """Split a batched DiffusionOutput into one output per request."""
     if num_outputs_per_prompt <= 0:
         raise ValueError(f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}.")
+    if result.media is not None and result.output is not None:
+        raise ValueError("DiffusionOutput cannot contain both media and legacy output")
+    if result.media is not None and not isinstance(result.media, DiffusionMediaOutput):
+        raise TypeError(f"DiffusionOutput.media must be DiffusionMediaOutput, got {type(result.media).__name__}")
+    if result.media is not None:
+        expected_batch = req.num_reqs * num_outputs_per_prompt
+        actual_batch = result.media.video.tensor.shape[0]
+        if actual_batch != expected_batch:
+            raise ValueError(
+                f"Video media batch dimension must equal request_count * num_outputs_per_prompt "
+                f"({expected_batch}), got {actual_batch}"
+            )
 
     return [
         DiffusionOutput(
@@ -43,6 +56,15 @@ def split_diffusion_output_by_request(
                 idx * num_outputs_per_prompt,
                 (idx + 1) * num_outputs_per_prompt,
             ),
+            media=(
+                slice_diffusion_media_output(
+                    result.media,
+                    idx * num_outputs_per_prompt,
+                    (idx + 1) * num_outputs_per_prompt,
+                )
+                if result.media is not None
+                else None
+            ),
             error=result.error,
             finished=result.finished,
             stage_durations=result.stage_durations,

vllm_omni/engine/arg_utils.py

@@ -566,6 +566,7 @@ class OrchestratorArgs:
     diffusion_compile_dynamic: bool | None = None
     cache_backend: str = "none"
     cache_config: str | None = None
+    video_output_transport: dict[str, object] | None = None
     enable_cache_dit_summary: bool = False
     step_execution: bool = False
     vae_use_slicing: bool = False

vllm_omni/entrypoints/cli/serve.py

@@ -71,6 +71,16 @@ def _nonneg_finite_float(value: str) -> float:
     return parsed
 
 
+def _json_object(value: str) -> dict[str, object]:
+    try:
+        parsed = json.loads(value)
+    except json.JSONDecodeError as exc:
+        raise argparse.ArgumentTypeError(f"must be valid JSON: {exc}") from exc
+    if not isinstance(parsed, dict):
+        raise argparse.ArgumentTypeError("must be a JSON object")
+    return parsed
+
+
 def _ensure_vllm_platform():
     """Ensure vLLM's current_platform is valid before arg parsing.
 
@@ -669,6 +679,14 @@ class OmniServeCommand(CLISubcommand):
             'MagCache: \'{"mag_threshold": 0.24, "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1}\'. '
             "Calibration mode: add '\"mag_calibrate\": true'",
         )
+        omni_config_group.add_argument(
+            "--video-output-transport",
+            type=_json_object,
+            default=None,
+            help=(
+                "JSON object configuring video output preparation, for example '{\"enable_device_postprocess\": true}'."
+            ),
+        )
         omni_config_group.add_argument(
             "--enable-cache-dit-summary",
             action="store_true",