diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index ee4fc6f3e..583f935dc 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.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 diff --git a/docs/.nav.yml b/docs/.nav.yml index aa84f2819..8dac28d22 100644 --- a/docs/.nav.yml +++ b/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 diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 4c3ad3373..ea0fc9e0c 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/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, diff --git a/docs/user_guide/diffusion/device_side_video_postprocess.md b/docs/user_guide/diffusion/device_side_video_postprocess.md new file mode 100644 index 000000000..910e64457 --- /dev/null +++ b/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. diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 40229f521..da8ffb452 100644 --- a/tests/config/test_omni_config.py +++ b/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 diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index de5abeab0..6f7833ff8 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/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 = {} diff --git a/tests/diffusion/test_async_output_worker.py b/tests/diffusion/test_async_output_worker.py index a4bff4625..8a81f30d4 100644 --- a/tests/diffusion/test_async_output_worker.py +++ b/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.""" diff --git a/tests/diffusion/test_device_output_preparation.py b/tests/diffusion/test_device_output_preparation.py new file mode 100644 index 000000000..0a6f1ab90 --- /dev/null +++ b/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() diff --git a/tests/diffusion/test_device_reduction.py b/tests/diffusion/test_device_reduction.py new file mode 100644 index 000000000..f69bb5eb5 --- /dev/null +++ b/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)) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 9f8d4dec6..69a9e4395 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/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(): diff --git a/tests/diffusion/test_ipc_async.py b/tests/diffusion/test_ipc_async.py index 03919095b..43650c581 100644 --- a/tests/diffusion/test_ipc_async.py +++ b/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, ) diff --git a/tests/diffusion/test_media_output_contract.py b/tests/diffusion/test_media_output_contract.py new file mode 100644 index 000000000..ae5c2fdbf --- /dev/null +++ b/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"])) diff --git a/tests/diffusion/test_video_transport_defaults_are_inert.py b/tests/diffusion/test_video_transport_defaults_are_inert.py new file mode 100644 index 000000000..fd7346e76 --- /dev/null +++ b/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 diff --git a/tests/diffusion/test_wan_device_reduction.py b/tests/diffusion/test_wan_device_reduction.py new file mode 100644 index 000000000..492aec7a5 --- /dev/null +++ b/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() diff --git a/tests/docs/test_generate_argparse.py b/tests/docs/test_generate_argparse.py new file mode 100644 index 000000000..16b92365f --- /dev/null +++ b/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} diff --git a/tests/e2e/offline_inference/test_device_postprocess_equivalence.py b/tests/e2e/offline_inference/test_device_postprocess_equivalence.py new file mode 100644 index 000000000..a7fbcafd9 --- /dev/null +++ b/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" + ) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index b64758e5d..0089330b2 100644 --- a/tests/entrypoints/test_serve.py +++ b/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") diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 81eb76f0b..2b016be0d 100644 --- a/vllm_omni/config/omni_config.py +++ b/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): diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2ac1f99dd..f7d714890 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c30529520..d27cc36e1 100644 --- a/vllm_omni/diffusion/data.py +++ b/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) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index c36ac0459..417090503 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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,26 +535,33 @@ class DiffusionEngine: raise RuntimeError(output.error) logger.debug("Generation completed successfully.") - if output.output is None: - logger.warning("Output is None, returning empty OmniRequestOutput") - return format_empty_diffusion_outputs(request, finished=output.finished) - - # When CPU offload is enabled, move output to CPU before - # post-processing to avoid device OOM — model weights may still - # reside on the device and leave no headroom for intermediates. - output_data = output.output - if self.od_config.enable_cpu_offload: - output_data = _move_tensor_tree_to_cpu(output_data) - - if self.post_process_func is not None: - # Some video pipelines need request-level controls during - # postprocess (for example worker-side frame interpolation). - postprocess_kwargs: dict[str, object] = {} - if self._post_process_accepts_sampling_params: - postprocess_kwargs["sampling_params"] = request.sampling_params - outputs = self.post_process_func(output_data, **postprocess_kwargs) + 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: - outputs = output_data + if output.output is None: + logger.warning("Output is None, returning empty OmniRequestOutput") + return format_empty_diffusion_outputs(request, finished=output.finished) + + # When CPU offload is enabled, move output to CPU before + # post-processing to avoid device OOM — model weights may still + # reside on the device and leave no headroom for intermediates. + output_data = output.output + if self.od_config.enable_cpu_offload: + output_data = _move_tensor_tree_to_cpu(output_data) + + if self.post_process_func is not None: + # Some video pipelines need request-level controls during + # postprocess (for example worker-side frame interpolation). + postprocess_kwargs: dict[str, object] = {} + if self._post_process_accepts_sampling_params: + postprocess_kwargs["sampling_params"] = request.sampling_params + outputs = self.post_process_func(output_data, **postprocess_kwargs) + else: + outputs = output_data postprocess_output = normalize_diffusion_postprocess_output(outputs) diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index 6eb24c540..f0e88a2c9 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/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) diff --git a/vllm_omni/diffusion/media.py b/vllm_omni/diffusion/media.py new file mode 100644 index 000000000..34f4a1224 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index c9ef325c4..a285e0009 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/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, diff --git a/vllm_omni/diffusion/postprocess/device_reduction.py b/vllm_omni/diffusion/postprocess/device_reduction.py new file mode 100644 index 000000000..ec4ef895f --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/postprocess/media.py b/vllm_omni/diffusion/postprocess/media.py new file mode 100644 index 000000000..fd2564434 --- /dev/null +++ b/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, + } diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 8804785fa..9dfcd6f3f 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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, diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 4443353a0..b9506f102 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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) diff --git a/vllm_omni/diffusion/worker/request_batch.py b/vllm_omni/diffusion/worker/request_batch.py index 079a213db..d5f1d5585 100644 --- a/vllm_omni/diffusion/worker/request_batch.py +++ b/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, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 44bd9058c..ee5552751 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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 diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 5c0204a33..348e41316 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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",