[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 = {}