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