# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for OpenAI-compatible video API encoding helpers.""" import base64 from io import BytesIO import httpx import numpy as np import pytest import torch from PIL import Image from vllm import envs from vllm_omni.diffusion.postprocess import rife_interpolator from vllm_omni.entrypoints.openai import video_api_utils from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _png_bytes() -> bytes: buffer = BytesIO() Image.new("RGB", (2, 1), color=(12, 34, 56)).save(buffer, format="PNG") return buffer.getvalue() def _install_http_transport(monkeypatch, handler, client_kwargs): async_client = httpx.AsyncClient def _client_factory(*args, **kwargs): client_kwargs.append(kwargs.copy()) return async_client(*args, transport=httpx.MockTransport(handler), **kwargs) monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory) @pytest.mark.asyncio async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch): requested_paths = [] client_kwargs = [] def _handler(request): requested_paths.append(request.url.path) if request.url.path == "/redirect.png": return httpx.Response(302, headers={"location": "/image.png"}) return httpx.Response(200, content=_png_bytes(), headers={"content-type": "image/png"}) monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True) _install_http_transport(monkeypatch, _handler, client_kwargs) image = await video_api_utils.decode_image_url("https://example.com/redirect.png") assert image.size == (2, 1) assert image.mode == "RGB" assert requested_paths == ["/redirect.png", "/image.png"] assert client_kwargs == [{"timeout": 60, "follow_redirects": True}] @pytest.mark.asyncio async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch): requested_paths = [] client_kwargs = [] def _handler(request): requested_paths.append(request.url.path) return httpx.Response(302, headers={"location": "/image.png"}) monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", False) _install_http_transport(monkeypatch, _handler, client_kwargs) with pytest.raises(InvalidInputReferenceError, match="redirect.*VLLM_MEDIA_URL_ALLOW_REDIRECTS"): await video_api_utils.decode_image_url("https://example.com/redirect.png") assert requested_paths == ["/redirect.png"] assert client_kwargs == [{"timeout": 60, "follow_redirects": False}] @pytest.mark.asyncio async def test_decode_image_url_reports_http_status(monkeypatch): client_kwargs = [] def _handler(request): return httpx.Response(404) monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True) _install_http_transport(monkeypatch, _handler, client_kwargs) with pytest.raises(InvalidInputReferenceError, match="server returned HTTP 404"): await video_api_utils.decode_image_url("https://example.com/missing.png") @pytest.mark.asyncio async def test_decode_image_url_reports_connection_failure(monkeypatch): client_kwargs = [] def _handler(request): raise httpx.ConnectError("connection refused", request=request) monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True) _install_http_transport(monkeypatch, _handler, client_kwargs) with pytest.raises(InvalidInputReferenceError, match="failed to download image"): await video_api_utils.decode_image_url("https://example.com/unreachable.png") @pytest.mark.asyncio async def test_decode_image_url_keeps_data_urls_local(monkeypatch): def _unexpected_http_client(*args, **kwargs): pytest.fail("data URLs must not create an HTTP client") monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _unexpected_http_client) encoded_image = base64.b64encode(_png_bytes()).decode() image = await video_api_utils.decode_image_url(f"data:image/png;base64,{encoded_image}") assert image.size == (2, 1) assert image.mode == "RGB" def _install_fake_video_mux(monkeypatch, mux_calls): def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None): mux_calls.append( { "frames": frames, "audio": audio, "fps": fps, "audio_sample_rate": audio_sample_rate, "video_codec_options": video_codec_options, } ) return b"fake-video" monkeypatch.setattr( "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes", _fake_mux_video_audio_bytes, ) def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): mux_calls = [] _install_fake_video_mux(monkeypatch, mux_calls) frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)] video_bytes = video_api_utils._encode_video_bytes( frames, fps=8, ) assert video_bytes == b"fake-video" assert mux_calls[0]["frames"].shape == (5, 2, 2, 3) assert mux_calls[0]["frames"].dtype == np.uint8 assert mux_calls[0]["fps"] == 8.0 assert mux_calls[0]["audio"] is None def test_float_frames_are_converted_without_stacking_full_video(monkeypatch): frame = np.array( [ [[0.0, 0.5, 1.0, 0.2], [1.5, -0.5, 0.5, 0.8]], ], dtype=np.float32, ) original = frame.copy() def fail_stack(*args, **kwargs): raise AssertionError("float video conversion must not stack all frames") monkeypatch.setattr(video_api_utils.np, "stack", fail_stack) frames = video_api_utils._coerce_video_to_uint8_frames([frame, frame]) expected = np.array([[[128, 191, 255], [255, 64, 191]]], dtype=np.uint8) assert frames.flags.c_contiguous np.testing.assert_array_equal(frames, np.array([expected, expected])) np.testing.assert_array_equal(frame, original) def test_two_dimensional_width_four_frames_are_not_treated_as_rgba(): frame = np.arange(8, dtype=np.float32).reshape(2, 4) / 8 frames = video_api_utils._coerce_video_to_uint8_frames([frame, frame]) expected = np.rint(frame * 255).astype(np.uint8) assert frames.shape == (2, 2, 4) np.testing.assert_array_equal(frames, np.stack([expected, expected])) def test_mixed_float_dtypes_preserve_stacked_rounding_semantics(): frames = [ np.full((1, 1, 3), 0.1, dtype=np.float16), np.full((1, 1, 3), 0.2, dtype=np.float32), ] stacked = np.stack(frames) expected = np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8) actual = video_api_utils._coerce_video_to_uint8_frames(frames) assert expected[0, 0, 0, 0] == 25 np.testing.assert_array_equal(actual, expected) @pytest.mark.parametrize("frame_count", [3, 4]) def test_channel_last_video_tensor_preserves_channel_sized_frame_count(frame_count): video = torch.arange(frame_count * 2 * 5 * 3, dtype=torch.uint8).reshape(frame_count, 2, 5, 3) frames = video_api_utils._coerce_video_to_uint8_frames(video) assert frames.shape == (frame_count, 2, 5, 3) np.testing.assert_array_equal(frames, video.numpy()) def test_channel_first_video_tensor_is_converted_to_channel_last(): video = torch.arange(3 * 5 * 2 * 6, dtype=torch.uint8).reshape(3, 5, 2, 6) frames = video_api_utils._coerce_video_to_uint8_frames(video) assert frames.shape == (5, 2, 6, 3) np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy()) def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch): muxers = [] class FakeFragmentedMP4Muxer: def __init__(self, *, width, height, fps, video_codec_options=None): self.calls = [] muxers.append( { "width": width, "height": height, "fps": fps, "video_codec_options": video_codec_options, "instance": self, } ) def mux_video_frames(self, frames): self.calls.append(frames) return f"fragment-{len(self.calls)}".encode() def close(self): return b"trailer" monkeypatch.setattr( "vllm_omni.diffusion.utils.media_utils.FragmentedMP4Muxer", FakeFragmentedMP4Muxer, ) encoder = video_api_utils.FragmentedMP4VideoEncoder( fps=12, video_codec_options={"preset": "ultrafast"}, ) assert encoder.encode([np.zeros((4, 6, 3), dtype=np.float32)]) == b"fragment-1" assert encoder.encode([np.ones((4, 6, 3), dtype=np.float32)]) == b"fragment-2" assert encoder.close() == b"trailer" assert len(muxers) == 1 assert muxers[0]["width"] == 6 assert muxers[0]["height"] == 4 assert muxers[0]["fps"] == 12.0 assert muxers[0]["video_codec_options"] == {"preset": "ultrafast"} assert len(muxers[0]["instance"].calls) == 2 def test_create_streaming_video_encoder_selects_requested_format(): assert isinstance( video_api_utils.create_streaming_video_encoder(output_format="m4s", fps=12), video_api_utils.FragmentedMP4VideoEncoder, ) def test_finalize_streaming_mp4_bytes_produces_progressive_mp4(): """Fragment MP4 chunks are remuxed into decodable progressive MP4 bytes.""" import numpy as np from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer, finalize_streaming_video_bytes def _read_mp4_video_info(mp4_bytes: bytes) -> tuple[int, float, float | None, int, int]: import io import av with av.open(io.BytesIO(mp4_bytes), mode="r", format="mp4") as container: stream = container.streams.video[0] frame_count = sum(1 for _ in container.decode(stream)) assert stream.average_rate is not None fps = float(stream.average_rate) duration_sec = None if stream.duration is not None: assert stream.time_base is not None duration_sec = float(stream.duration * stream.time_base) return frame_count, fps, duration_sec, stream.width, stream.height width = 32 height = 32 fps = 16.0 input_frame_count = 2 muxer = FragmentedMP4Muxer(width=width, height=height, fps=fps) frames = np.zeros((input_frame_count, height, width, 3), dtype=np.uint8) streamed = muxer.mux_video_frames(frames) + muxer.close() finalized = finalize_streaming_video_bytes(streamed, input_format="m4s", fps=fps) assert finalized assert finalized != streamed streamed_info = _read_mp4_video_info(streamed) final_info = _read_mp4_video_info(finalized) expected_duration = input_frame_count / fps assert streamed_info[0] == input_frame_count assert final_info[0] == input_frame_count assert streamed_info[1] == pytest.approx(fps) assert final_info[1] == pytest.approx(fps) assert streamed_info[3:] == (width, height) assert final_info[3:] == (width, height) assert final_info[2] == pytest.approx(expected_duration, rel=0.05) assert streamed_info[2] == pytest.approx(expected_duration, rel=0.05) def test_rife_model_inference_runs_on_dummy_tensors(): model = rife_interpolator.Model().eval() img0 = torch.rand(1, 3, 32, 32) img1 = torch.rand(1, 3, 32, 32) output = model.inference(img0, img1, scale=1.0) assert output.shape == (1, 3, 32, 32) assert torch.isfinite(output).all() def test_frame_interpolator_runs_actual_torch_tensor_path(monkeypatch): model = rife_interpolator.Model().eval() interpolator = rife_interpolator.FrameInterpolator() monkeypatch.setattr(interpolator, "_ensure_model_loaded", lambda preferred_device=None: model) video = torch.zeros(1, 3, 2, 32, 32) output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0) assert multiplier == 2 assert output_video.shape == (1, 3, 3, 32, 32) assert torch.isfinite(output_video).all() def test_frame_interpolator_uses_platform_device_when_tensor_is_cpu(monkeypatch): chosen_devices = [] model = rife_interpolator.Model().eval() def _fake_ensure_model_loaded(*, preferred_device=None): chosen_devices.append(preferred_device) return model interpolator = rife_interpolator.FrameInterpolator() monkeypatch.setattr(interpolator, "_ensure_model_loaded", _fake_ensure_model_loaded) monkeypatch.setattr(model.flownet, "to", lambda device: model.flownet) monkeypatch.setattr(rife_interpolator, "_select_torch_device", lambda: torch.device("cuda")) video = torch.zeros(1, 3, 2, 32, 32) output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0) assert chosen_devices == [torch.device("cuda")] assert multiplier == 2 assert output_video.shape == (1, 3, 3, 32, 32)