diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 86757e56d..65d285217 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -548,6 +548,23 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation( assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0 +def test_decode_video_bytes_can_keep_first_frames(): + from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes + + frames = _decode_video_bytes( + _make_test_video_bytes((32, 24), num_frames=6), + source="input_reference", + max_frames=2, + keep="first", + ) + + assert len(frames) == 2 + assert frames.fps == pytest.approx(8.0) + red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames] + assert red_means[0] < red_means[1] + assert red_means[1] < 100 + + def test_decode_video_bytes_can_keep_last_frames(): from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index f0d2a59f2..4421d74a2 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -8,7 +8,6 @@ from __future__ import annotations import base64 import binascii -from collections import deque from io import BytesIO from typing import Any, Literal @@ -16,6 +15,12 @@ import httpx import numpy as np import torch from PIL import Image, UnidentifiedImageError +from vllm.multimodal.video import ( + VIDEO_LOADER_REGISTRY, + VideoBackend, + VideoSourceMetadata, + VideoTargetMetadata, +) from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError from vllm_omni.entrypoints.openai.protocol.videos import ( @@ -58,6 +63,30 @@ def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image: raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc +@VIDEO_LOADER_REGISTRY.register("omni") +class OmniVideoBackend(VideoBackend): + """Video backend that selects sequential first-N or last-N frames.""" + + @classmethod + def compute_frames_index_to_sample( + cls, + source: VideoSourceMetadata, + target: VideoTargetMetadata, + *, + keep: Literal["first", "last"] = "first", + **kwargs, + ) -> list[int]: + num_frames_to_sample = source.total_frames_num + if target.num_frames > 0: + num_frames_to_sample = min(num_frames_to_sample, target.num_frames) + num_frames_to_sample = max(1, num_frames_to_sample) + if num_frames_to_sample >= source.total_frames_num: + return list(range(source.total_frames_num)) + if keep == "last": + return list(range(source.total_frames_num - num_frames_to_sample, source.total_frames_num)) + return list(range(num_frames_to_sample)) + + def _decode_video_bytes( video_bytes: bytes, *, @@ -65,45 +94,29 @@ def _decode_video_bytes( max_frames: int | None = None, keep: Literal["first", "last"] = "first", ) -> VideoFrames: - try: - import av - except ImportError as exc: # pragma: no cover - av is a serving dependency via media_utils - raise InvalidInputReferenceError(f"Invalid {source}: video decoding requires PyAV.") from exc - if keep not in {"first", "last"}: raise InvalidInputReferenceError(f"Invalid {source}: video frame selection must be 'first' or 'last'.") if max_frames is not None and max_frames <= 0: raise InvalidInputReferenceError(f"Invalid {source}: max video frames must be positive.") - frames: list[Image.Image] = [] - tail_frames: deque[Image.Image] | None = ( - deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None - ) - fps: float | None = None + loader = VIDEO_LOADER_REGISTRY.load("omni") + num_frames = max_frames if max_frames is not None else -1 try: - with av.open(BytesIO(video_bytes)) as container: - video_stream = container.streams.video[0] if container.streams.video else None - if video_stream is not None: - fps = ( - positive_float(getattr(video_stream, "average_rate", None)) - or positive_float(getattr(video_stream, "base_rate", None)) - or positive_float(getattr(video_stream, "guessed_rate", None)) - ) - for frame in container.decode(video=0): - image = frame.to_image().convert("RGB") - if tail_frames is not None: - tail_frames.append(image) - else: - frames.append(image) - if keep == "first" and max_frames is not None and len(frames) >= max_frames: - break + frames_array, metadata = loader.load_bytes( + video_bytes, + num_frames=num_frames, + backend="pyav", + keep=keep, + ) except Exception as exc: raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") from exc - if tail_frames is not None: - frames = list(tail_frames) + fps: float | None = positive_float(metadata.get("fps")) + + frames = [Image.fromarray(f, "RGB") for f in frames_array] if not frames: raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") + return VideoFrames(frames, fps=fps)