| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import bisect |
| 5 | import mimetypes |
| 6 | from collections import defaultdict |
| 7 | from collections.abc import Generator, Sequence |
| 8 | from dataclasses import replace |
| 9 | from itertools import groupby |
| 10 | from typing import TYPE_CHECKING, Any |
| 11 | |
| 12 | import numpy as np |
| 13 | import numpy.typing as npt |
| 14 | from PIL import Image |
| 15 | |
| 16 | from vllm.inputs import MultiModalPlaceholders |
| 17 | from vllm.utils.import_utils import LazyLoader |
| 18 | |
| 19 | from .inputs import ( |
| 20 | BatchedTensorInputs, |
| 21 | MultiModalFeatureSpec, |
| 22 | MultiModalFieldElem, |
| 23 | MultiModalKwargsItem, |
| 24 | MultiModalSharedField, |
| 25 | nested_tensors_equal, |
| 26 | ) |
| 27 | from .media import ( |
| 28 | AudioMediaIO, |
| 29 | ImageMediaIO, |
| 30 | MediaConnector, |
| 31 | MediaWithBytes, |
| 32 | VideoMediaIO, |
| 33 | ) |
| 34 | |
| 35 | if TYPE_CHECKING: |
| 36 | import torch.types |
| 37 | else: |
| 38 | torch = LazyLoader("torch", globals(), "torch") |
| 39 | |
| 40 | |
| 41 | def encode_audio_base64( |
| 42 | audio: np.ndarray, |
| 43 | sampling_rate: int, |
| 44 | *, |
| 45 | format: str = "WAV", |
| 46 | ) -> str: |
| 47 | """Encode audio as base64.""" |
| 48 | audio_io = AudioMediaIO() |
| 49 | return audio_io.encode_base64((audio, sampling_rate), audio_format=format) |
| 50 | |
| 51 | |
| 52 | def encode_audio_url( |
| 53 | audio: np.ndarray, |
| 54 | sampling_rate: int, |
| 55 | *, |
| 56 | format: str = "WAV", |
| 57 | ) -> str: |
| 58 | """Encode audio as a data URL.""" |
| 59 | audio_b64 = encode_audio_base64(audio, sampling_rate, format=format) |
| 60 | mimetype = mimetypes.types_map.get("." + format.lower(), f"audio/{format.lower()}") |
| 61 | return f"data:{mimetype};base64,{audio_b64}" |
| 62 | |
| 63 | |
| 64 | def encode_image_base64( |
| 65 | image: Image.Image, |
| 66 | *, |
| 67 | image_mode: str | None = "RGB", |
| 68 | format: str = "PNG", |
| 69 | ) -> str: |
| 70 | """Encode a pillow image to base64 format. |
| 71 | |
| 72 | By default, the image is converted into RGB format before being encoded. |
| 73 | Pass `image_mode=None` to keep the original image mode. |
| 74 | """ |
| 75 | image_io = ImageMediaIO(image_mode=image_mode) |
| 76 | return image_io.encode_base64(image, image_format=format) |
| 77 | |
| 78 | |
| 79 | def encode_image_url( |
| 80 | image: Image.Image, |
| 81 | *, |
| 82 | image_mode: str | None = "RGB", |
| 83 | format: str = "PNG", |
| 84 | ) -> str: |
| 85 | """Encode a pillow image as a data URL. |
| 86 | |
| 87 | By default, the image is converted into RGB format before being encoded. |
| 88 | Pass `image_mode=None` to keep the original image mode. |
| 89 | """ |
| 90 | image_b64 = encode_image_base64(image, image_mode=image_mode, format=format) |
| 91 | mimetype = mimetypes.types_map.get("." + format.lower(), f"image/{format.lower()}") |
| 92 | return f"data:{mimetype};base64,{image_b64}" |
| 93 | |
| 94 | |
| 95 | def encode_video_base64( |
| 96 | frames: npt.NDArray, |
| 97 | *, |
| 98 | format: str = "JPEG", |
| 99 | ) -> str: |
| 100 | image_io = ImageMediaIO() |
| 101 | video_io = VideoMediaIO(image_io) |
| 102 | return video_io.encode_base64(frames, video_format=format) |
| 103 | |
| 104 | |
| 105 | def encode_video_url( |
| 106 | frames: npt.NDArray, |
| 107 | *, |
| 108 | format: str = "JPEG", |
| 109 | ) -> str: |
| 110 | video_b64 = encode_video_base64(frames, format=format) |
| 111 | |
| 112 | if format.lower() == "jpeg": |
| 113 | mimetype = "video/jpeg" |
| 114 | else: |
| 115 | mimetype = mimetypes.types_map.get( |
| 116 | "." + format.lower(), f"video/{format.lower()}" |
| 117 | ) |
| 118 | |
| 119 | return f"data:{mimetype};base64,{video_b64}" |
| 120 | |
| 121 | |
| 122 | def get_mm_features_in_window( |
| 123 | mm_features: list[MultiModalFeatureSpec], |
| 124 | start: int, |
| 125 | end: int, |
| 126 | ) -> tuple[int, int]: |
| 127 | """Return (lo, hi) indices for features overlapping [start, end). |
| 128 | |
| 129 | Assumes mm_features are sorted by offset and non-overlapping, so |
| 130 | offset + length is also sorted. |
| 131 | """ |
| 132 | lo = bisect.bisect_left( |
| 133 | mm_features, |
| 134 | start + 1, |
| 135 | key=lambda f: f.mm_position.offset + f.mm_position.length, |
| 136 | ) |
| 137 | hi = bisect.bisect_left( |
| 138 | mm_features, |
| 139 | end, |
| 140 | key=lambda f: f.mm_position.offset, |
| 141 | ) |
| 142 | return lo, hi |
| 143 | |
| 144 | |
| 145 | def argsort_mm_positions( |
| 146 | mm_positions: MultiModalPlaceholders, |
| 147 | ) -> list[tuple[str, int]]: |
| 148 | """Given a `MultiModalPlaceholders`, output a sequence of keys to |
| 149 | sort the dictionary by `offset` (starting index in the input sequence) |
| 150 | in ascending order. |
| 151 | |
| 152 | Returns: |
| 153 | A list of `(modality, idx)`, which can be used to access an item |
| 154 | by `mm_positions[modality][idx]`. |
| 155 | |
| 156 | """ |
| 157 | flat_items = ( |
| 158 | (modality, idx, item) |
| 159 | for modality, items in mm_positions.items() |
| 160 | for idx, item in enumerate(items) |
| 161 | ) |
| 162 | |
| 163 | sorted_flat_items = sorted(flat_items, key=lambda x: x[2].offset) |
| 164 | |
| 165 | return [(modality, idx) for modality, idx, _ in sorted_flat_items] |
| 166 | |
| 167 | |
| 168 | def _can_batch_mm_items( |
| 169 | left: MultiModalKwargsItem, |
| 170 | right: MultiModalKwargsItem, |
| 171 | ) -> bool: |
| 172 | if left.keys() != right.keys(): |
| 173 | return False |
| 174 | |
| 175 | for key, left_elem in left.items(): |
| 176 | right_elem = right[key] |
| 177 | left_field, right_field = left_elem.field, right_elem.field |
| 178 | is_shared_field = isinstance(left_field, MultiModalSharedField) and isinstance( |
| 179 | right_field, MultiModalSharedField |
| 180 | ) |
| 181 | if (type(left_field) is not type(right_field)) or ( |
| 182 | is_shared_field |
| 183 | and not nested_tensors_equal( |
| 184 | left_elem.data, right_elem.data, check_dtype=True |
| 185 | ) |
| 186 | ): |
| 187 | return False |
| 188 | |
| 189 | return True |
| 190 | |
| 191 | |
| 192 | def _batch_mm_items( |
| 193 | items: Sequence[MultiModalKwargsItem], |
| 194 | *, |
| 195 | device: torch.types.Device = None, |
| 196 | pin_memory: bool = False, |
| 197 | ): |
| 198 | elems = defaultdict[str, list[MultiModalFieldElem]](list) |
| 199 | for item in items: |
| 200 | for key, elem in item.items(): |
| 201 | elems[key].append(elem) |
| 202 | |
| 203 | return { |
| 204 | key: elems[0].field.reduce_data( |
| 205 | elems, |
| 206 | device=device, |
| 207 | pin_memory=pin_memory, |
| 208 | ) |
| 209 | for key, elems in elems.items() |
| 210 | } |
| 211 | |
| 212 | |
| 213 | def strip_covered_mm_data( |
| 214 | mm_features: list[MultiModalFeatureSpec], |
| 215 | num_computed_tokens: int, |
| 216 | uses_mrope: bool = False, |
| 217 | ) -> list[MultiModalFeatureSpec]: |
| 218 | """Drop the tensor data of mm items whose placeholder span is fully inside |
| 219 | a prefix-cache-covered region: no encoder run can be scheduled for them, |
| 220 | so the workers never consume the payload fields. M-RoPE models keep |
| 221 | CPU-side metadata fields used to compute positions. SHM address items |
| 222 | are also kept so workers can balance the sender's reference count. The |
| 223 | scheduler-side ``Request`` keeps the full features.""" |
| 224 | if not mm_features or num_computed_tokens == 0: |
| 225 | return mm_features |
| 226 | |
| 227 | def maybe_strip(f: MultiModalFeatureSpec) -> MultiModalFeatureSpec: |
| 228 | if ( |
| 229 | f.data is None |
| 230 | # SHM address items must reach the worker so it can acknowledge |
| 231 | # the sender's reference count before the item is evicted. This |
| 232 | # mirrors how the SHM receiver identifies them ("address" in item). |
| 233 | or "address" in f.data |
| 234 | or (f.mm_position.offset + f.mm_position.length > num_computed_tokens) |
| 235 | ): |
| 236 | return f |
| 237 | |
| 238 | data = None |
| 239 | if uses_mrope: |
| 240 | data = MultiModalKwargsItem( |
| 241 | {k: elem for k, elem in f.data.items() if elem.field.keep_on_cpu} |
| 242 | ) |
| 243 | return replace(f, data=data) |
| 244 | |
| 245 | return [maybe_strip(f) for f in mm_features] |
| 246 | |
| 247 | |
| 248 | def group_and_batch_mm_items( |
| 249 | items: Sequence[MultiModalKwargsItem], |
| 250 | *, |
| 251 | device: torch.types.Device = None, |
| 252 | pin_memory: bool = False, |
| 253 | ) -> Generator[tuple[int, BatchedTensorInputs]]: |
| 254 | """Group consecutive items (possibly from different requests) into batches. |
| 255 | |
| 256 | Items must be split across groups if any of the following occurs, |
| 257 | as the batch would otherwise be invalid: |
| 258 | - They have different fields (e.g. mixed image and embedding inputs). |
| 259 | - They have different values in `MultiModalSharedField`. |
| 260 | |
| 261 | Args: |
| 262 | items: List of `MultiModalKwargsItem`. |
| 263 | device: The device to place the grouped tensors on. |
| 264 | pin_memory: Whether to pin memory for faster host-to-device transfer. |
| 265 | |
| 266 | Yields: |
| 267 | A tuple `(num_items, grouped_kwargs)`, where: |
| 268 | - `kwargs` is a dictionary of keyword arguments to pass to the model; |
| 269 | - `num_items` is the corresponding number of items. |
| 270 | |
| 271 | """ |
| 272 | start_idx = 0 |
| 273 | for end_idx in range(1, len(items) + 1): |
| 274 | if end_idx < len(items) and _can_batch_mm_items( |
| 275 | items[end_idx - 1], items[end_idx] |
| 276 | ): |
| 277 | continue |
| 278 | |
| 279 | group_data = _batch_mm_items( |
| 280 | items[start_idx:end_idx], |
| 281 | device=device, |
| 282 | pin_memory=pin_memory, |
| 283 | ) |
| 284 | |
| 285 | yield end_idx - start_idx, group_data |
| 286 | start_idx = end_idx |
| 287 | |
| 288 | assert start_idx == len(items) |
| 289 | |
| 290 | |
| 291 | def group_and_batch_mm_kwargs( |
| 292 | mm_kwargs: list[tuple[str, MultiModalKwargsItem]], |
| 293 | *, |
| 294 | device: torch.types.Device = None, |
| 295 | pin_memory: bool = False, |
| 296 | ) -> Generator[tuple[str, int, BatchedTensorInputs], None, None]: |
| 297 | """Group consecutive items (possibly from different requests) into batches. |
| 298 | |
| 299 | Items must be split across groups if any of the following occurs, |
| 300 | as the batch would otherwise be invalid: |
| 301 | - They have different fields (e.g. mixed image and embedding inputs). |
| 302 | - They have different values in `MultiModalSharedField`. |
| 303 | |
| 304 | To simplify the implementation of `embed_multimodal`, we add another |
| 305 | restriction that the items in a batch must belong to the same modality. |
| 306 | |
| 307 | Args: |
| 308 | mm_kwargs: List of `(modality, item)`. |
| 309 | device: The device to place the grouped tensors on. |
| 310 | pin_memory: Whether to pin memory for faster host-to-device transfer. |
| 311 | |
| 312 | Yields: |
| 313 | A tuple `(modality, num_items, grouped_kwargs)`, where: |
| 314 | - `modality` is the modality of the batch; |
| 315 | - `kwargs` is a dictionary of keyword arguments to pass to the model; |
| 316 | - `num_items` is the corresponding number of items. |
| 317 | |
| 318 | """ |
| 319 | for modality, group in groupby(mm_kwargs, key=lambda x: x[0]): |
| 320 | items_lst = [item for _, item in group] |
| 321 | |
| 322 | for num_items, mm_kwargs_batch in group_and_batch_mm_items( |
| 323 | items_lst, |
| 324 | device=device, |
| 325 | pin_memory=pin_memory, |
| 326 | ): |
| 327 | yield modality, num_items, mm_kwargs_batch |
| 328 | |
| 329 | |
| 330 | def fetch_audio( |
| 331 | audio_url: str, |
| 332 | audio_io_kwargs: dict[str, Any] | None = None, |
| 333 | ) -> tuple[np.ndarray, int | float]: |
| 334 | """Args: |
| 335 | audio_url: URL of the audio file to fetch. |
| 336 | audio_io_kwargs: Additional kwargs passed to handle audio IO. |
| 337 | |
| 338 | Warning: |
| 339 | This method has direct access to local files and is only intended |
| 340 | to be called by user code. Never call this from the online server! |
| 341 | |
| 342 | """ |
| 343 | media_io_kwargs = None if not audio_io_kwargs else {"audio": audio_io_kwargs} |
| 344 | media_connector = MediaConnector( |
| 345 | media_io_kwargs=media_io_kwargs, |
| 346 | allowed_local_media_path="/", |
| 347 | ) |
| 348 | return media_connector.fetch_audio(audio_url) |
| 349 | |
| 350 | |
| 351 | def fetch_image( |
| 352 | image_url: str, |
| 353 | image_io_kwargs: dict[str, Any] | None = None, |
| 354 | ) -> Image.Image: |
| 355 | """Args: |
| 356 | image_url: URL of the image file to fetch. |
| 357 | image_io_kwargs: Additional kwargs passed to handle image IO. |
| 358 | |
| 359 | Warning: |
| 360 | This method has direct access to local files and is only intended |
| 361 | to be called by user code. Never call this from the online server! |
| 362 | |
| 363 | """ |
| 364 | media_io_kwargs = None if not image_io_kwargs else {"image": image_io_kwargs} |
| 365 | media_connector = MediaConnector( |
| 366 | media_io_kwargs=media_io_kwargs, |
| 367 | allowed_local_media_path="/", |
| 368 | ) |
| 369 | return media_connector.fetch_image(image_url) |
| 370 | |
| 371 | |
| 372 | def fetch_video( |
| 373 | video_url: str, |
| 374 | video_io_kwargs: dict[str, Any] | None = None, |
| 375 | ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]: |
| 376 | """Args: |
| 377 | video_url: URL of the video file to fetch. |
| 378 | video_io_kwargs: Additional kwargs passed to handle video IO. |
| 379 | |
| 380 | Warning: |
| 381 | This method has direct access to local files and is only intended |
| 382 | to be called by user code. Never call this from the online server! |
| 383 | |
| 384 | """ |
| 385 | media_io_kwargs = None if not video_io_kwargs else {"video": video_io_kwargs} |
| 386 | media_connector = MediaConnector( |
| 387 | media_io_kwargs=media_io_kwargs, |
| 388 | allowed_local_media_path="/", |
| 389 | ) |
| 390 | return media_connector.fetch_video(video_url) |
| 391 | |
| 392 | |
| 393 | def set_mm_embedding_modality(embed: "torch.Tensor", modality: str) -> "torch.Tensor": |
| 394 | """Attach modality metadata to a gathered multimodal embedding tensor. |
| 395 | |
| 396 | Used by interleaved Omni merge paths that need to group embeddings by |
| 397 | modality without threading a parallel modalities list through |
| 398 | ``embed_input_ids``. |
| 399 | """ |
| 400 | embed.modality = modality # type: ignore[attr-defined] |
| 401 | return embed |
| 402 | |
| 403 | |
| 404 | def copy_mm_embedding_modality( |
| 405 | src: "torch.Tensor", dst: "torch.Tensor" |
| 406 | ) -> "torch.Tensor": |
| 407 | """Copy ``modality`` from ``src`` onto ``dst`` if present.""" |
| 408 | modality = getattr(src, "modality", None) |
| 409 | if modality is not None: |
| 410 | dst.modality = modality # type: ignore[attr-defined] |
| 411 | return dst |
| 412 | |
| 413 | |
| 414 | def get_mm_embedding_modalities( |
| 415 | multimodal_embeddings: Sequence["torch.Tensor"], |
| 416 | ) -> list[str]: |
| 417 | """Collect per-embedding modalities previously set on the tensors.""" |
| 418 | modalities: list[str] = [] |
| 419 | for i, emb in enumerate(multimodal_embeddings): |
| 420 | modality = getattr(emb, "modality", None) |
| 421 | if modality is None: |
| 422 | raise ValueError( |
| 423 | f"Missing modality on multimodal embedding at index {i}. " |
| 424 | "Encoder gather must set embed.modality before interleaved " |
| 425 | "audio-in-video merge." |
| 426 | ) |
| 427 | modalities.append(modality) |
| 428 | return modalities |