| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from functools import partial |
| 4 | from io import BytesIO |
| 5 | from pathlib import Path |
| 6 | from typing import Any |
| 7 | |
| 8 | import numpy as np |
| 9 | import numpy.typing as npt |
| 10 | import pybase64 |
| 11 | import torch |
| 12 | from PIL import Image |
| 13 | |
| 14 | from vllm import envs |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.utils.serial_utils import tensor2base64 |
| 17 | from vllm.utils.sparse_utils import ( |
| 18 | check_sparse_tensor_invariants_threadsafe, |
| 19 | safe_to_dense, |
| 20 | ) |
| 21 | |
| 22 | from ..video import VIDEO_LOADER_REGISTRY, DecodedFrames |
| 23 | from ..video_decoders import incompatible_backend_options |
| 24 | from .base import MediaIO, MediaWithBytes |
| 25 | from .image import MAGIC_NUMPY_PREFIX, ImageMediaIO |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | |
| 30 | class VideoMediaIO(MediaIO[MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]]): |
| 31 | """Configuration values can be user-provided either by --media-io-kwargs or |
| 32 | by the runtime API field "media_io_kwargs". Ensure proper validation and |
| 33 | error handling. |
| 34 | |
| 35 | Use --media-io-kwargs '{"video": {"backend": "torchcodec", |
| 36 | "device": "cuda"}}' to decode with NVDEC and keep frames on the GPU, |
| 37 | which pairs with a device-side multi-modal processor |
| 38 | (--mm-processor-device, EPD encode-only instances). |
| 39 | """ |
| 40 | |
| 41 | @classmethod |
| 42 | def merge_kwargs( |
| 43 | cls, |
| 44 | default_kwargs: dict[str, Any] | None, |
| 45 | runtime_kwargs: dict[str, Any] | None, |
| 46 | ) -> dict[str, Any]: |
| 47 | if runtime_kwargs: |
| 48 | # Decoder GPU memory is reserved from the startup value. |
| 49 | runtime_kwargs = dict(runtime_kwargs) |
| 50 | runtime_kwargs.pop("hw_decoders", None) |
| 51 | runtime_kwargs.pop("pool_size", None) |
| 52 | # The decode device determines which device the decoded frames |
| 53 | # (and thus the processor pipeline) live on; startup-only. |
| 54 | runtime_kwargs.pop("device", None) |
| 55 | |
| 56 | # Block request-level selection of GPU video backends that |
| 57 | # were not configured (and VRAM-reserved) at startup. |
| 58 | for key in ("video_backend", "backend"): |
| 59 | requested = runtime_kwargs.get(key) |
| 60 | if requested and VIDEO_LOADER_REGISTRY.backend_requires_gpu(requested): |
| 61 | static_val = (default_kwargs or {}).get(key) |
| 62 | if static_val != requested: |
| 63 | logger.warning_once( |
| 64 | "Stripping request-level %s=%r: GPU video " |
| 65 | "backend not configured at startup.", |
| 66 | key, |
| 67 | requested, |
| 68 | ) |
| 69 | runtime_kwargs = { |
| 70 | k: v for k, v in runtime_kwargs.items() if k != key |
| 71 | } |
| 72 | |
| 73 | merged = super().merge_kwargs(default_kwargs, runtime_kwargs) |
| 74 | # fps and num_frames interact with each other, so if either is |
| 75 | # overridden at request time, wipe the other from defaults to |
| 76 | # avoid unintuitive cross-field interactions. |
| 77 | if runtime_kwargs: |
| 78 | if "num_frames" in runtime_kwargs and "fps" not in runtime_kwargs: |
| 79 | merged.pop("fps", None) |
| 80 | elif "fps" in runtime_kwargs and "num_frames" not in runtime_kwargs: |
| 81 | merged.pop("num_frames", None) |
| 82 | |
| 83 | requested_backend = runtime_kwargs.get("backend") |
| 84 | static_backend = (default_kwargs or {}).get("backend") |
| 85 | if requested_backend and requested_backend != static_backend: |
| 86 | for name in incompatible_backend_options( |
| 87 | static_backend, requested_backend |
| 88 | ): |
| 89 | merged.pop(name, None) |
| 90 | return merged |
| 91 | |
| 92 | def __init__( |
| 93 | self, |
| 94 | image_io: ImageMediaIO, |
| 95 | num_frames: int = 32, |
| 96 | **kwargs, |
| 97 | ) -> None: |
| 98 | super().__init__() |
| 99 | |
| 100 | self.image_io = image_io |
| 101 | self.num_frames = num_frames |
| 102 | # `kwargs` contains custom arguments from |
| 103 | # --media-io-kwargs for this modality, merged with |
| 104 | # per-request runtime media_io_kwargs via merge_kwargs(). |
| 105 | # They can be passed to the underlying |
| 106 | # media loaders (e.g. custom implementations) |
| 107 | # for flexible control. |
| 108 | |
| 109 | # Allow per-request override of video backend via kwargs. |
| 110 | # This enables users to specify a different backend than the |
| 111 | # global VLLM_VIDEO_LOADER_BACKEND env var, e.g.: |
| 112 | # --media-io-kwargs '{"video": {"video_backend": "torchcodec"}}' |
| 113 | video_loader_backend = ( |
| 114 | kwargs.pop("video_backend", None) or envs.VLLM_VIDEO_LOADER_BACKEND |
| 115 | ) |
| 116 | self.kwargs = kwargs |
| 117 | self.video_loader = VIDEO_LOADER_REGISTRY.load(video_loader_backend) |
| 118 | |
| 119 | def load_bytes( |
| 120 | self, data: bytes |
| 121 | ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: |
| 122 | video = self.video_loader.load_bytes( |
| 123 | data, num_frames=self.num_frames, **self.kwargs |
| 124 | ) |
| 125 | return MediaWithBytes(video, data) |
| 126 | |
| 127 | def load_base64( |
| 128 | self, media_type: str, data: str |
| 129 | ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: |
| 130 | if media_type.lower() == "video/jpeg": |
| 131 | load_frame = partial( |
| 132 | self.image_io.load_base64, |
| 133 | "image/jpeg", |
| 134 | ) |
| 135 | |
| 136 | if self.num_frames > 0: |
| 137 | frame_parts = data.split(",", self.num_frames)[: self.num_frames] |
| 138 | elif self.num_frames == 0: |
| 139 | raise ValueError("num_frames must be greater than 0 or -1") |
| 140 | else: |
| 141 | frame_parts = data.split(",") |
| 142 | |
| 143 | frames = np.stack( |
| 144 | [np.asarray(load_frame(frame_data)) for frame_data in frame_parts] |
| 145 | ) |
| 146 | total = int(frames.shape[0]) |
| 147 | fps = float(self.kwargs.get("fps", 1)) |
| 148 | |
| 149 | # validate and extract frames_indices |
| 150 | frames_indices = self.kwargs.get("frames_indices") |
| 151 | if frames_indices is not None: |
| 152 | if not ( |
| 153 | isinstance(frames_indices, list) |
| 154 | and all(isinstance(i, int) for i in frames_indices) |
| 155 | ): |
| 156 | raise ValueError("frames_indices must be a list of integers") |
| 157 | if len(frames_indices) != total: |
| 158 | raise ValueError( |
| 159 | f"frames_indices length ({len(frames_indices)}) must " |
| 160 | f"match number of frames sent ({total})" |
| 161 | ) |
| 162 | else: |
| 163 | frames_indices = list(range(total)) |
| 164 | |
| 165 | # validate and extract total_num_frames |
| 166 | total_num_frames = self.kwargs.get("total_num_frames", total) |
| 167 | if not isinstance(total_num_frames, int) or total_num_frames < 1: |
| 168 | raise ValueError("total_num_frames must be a positive integer") |
| 169 | if total_num_frames < total: |
| 170 | raise ValueError( |
| 171 | f"total_num_frames ({total_num_frames}) must be >= " |
| 172 | f"number of frames sent ({total})" |
| 173 | ) |
| 174 | |
| 175 | # validate and extract duration |
| 176 | duration = self.kwargs.get("duration") |
| 177 | if duration is not None: |
| 178 | if not isinstance(duration, (int, float)) or duration < 0: |
| 179 | raise ValueError("duration must be a non-negative number") |
| 180 | else: |
| 181 | duration = total_num_frames / fps if fps > 0 else 0.0 |
| 182 | |
| 183 | metadata = { |
| 184 | "total_num_frames": total_num_frames, |
| 185 | "fps": fps, |
| 186 | "duration": duration, |
| 187 | "video_backend": "jpeg_sequence", |
| 188 | "frames_indices": frames_indices, |
| 189 | "do_sample_frames": self.kwargs.get("do_sample_frames", False), |
| 190 | } |
| 191 | return MediaWithBytes((frames, metadata), data.encode()) |
| 192 | |
| 193 | return self.load_bytes(pybase64.b64decode(data, validate=True)) |
| 194 | |
| 195 | def load_file( |
| 196 | self, filepath: Path |
| 197 | ) -> MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]: |
| 198 | with filepath.open("rb") as f: |
| 199 | data = f.read() |
| 200 | |
| 201 | return self.load_bytes(data) |
| 202 | |
| 203 | def encode_base64( |
| 204 | self, |
| 205 | media: npt.NDArray, |
| 206 | *, |
| 207 | video_format: str = "JPEG", |
| 208 | ) -> str: |
| 209 | video = media |
| 210 | |
| 211 | if video_format == "JPEG": |
| 212 | encode_frame = partial( |
| 213 | self.image_io.encode_base64, |
| 214 | image_format=video_format, |
| 215 | ) |
| 216 | |
| 217 | return ",".join(encode_frame(Image.fromarray(frame)) for frame in video) |
| 218 | |
| 219 | msg = "Only JPEG format is supported for now." |
| 220 | raise NotImplementedError(msg) |
| 221 | |
| 222 | |
| 223 | class VideoEmbeddingMediaIO(MediaIO[torch.Tensor]): |
| 224 | """Video embedding MediaIO implementation. |
| 225 | |
| 226 | Configuration values can be user-provided either by --media-io-kwargs or |
| 227 | by the runtime API field "media_io_kwargs". Ensure proper validation and |
| 228 | error handling. |
| 229 | """ |
| 230 | |
| 231 | def __init__(self) -> None: |
| 232 | super().__init__() |
| 233 | |
| 234 | def _load_pickled_torch(self, data: bytes) -> torch.Tensor: |
| 235 | buffer = BytesIO(data) |
| 236 | with check_sparse_tensor_invariants_threadsafe(): |
| 237 | tensor = torch.load(buffer, weights_only=True) |
| 238 | return safe_to_dense(tensor, parameter="video_embeds") |
| 239 | |
| 240 | def _load_numpy(self, data: bytes) -> torch.Tensor: |
| 241 | with BytesIO(data) as buffer: |
| 242 | return torch.from_numpy(np.load(buffer)) |
| 243 | |
| 244 | def load_bytes(self, data: bytes) -> torch.Tensor: |
| 245 | if data[:6] == MAGIC_NUMPY_PREFIX: |
| 246 | return self._load_numpy(data) |
| 247 | |
| 248 | return self._load_pickled_torch(data) |
| 249 | |
| 250 | def load_base64(self, media_type: str, data: str) -> torch.Tensor: |
| 251 | return self.load_bytes(pybase64.b64decode(data, validate=True)) |
| 252 | |
| 253 | def load_file(self, filepath: Path) -> torch.Tensor: |
| 254 | if filepath.suffix == ".npy": |
| 255 | return torch.from_numpy(np.load(filepath)) |
| 256 | |
| 257 | with check_sparse_tensor_invariants_threadsafe(): |
| 258 | tensor = torch.load(filepath, weights_only=True) |
| 259 | return safe_to_dense(tensor, parameter="video_embeds") |
| 260 | |
| 261 | def encode_base64(self, media: torch.Tensor) -> str: |
| 262 | return tensor2base64(media) |