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

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