| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import ABC, abstractmethod |
| 5 | from collections import UserDict |
| 6 | from collections.abc import Callable, Iterator, Mapping, Sequence, Set |
| 7 | from typing import ( |
| 8 | TYPE_CHECKING, |
| 9 | Any, |
| 10 | Generic, |
| 11 | Literal, |
| 12 | NamedTuple, |
| 13 | TypeAlias, |
| 14 | TypeGuard, |
| 15 | TypeVar, |
| 16 | ) |
| 17 | |
| 18 | import numpy as np |
| 19 | import torch |
| 20 | from typing_extensions import assert_never |
| 21 | |
| 22 | from vllm.inputs import ModalityData, MultiModalDataDict, MultiModalUUIDDict |
| 23 | from vllm.utils.collection_utils import is_list_of |
| 24 | from vllm.utils.import_utils import LazyLoader |
| 25 | |
| 26 | from .audio import AudioResampler, AudioSpec, normalize_audio |
| 27 | from .image import convert_image_mode, normalize_image |
| 28 | from .inputs import ( |
| 29 | AudioItem, |
| 30 | HfAudioItem, |
| 31 | HfImageItem, |
| 32 | HfVideoItem, |
| 33 | ImageItem, |
| 34 | MultiModalFieldConfig, |
| 35 | MultiModalKwargsItems, |
| 36 | VideoItem, |
| 37 | ) |
| 38 | from .media import MediaWithBytes |
| 39 | from .video import DecodedFrames |
| 40 | |
| 41 | _T = TypeVar("_T") |
| 42 | _I = TypeVar("_I") |
| 43 | |
| 44 | EmbeddingFieldRole: TypeAlias = Literal["values", "metadata"] |
| 45 | """What a field of a pre-computed-embedding input carries. |
| 46 | |
| 47 | `"values"` is the embedding tensor itself; `"metadata"` is a key that sizes the |
| 48 | prompt's placeholder range. The distinction is what lets one declaration serve |
| 49 | both an ordinary request (everything present) and an EC consumer (embeddings |
| 50 | arrive through the connector, metadata still in the request). |
| 51 | """ |
| 52 | |
| 53 | if TYPE_CHECKING: |
| 54 | import PIL.Image as PILImage |
| 55 | |
| 56 | else: |
| 57 | PILImage = LazyLoader("PILImage", globals(), "PIL.Image") |
| 58 | |
| 59 | |
| 60 | class ModalityDataItems(ABC, Generic[_T, _I]): |
| 61 | """Represents data items for a modality in |
| 62 | [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems]. |
| 63 | """ |
| 64 | |
| 65 | def __init__(self, data: _T, modality: str) -> None: |
| 66 | super().__init__() |
| 67 | |
| 68 | self.data: _T = data |
| 69 | self.modality = modality |
| 70 | |
| 71 | def __repr__(self) -> str: |
| 72 | return f"{type(self).__name__}(modality={self.modality!r}, len={len(self)})" |
| 73 | |
| 74 | def __len__(self) -> int: |
| 75 | return self.get_count() |
| 76 | |
| 77 | def __getitem__(self, index: int) -> _I: |
| 78 | return self.get(index) |
| 79 | |
| 80 | if TYPE_CHECKING: |
| 81 | # Auto-generated |
| 82 | def __iter__(self) -> Iterator[_I]: ... |
| 83 | |
| 84 | @abstractmethod |
| 85 | def get_count(self) -> int: |
| 86 | """Get the number of data items.""" |
| 87 | raise NotImplementedError |
| 88 | |
| 89 | @abstractmethod |
| 90 | def get(self, index: int) -> _I: |
| 91 | """Get a data item by its index.""" |
| 92 | raise NotImplementedError |
| 93 | |
| 94 | def get_all(self) -> list[_I]: |
| 95 | """Get all data items.""" |
| 96 | return [self.get(idx) for idx in range(self.get_count())] |
| 97 | |
| 98 | def get_item_for_hash(self, index: int) -> object: |
| 99 | return self.get(index) |
| 100 | |
| 101 | def get_all_items_for_hash(self) -> list[object]: |
| 102 | return [self.get_item_for_hash(idx) for idx in range(self.get_count())] |
| 103 | |
| 104 | @abstractmethod |
| 105 | def get_processor_data(self) -> Mapping[str, object]: |
| 106 | """Get the data to pass to the HF processor.""" |
| 107 | raise NotImplementedError |
| 108 | |
| 109 | @abstractmethod |
| 110 | def get_passthrough_data(self) -> Mapping[str, object]: |
| 111 | """Get the data to pass directly to the model.""" |
| 112 | raise NotImplementedError |
| 113 | |
| 114 | |
| 115 | class ProcessorBatchItems(ModalityDataItems[Sequence[_T], _T]): |
| 116 | """Base class for data items that are arranged in a list.""" |
| 117 | |
| 118 | def _unwrap(self, item: _T | MediaWithBytes[_T]) -> _T: |
| 119 | """Extract media from wrapper if present.""" |
| 120 | return item.media if isinstance(item, MediaWithBytes) else item |
| 121 | |
| 122 | def get_count(self) -> int: |
| 123 | return len(self.data) |
| 124 | |
| 125 | def get(self, index: int) -> _T: |
| 126 | return self._unwrap(self.data[index]) |
| 127 | |
| 128 | def get_item_for_hash(self, index: int) -> _T | MediaWithBytes[_T]: |
| 129 | # Return raw item for hashing (preserves original_bytes if present) |
| 130 | return self.data[index] |
| 131 | |
| 132 | def get_processor_data(self) -> Mapping[str, object]: |
| 133 | return {f"{self.modality}s": self.get_all()} |
| 134 | |
| 135 | def get_passthrough_data(self) -> Mapping[str, object]: |
| 136 | return {} |
| 137 | |
| 138 | |
| 139 | def validate_embedding_ndim( |
| 140 | tensor: torch.Tensor, |
| 141 | modality: str, |
| 142 | index: int | None = None, |
| 143 | ) -> None: |
| 144 | """Validate tensor ndim for multimodal embeddings. |
| 145 | |
| 146 | Single embeddings should be 2D (seq_len, hidden_size). |
| 147 | Batched embeddings should be 3D (batch, seq_len, hidden_size). |
| 148 | |
| 149 | Args: |
| 150 | tensor: The tensor to validate. |
| 151 | modality: The modality name for error messages (e.g., "image", "audio"). |
| 152 | index: Optional index for list items, included in error messages. |
| 153 | |
| 154 | """ |
| 155 | if tensor.ndim < 2 or tensor.ndim > 3: |
| 156 | idx_str = f" [{index}]" if index is not None else "" |
| 157 | raise ValueError( |
| 158 | f"{modality.capitalize()} embedding{idx_str} must be 2D " |
| 159 | f"(seq_len, hidden_size) or 3D (batch, seq_len, hidden_size), " |
| 160 | f"got {tensor.ndim}D tensor with shape {tuple(tensor.shape)}" |
| 161 | ) |
| 162 | |
| 163 | |
| 164 | class EmbeddingItems( |
| 165 | ModalityDataItems[torch.Tensor | list[torch.Tensor], torch.Tensor] |
| 166 | ): |
| 167 | """Base class for data items that are expressed as a batched embedding tensor, |
| 168 | or a list of embedding tensors (one per item). |
| 169 | """ |
| 170 | |
| 171 | def __init__( |
| 172 | self, |
| 173 | data: torch.Tensor | list[torch.Tensor], |
| 174 | modality: str, |
| 175 | expected_hidden_size: int | None = None, |
| 176 | ) -> None: |
| 177 | super().__init__(data, modality) |
| 178 | |
| 179 | # Validate ndim first (before hidden_size which depends on correct ndim) |
| 180 | self._validate_ndim() |
| 181 | |
| 182 | # Validate hidden dimension if expected size is provided |
| 183 | if expected_hidden_size is not None: |
| 184 | self._validate_hidden_size(expected_hidden_size) |
| 185 | |
| 186 | def _validate_ndim(self) -> None: |
| 187 | """Validate that embedding tensors have correct ndim (2D or 3D).""" |
| 188 | if isinstance(self.data, torch.Tensor): |
| 189 | validate_embedding_ndim(self.data, self.modality) |
| 190 | else: |
| 191 | # List of tensors: each should be 2D (seq_len, hidden_size) |
| 192 | for idx, tensor in enumerate(self.data): |
| 193 | if tensor.ndim != 2: |
| 194 | raise ValueError( |
| 195 | f"{self.modality.capitalize()} embedding [{idx}] must be " |
| 196 | f"2D (seq_len, hidden_size), got {tensor.ndim}D tensor " |
| 197 | f"with shape {tuple(tensor.shape)}" |
| 198 | ) |
| 199 | |
| 200 | def _validate_hidden_size(self, expected_hidden_size: int) -> None: |
| 201 | """Validate that embedding hidden dimension matches expected size. |
| 202 | |
| 203 | This validates hidden dimensions to prevent vulnerabilities: Embeddings |
| 204 | with correct ndim but wrong hidden dimension could bypass initial |
| 205 | checks and cause crashes during model inference when dimensions don't match. |
| 206 | """ |
| 207 | if isinstance(self.data, torch.Tensor): |
| 208 | # Batched tensor: shape is (batch, seq_len, hidden_size) |
| 209 | actual_hidden_size = self.data.shape[-1] |
| 210 | if actual_hidden_size != expected_hidden_size: |
| 211 | raise ValueError( |
| 212 | f"{self.modality.capitalize()} embedding hidden dimension " |
| 213 | f"mismatch: got {actual_hidden_size}, but model expects " |
| 214 | f"{expected_hidden_size}. Embedding shape: {tuple(self.data.shape)}" |
| 215 | ) |
| 216 | else: |
| 217 | # List of tensors: each has shape (seq_len, hidden_size) |
| 218 | for idx, tensor in enumerate(self.data): |
| 219 | actual_hidden_size = tensor.shape[-1] |
| 220 | if actual_hidden_size != expected_hidden_size: |
| 221 | raise ValueError( |
| 222 | f"{self.modality.capitalize()} embedding [{idx}] hidden " |
| 223 | f"dimension mismatch: got {actual_hidden_size}, but model " |
| 224 | f"expects {expected_hidden_size}. " |
| 225 | f"Embedding shape: {tuple(tensor.shape)}" |
| 226 | ) |
| 227 | |
| 228 | def _unwrap( |
| 229 | self, item: torch.Tensor | MediaWithBytes[torch.Tensor] |
| 230 | ) -> torch.Tensor: |
| 231 | """Extract media from wrapper if present.""" |
| 232 | return item.media if isinstance(item, MediaWithBytes) else item |
| 233 | |
| 234 | def get_count(self) -> int: |
| 235 | return len(self.data) |
| 236 | |
| 237 | def get(self, index: int) -> torch.Tensor: |
| 238 | return self._unwrap(self.data[index]) |
| 239 | |
| 240 | def get_processor_data(self) -> Mapping[str, object]: |
| 241 | return {} |
| 242 | |
| 243 | def get_passthrough_data(self) -> Mapping[str, object]: |
| 244 | return {f"{self.modality}_embeds": self.data} |
| 245 | |
| 246 | def get_feature_size(self, item_idx: int) -> int: |
| 247 | return len(self.get(item_idx)) |
| 248 | |
| 249 | |
| 250 | class DictEmbeddingItems( |
| 251 | ModalityDataItems[Mapping[str, torch.Tensor], Mapping[str, torch.Tensor]] |
| 252 | ): |
| 253 | """Base class for data items that are expressed as a dictionary of tensors. |
| 254 | |
| 255 | Usually, the dictionary keys correspond to the outputs of HF processor. |
| 256 | """ |
| 257 | |
| 258 | def __init__( |
| 259 | self, |
| 260 | data: Mapping[str, torch.Tensor], |
| 261 | modality: str, |
| 262 | required_fields: Set[str], |
| 263 | fields_factory: Callable[ |
| 264 | [Mapping[str, torch.Tensor]], |
| 265 | Mapping[str, MultiModalFieldConfig], |
| 266 | ], |
| 267 | optional_fields: Set[str] = frozenset(), |
| 268 | ) -> None: |
| 269 | """Args: |
| 270 | data: The dictionary of tensors for this modality. |
| 271 | modality: The modality these items belong to. |
| 272 | required_fields: Fields `data` must contain. |
| 273 | fields_factory: Builds the field config from the data. |
| 274 | optional_fields: Fields `data` may omit. Which fields these are is |
| 275 | the caller's decision -- see |
| 276 | `MultiModalDataParser.embedding_field_sets`, where a deployment |
| 277 | that receives embeddings through an EC connector makes the |
| 278 | embeddings optional. They still need a field config, since they |
| 279 | are used whenever they *are* supplied. |
| 280 | |
| 281 | """ |
| 282 | from transformers.feature_extraction_utils import BatchFeature |
| 283 | |
| 284 | super().__init__(data, modality) |
| 285 | |
| 286 | # Nothing required would leave nothing to size the placeholder range |
| 287 | # from, so the item would parse into zero entries and silently produce a |
| 288 | # wrong prompt instead of failing here. |
| 289 | if not required_fields: |
| 290 | raise ValueError( |
| 291 | f"Cannot parse {modality!r} embeddings: every declared field is " |
| 292 | f"optional ({sorted(optional_fields)}), so nothing is left to " |
| 293 | "size the placeholder range." |
| 294 | ) |
| 295 | |
| 296 | declared_fields = set(required_fields) | set(optional_fields) |
| 297 | |
| 298 | missing_required_data_keys = set(required_fields) - data.keys() |
| 299 | if missing_required_data_keys: |
| 300 | data_keys = set(data.keys()) |
| 301 | msg = ( |
| 302 | f"The data should contain the fields: {set(required_fields)}, " |
| 303 | f"but only found the following keys: {data_keys}" |
| 304 | ) |
| 305 | raise ValueError(msg) |
| 306 | |
| 307 | fields_config = fields_factory(data) |
| 308 | # Check every declared field, not just the required ones: an optional |
| 309 | # field still needs a config for when it is supplied. |
| 310 | missing_required_fields = declared_fields - fields_config.keys() |
| 311 | if missing_required_fields: |
| 312 | fields = set(fields_config.keys()) |
| 313 | msg = f"{declared_fields=} should be a subset of {fields=}" |
| 314 | raise ValueError(msg) |
| 315 | |
| 316 | self.fields_config = fields_config |
| 317 | self.required_fields = set(required_fields) |
| 318 | |
| 319 | self._kwargs = MultiModalKwargsItems.from_hf_inputs( |
| 320 | BatchFeature(dict(data)), |
| 321 | fields_config, |
| 322 | ) |
| 323 | |
| 324 | def get_count(self) -> int: |
| 325 | return len(self._kwargs[self.modality]) |
| 326 | |
| 327 | def get(self, index: int) -> Mapping[str, torch.Tensor]: |
| 328 | return self._kwargs[self.modality][index].get_data() |
| 329 | |
| 330 | def get_processor_data(self) -> Mapping[str, object]: |
| 331 | return {} |
| 332 | |
| 333 | def get_passthrough_data(self) -> Mapping[str, object]: |
| 334 | return self.data |
| 335 | |
| 336 | |
| 337 | class AudioProcessorItems(ProcessorBatchItems[HfAudioItem | None]): |
| 338 | def __init__(self, data: Sequence[HfAudioItem | None]) -> None: |
| 339 | super().__init__(data, "audio") |
| 340 | |
| 341 | def get_audio_length(self, item_idx: int) -> int: |
| 342 | audio = self.get(item_idx) |
| 343 | if audio is None: |
| 344 | raise ValueError(f"Cannot get length of cached audio at {item_idx}") |
| 345 | |
| 346 | return len(audio) |
| 347 | |
| 348 | |
| 349 | class AudioEmbeddingItems(EmbeddingItems): |
| 350 | def __init__( |
| 351 | self, |
| 352 | data: torch.Tensor | list[torch.Tensor], |
| 353 | expected_hidden_size: int | None = None, |
| 354 | ) -> None: |
| 355 | super().__init__(data, "audio", expected_hidden_size) |
| 356 | |
| 357 | |
| 358 | class ImageSize(NamedTuple): |
| 359 | width: int |
| 360 | height: int |
| 361 | |
| 362 | |
| 363 | class ImageProcessorItems(ProcessorBatchItems[HfImageItem | None]): |
| 364 | def __init__(self, data: Sequence[HfImageItem | None]) -> None: |
| 365 | super().__init__(data, "image") |
| 366 | |
| 367 | def get_image_size(self, item_idx: int) -> ImageSize: |
| 368 | image = self.get(item_idx) |
| 369 | if image is None: |
| 370 | raise ValueError(f"Cannot get size of cached image at {item_idx}") |
| 371 | |
| 372 | if isinstance(image, PILImage.Image): |
| 373 | return ImageSize(*image.size) |
| 374 | if isinstance(image, (np.ndarray, torch.Tensor)): |
| 375 | if image.ndim == 3 and image.shape[-1] in (1, 3, 4): |
| 376 | # HWC format (e.g. from np.array(PIL.Image)). |
| 377 | # PIL images are always channels-last. |
| 378 | h, w = image.shape[0], image.shape[1] |
| 379 | else: |
| 380 | # CHW format (standard PyTorch / numpy convention). |
| 381 | _, h, w = image.shape |
| 382 | return ImageSize(w, h) |
| 383 | |
| 384 | assert_never(image) |
| 385 | |
| 386 | |
| 387 | class ImageEmbeddingItems(EmbeddingItems): |
| 388 | def __init__( |
| 389 | self, |
| 390 | data: torch.Tensor | list[torch.Tensor], |
| 391 | expected_hidden_size: int | None = None, |
| 392 | ) -> None: |
| 393 | super().__init__(data, "image", expected_hidden_size) |
| 394 | |
| 395 | |
| 396 | class VideoProcessorItems(ProcessorBatchItems[HfVideoItem | None]): |
| 397 | def __init__( |
| 398 | self, |
| 399 | data: Sequence[HfVideoItem | None], |
| 400 | metadata: dict[str, Any] | list[dict[str, Any] | None] | None = None, |
| 401 | ) -> None: |
| 402 | super().__init__(data, "video") |
| 403 | |
| 404 | self.metadata = metadata |
| 405 | |
| 406 | def _unwrap(self, item: Any) -> Any: |
| 407 | if isinstance(item, tuple): |
| 408 | frames, metadata = item |
| 409 | return super()._unwrap(frames), metadata |
| 410 | return super()._unwrap(item) |
| 411 | |
| 412 | def get_item_for_hash(self, index: int) -> Any: |
| 413 | item = self.data[index] |
| 414 | if isinstance(item, MediaWithBytes) and isinstance(self.metadata, list): |
| 415 | metadata = self.metadata[index] |
| 416 | if metadata is not None: |
| 417 | return item, metadata |
| 418 | return item |
| 419 | |
| 420 | def get_num_frames(self, item_idx: int) -> int: |
| 421 | video = self.get(item_idx) |
| 422 | if video is None: |
| 423 | raise ValueError(f"Cannot get length of cached video at {item_idx}") |
| 424 | |
| 425 | return len(video) |
| 426 | |
| 427 | def get_frame_size(self, item_idx: int) -> ImageSize: |
| 428 | video = self.get(item_idx) |
| 429 | if video is None: |
| 430 | raise ValueError(f"Cannot get size of cached video at {item_idx}") |
| 431 | if len(video) == 0: |
| 432 | raise ValueError(f"Cannot get size of empty video at {item_idx}") |
| 433 | |
| 434 | image = video[0] |
| 435 | |
| 436 | if isinstance(image, PILImage.Image): |
| 437 | return ImageSize(*image.size) |
| 438 | if isinstance(image, (np.ndarray, torch.Tensor)): |
| 439 | if image.ndim == 3 and image.shape[-1] in (1, 3, 4): |
| 440 | # HWC format (e.g. from np.array(PIL.Image) via |
| 441 | # _get_video_with_metadata). PIL images are always |
| 442 | # channels-last. |
| 443 | h, w = image.shape[0], image.shape[1] |
| 444 | else: |
| 445 | # CHW format (standard PyTorch / numpy convention). |
| 446 | _, h, w = image.shape |
| 447 | return ImageSize(w, h) |
| 448 | |
| 449 | assert_never(image) |
| 450 | |
| 451 | |
| 452 | class VideoEmbeddingItems(EmbeddingItems): |
| 453 | def __init__( |
| 454 | self, |
| 455 | data: torch.Tensor | list[torch.Tensor], |
| 456 | expected_hidden_size: int | None = None, |
| 457 | ) -> None: |
| 458 | super().__init__(data, "video", expected_hidden_size) |
| 459 | |
| 460 | |
| 461 | class VisionChunkProcessorItems(ProcessorBatchItems[Any]): |
| 462 | """Processor items for vision chunks (unified image and video chunks).""" |
| 463 | |
| 464 | def __init__(self, data: Sequence[Any]) -> None: |
| 465 | super().__init__(data, "vision_chunk") |
| 466 | |
| 467 | |
| 468 | _D = TypeVar("_D", bound=ModalityDataItems[Any, Any]) |
| 469 | |
| 470 | |
| 471 | class MultiModalDataItems(UserDict[str, ModalityDataItems[Any, Any]]): |
| 472 | """A normalized [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] |
| 473 | such that each entry corresponds to a list. |
| 474 | """ |
| 475 | |
| 476 | def select(self, modalities: Set[str]): |
| 477 | """Construct a new `MultiModalDataItems` instance containing only the |
| 478 | selected modalities. |
| 479 | """ |
| 480 | return MultiModalDataItems( |
| 481 | {modality: self[modality] for modality in modalities} |
| 482 | ) |
| 483 | |
| 484 | def get_count(self, modality: str, *, strict: bool = True) -> int: |
| 485 | """Get the number of data items belonging to a modality. |
| 486 | |
| 487 | If `strict=False`, return `0` instead of raising [`KeyError`][] |
| 488 | even if the modality is not found. |
| 489 | """ |
| 490 | if modality not in self: |
| 491 | if strict: |
| 492 | available_modalities = set(self.keys()) |
| 493 | raise KeyError( |
| 494 | f"Modality {modality!r} not found. " |
| 495 | f"Available modalities: {available_modalities}" |
| 496 | ) |
| 497 | |
| 498 | return 0 |
| 499 | |
| 500 | return self[modality].get_count() |
| 501 | |
| 502 | def get_all_counts(self) -> Mapping[str, int]: |
| 503 | """Get the number of items belonging to each modality.""" |
| 504 | return {m: items.get_count() for m, items in self.items()} |
| 505 | |
| 506 | def get_items( |
| 507 | self, |
| 508 | modality: str, |
| 509 | typ: type[_D] | tuple[type[_D], ...], |
| 510 | ) -> _D: |
| 511 | """Get the data items belonging to a modality, |
| 512 | requiring that they belong to a certain type. |
| 513 | """ |
| 514 | if modality not in self: |
| 515 | available_modalities = set(self.keys()) |
| 516 | raise KeyError( |
| 517 | f"Modality {modality!r} not found. " |
| 518 | f"Available modalities: {available_modalities}" |
| 519 | ) |
| 520 | |
| 521 | items = self[modality] |
| 522 | if not isinstance(items, typ): |
| 523 | raise TypeError( |
| 524 | f"Invalid type of data items for {modality=}. " |
| 525 | f"Expected type: {typ}, but " |
| 526 | f"found type: {type(items)}" |
| 527 | ) |
| 528 | |
| 529 | return items # type: ignore[return-value] |
| 530 | |
| 531 | |
| 532 | ModalityDataParser: TypeAlias = Callable[ |
| 533 | [ModalityData[Any]], ModalityDataItems[Any, Any] | None |
| 534 | ] |
| 535 | |
| 536 | |
| 537 | class MultiModalDataParser: |
| 538 | """Parses [`MultiModalDataDict`][vllm.inputs.MultiModalDataDict] |
| 539 | into [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems]. |
| 540 | |
| 541 | Args: |
| 542 | target_sr (float, optional): Enables automatic resampling of audio |
| 543 | items to the model's expected sampling rate. |
| 544 | audio_resample_method (str): Backend used for the resampling above. |
| 545 | Defaults to torchaudio; models with specific needs may override |
| 546 | (e.g. phi4mm uses scipy). |
| 547 | target_channels (int, optional): Target number of audio channels. |
| 548 | If provided, normalizes audio to this many channels (e.g., 1 for mono). |
| 549 | If None, audio channels are passed through unchanged. |
| 550 | expected_hidden_size (int, optional): Expected hidden dimension for |
| 551 | embedding inputs. If provided, validates that user-supplied |
| 552 | embeddings have the correct hidden size to prevent crashes |
| 553 | during model inference. |
| 554 | allow_missing_mm_embeddings (bool): Whether pre-computed embedding |
| 555 | tensors may be absent from the request on a disaggregated consumer. |
| 556 | Derived by `BaseProcessingInfo.allow_missing_mm_embeddings`. |
| 557 | |
| 558 | """ |
| 559 | |
| 560 | embedding_fields: Mapping[str, Mapping[str, EmbeddingFieldRole]] = {} |
| 561 | """Per-modality field roles for pre-computed-embedding inputs. |
| 562 | |
| 563 | Declared here rather than inside `_parse_*_data` so an EC producer can read |
| 564 | it too: on a producer the request carries real media, so the branch that |
| 565 | builds `DictEmbeddingItems` never runs, yet the producer still has to know |
| 566 | which processed keys to publish. One declaration, both sides, no drift. |
| 567 | |
| 568 | A modality absent from this mapping can only be sent whole in the request. |
| 569 | """ |
| 570 | |
| 571 | @classmethod |
| 572 | def placeholder_metadata_fields(cls, modality: str) -> set[str]: |
| 573 | """The keys that size `modality`'s placeholder range. |
| 574 | |
| 575 | What an EC producer publishes alongside the embedding, and what a |
| 576 | consumer keeps requiring once the embedding itself is gone. |
| 577 | """ |
| 578 | return { |
| 579 | field |
| 580 | for field, role in cls.embedding_fields.get(modality, {}).items() |
| 581 | if role == "metadata" |
| 582 | } |
| 583 | |
| 584 | def embedding_field_sets(self, modality: str) -> tuple[set[str], set[str]]: |
| 585 | """`modality`'s (required, optional) fields for this deployment. |
| 586 | |
| 587 | Resolves the static roles in `embedding_fields` against where the |
| 588 | embeddings actually come from: on an EC consumer they arrive through the |
| 589 | connector, so the request may omit them; anywhere else a request that |
| 590 | claims to carry pre-computed embeddings has to actually carry them. |
| 591 | """ |
| 592 | metadata = self.placeholder_metadata_fields(modality) |
| 593 | values = set(self.embedding_fields.get(modality, {})) - metadata |
| 594 | if self.allow_missing_mm_embeddings: |
| 595 | return metadata, values |
| 596 | return metadata | values, set() |
| 597 | |
| 598 | def __init__( |
| 599 | self, |
| 600 | *, |
| 601 | target_sr: float | None = None, |
| 602 | target_channels: int | None = None, |
| 603 | audio_resample_method: Literal["pyav", "scipy", "soxr", "torchaudio"] = ( |
| 604 | "torchaudio" |
| 605 | ), |
| 606 | video_needs_metadata: bool = False, |
| 607 | expected_hidden_size: int | None = None, |
| 608 | allow_missing_mm_embeddings: bool = False, |
| 609 | ) -> None: |
| 610 | super().__init__() |
| 611 | |
| 612 | self.allow_missing_mm_embeddings = allow_missing_mm_embeddings |
| 613 | |
| 614 | self.audio_resampler = AudioResampler( |
| 615 | target_sr=target_sr, |
| 616 | method=audio_resample_method, |
| 617 | ) |
| 618 | self.target_channels = target_channels |
| 619 | self.video_needs_metadata = video_needs_metadata |
| 620 | self.expected_hidden_size = expected_hidden_size |
| 621 | |
| 622 | @classmethod |
| 623 | def is_embeddings( |
| 624 | cls, data: object |
| 625 | ) -> TypeGuard[torch.Tensor | list[torch.Tensor]]: |
| 626 | if isinstance(data, torch.Tensor): |
| 627 | return data.ndim == 3 |
| 628 | if is_list_of(data, torch.Tensor) and len(data) > 0: |
| 629 | return data[0].ndim == 2 # type: ignore[index] |
| 630 | |
| 631 | return False |
| 632 | |
| 633 | def _get_audio_with_sr( |
| 634 | self, |
| 635 | audio: AudioItem, |
| 636 | ) -> tuple[np.ndarray, float | None]: |
| 637 | if isinstance(audio, tuple): |
| 638 | return audio |
| 639 | if isinstance(audio, list): |
| 640 | return np.array(audio), None |
| 641 | if isinstance(audio, np.ndarray): |
| 642 | return audio, None |
| 643 | if isinstance(audio, torch.Tensor): |
| 644 | return audio.numpy(), None |
| 645 | |
| 646 | assert_never(audio) |
| 647 | |
| 648 | def _get_video_with_metadata( |
| 649 | self, |
| 650 | video: VideoItem, |
| 651 | ) -> tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any] | None]: |
| 652 | if isinstance(video, MediaWithBytes): |
| 653 | new_video, metadata = self._get_video_with_metadata(video.media) |
| 654 | return MediaWithBytes(new_video, video.original_bytes), metadata |
| 655 | if isinstance(video, tuple): |
| 656 | return video |
| 657 | if isinstance(video, list): |
| 658 | return np.array(video), None |
| 659 | if isinstance(video, (np.ndarray, torch.Tensor)): |
| 660 | # Tensors pass through untouched: HF video processors accept them |
| 661 | # directly, and device tensors (e.g. NVDEC-decoded frames) must |
| 662 | # stay on-device to avoid a D2H round-trip. |
| 663 | return video, None |
| 664 | |
| 665 | assert_never(video) |
| 666 | |
| 667 | def _parse_audio_embedding_data( |
| 668 | self, data: dict[str, torch.Tensor] |
| 669 | ) -> DictEmbeddingItems: |
| 670 | counts = data.get("audio_num_tokens") |
| 671 | if ( |
| 672 | not isinstance(counts, torch.Tensor) |
| 673 | or counts.dtype not in (torch.int32, torch.int64) |
| 674 | or counts.ndim not in (1, 2) |
| 675 | or (counts.ndim == 2 and counts.shape[1] != 1) |
| 676 | or not torch.all(counts > 0) |
| 677 | ): |
| 678 | raise ValueError( |
| 679 | "audio_num_tokens must contain one positive integer per audio" |
| 680 | ) |
| 681 | |
| 682 | counts = counts.flatten() |
| 683 | if "audio_embeds" in data: |
| 684 | embeds = data["audio_embeds"] |
| 685 | if len(embeds) != len(counts) or any( |
| 686 | len(embedding) != count for embedding, count in zip(embeds, counts) |
| 687 | ): |
| 688 | raise ValueError("audio_num_tokens does not match audio_embeds") |
| 689 | |
| 690 | required, optional = self.embedding_field_sets("audio") |
| 691 | return DictEmbeddingItems( |
| 692 | {**data, "audio_num_tokens": counts}, |
| 693 | modality="audio", |
| 694 | required_fields=required, |
| 695 | optional_fields=optional, |
| 696 | fields_factory=lambda _: { |
| 697 | "audio_embeds": MultiModalFieldConfig.batched("audio"), |
| 698 | "audio_num_tokens": MultiModalFieldConfig.batched( |
| 699 | "audio", keep_on_cpu=True |
| 700 | ), |
| 701 | }, |
| 702 | ) |
| 703 | |
| 704 | def _parse_audio_data( |
| 705 | self, |
| 706 | data: dict[str, torch.Tensor] | ModalityData[AudioItem], |
| 707 | ) -> ModalityDataItems[Any, Any] | None: |
| 708 | if data is None: |
| 709 | return None |
| 710 | |
| 711 | if self.is_embeddings(data): |
| 712 | return AudioEmbeddingItems(data, self.expected_hidden_size) |
| 713 | |
| 714 | if isinstance(data, dict) and "audio_num_tokens" in self.embedding_fields.get( |
| 715 | "audio", {} |
| 716 | ): |
| 717 | return self._parse_audio_embedding_data(data) |
| 718 | |
| 719 | data_items: list[AudioItem | None] |
| 720 | if ( |
| 721 | (is_list_of(data, float) and len(data) > 0) |
| 722 | or (isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 1) |
| 723 | or isinstance(data, tuple) |
| 724 | ): |
| 725 | data_items = [data] |
| 726 | elif isinstance(data, (np.ndarray, torch.Tensor)): |
| 727 | data_items = [elem for elem in data] |
| 728 | else: |
| 729 | data_items = data # type: ignore[assignment] |
| 730 | |
| 731 | new_audios = list[np.ndarray | None]() |
| 732 | for data_item in data_items: |
| 733 | # Requests can omit audio samples when reusing a cached UUID. |
| 734 | if data_item is None: |
| 735 | new_audios.append(None) |
| 736 | continue |
| 737 | |
| 738 | audio, orig_sr = self._get_audio_with_sr(data_item) |
| 739 | if orig_sr is None: |
| 740 | new_audio = audio |
| 741 | else: |
| 742 | new_audio = self.audio_resampler.resample(audio, orig_sr=orig_sr) |
| 743 | |
| 744 | # Apply channel normalization if target_channels is set |
| 745 | if self.target_channels is not None: |
| 746 | spec = AudioSpec(target_channels=self.target_channels) |
| 747 | new_audio = normalize_audio(new_audio, spec) |
| 748 | |
| 749 | new_audios.append(new_audio) |
| 750 | |
| 751 | return AudioProcessorItems(new_audios) |
| 752 | |
| 753 | def _parse_image_data( |
| 754 | self, |
| 755 | data: ModalityData[ImageItem], |
| 756 | ) -> ModalityDataItems[Any, Any] | None: |
| 757 | if data is None: |
| 758 | return None |
| 759 | |
| 760 | if self.is_embeddings(data): |
| 761 | return ImageEmbeddingItems(data, self.expected_hidden_size) |
| 762 | |
| 763 | if isinstance(data, (PILImage.Image, MediaWithBytes)) or ( |
| 764 | isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 3 |
| 765 | ): |
| 766 | data_items = [data] |
| 767 | elif isinstance(data, (np.ndarray, torch.Tensor)): |
| 768 | data_items = [elem for elem in data] |
| 769 | else: |
| 770 | data_items = data |
| 771 | |
| 772 | data_items = [ |
| 773 | convert_image_mode(normalize_image(item), "RGB") |
| 774 | if isinstance(item, PILImage.Image) |
| 775 | else item |
| 776 | for item in data_items |
| 777 | ] |
| 778 | |
| 779 | return ImageProcessorItems(data_items) |
| 780 | |
| 781 | def _parse_video_data( |
| 782 | self, |
| 783 | data: ModalityData[VideoItem], |
| 784 | ) -> ModalityDataItems[Any, Any] | None: |
| 785 | if data is None: |
| 786 | return None |
| 787 | |
| 788 | if self.is_embeddings(data): |
| 789 | return VideoEmbeddingItems(data, self.expected_hidden_size) |
| 790 | |
| 791 | data_items: list[VideoItem] |
| 792 | if ( |
| 793 | (is_list_of(data, PILImage.Image) and len(data) > 0) |
| 794 | or ( |
| 795 | is_list_of(data, (np.ndarray, torch.Tensor), check="all") |
| 796 | and len(data) > 0 |
| 797 | and all(item.ndim == 3 for item in data) |
| 798 | ) |
| 799 | or (isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4) |
| 800 | ): |
| 801 | data_items = [data] |
| 802 | elif isinstance(data, (np.ndarray, torch.Tensor)): |
| 803 | data_items = [elem for elem in data] |
| 804 | elif isinstance(data, tuple) and len(data) == 2: |
| 805 | data_items = [data] |
| 806 | else: |
| 807 | data_items = data # type: ignore[assignment] |
| 808 | |
| 809 | new_videos = list[ |
| 810 | DecodedFrames |
| 811 | | MediaWithBytes[DecodedFrames] |
| 812 | | tuple[DecodedFrames | MediaWithBytes[DecodedFrames], dict[str, Any]] |
| 813 | | None |
| 814 | ]() |
| 815 | metadata_lst: list[dict[str, Any] | None] = [] |
| 816 | for data_item in data_items: |
| 817 | # Allow None video items, valid requests can contain empty URLs |
| 818 | # if they use multi-modal uuids. |
| 819 | if data_item is None: |
| 820 | new_videos.append(None) |
| 821 | metadata_lst.append(None) |
| 822 | continue |
| 823 | video, metadata = self._get_video_with_metadata(data_item) |
| 824 | if self.video_needs_metadata: |
| 825 | if metadata is None: |
| 826 | raise ValueError( |
| 827 | "Video metadata is required but not found in mm input. " |
| 828 | "Please check your video input in `multi_modal_data`" |
| 829 | ) |
| 830 | new_videos.append((video, metadata)) |
| 831 | else: |
| 832 | new_videos.append(video) |
| 833 | metadata_lst.append(metadata) |
| 834 | |
| 835 | return VideoProcessorItems(new_videos, metadata=metadata_lst) |
| 836 | |
| 837 | def _parse_vision_chunk_data( |
| 838 | self, |
| 839 | data: ModalityData[Any], |
| 840 | ) -> ModalityDataItems[Any, Any] | None: |
| 841 | """Parse vision chunk data (unified image and video chunks).""" |
| 842 | if data is None: |
| 843 | return None |
| 844 | |
| 845 | if self.is_embeddings(data): |
| 846 | raise ValueError("Do not support embedding data for vision_chunk right now") |
| 847 | |
| 848 | if isinstance(data, dict): |
| 849 | data = [data] |
| 850 | |
| 851 | return VisionChunkProcessorItems(data) |
| 852 | |
| 853 | def _get_subparsers(self) -> Mapping[str, ModalityDataParser]: |
| 854 | return { |
| 855 | "audio": self._parse_audio_data, |
| 856 | "image": self._parse_image_data, |
| 857 | "video": self._parse_video_data, |
| 858 | "vision_chunk": self._parse_vision_chunk_data, |
| 859 | } |
| 860 | |
| 861 | def parse_mm_data(self, mm_data: MultiModalDataDict) -> MultiModalDataItems: |
| 862 | subparsers = self._get_subparsers() |
| 863 | |
| 864 | mm_items = MultiModalDataItems() |
| 865 | for k, v in mm_data.items(): |
| 866 | if k not in subparsers: |
| 867 | raise ValueError(f"Unsupported modality: {k}") |
| 868 | |
| 869 | # ignore empty embedding data |
| 870 | if (parsed_data := subparsers[k](v)) is not None: |
| 871 | mm_items[k] = parsed_data |
| 872 | |
| 873 | return mm_items |
| 874 | |
| 875 | |
| 876 | MultiModalUUIDItems: TypeAlias = dict[str, Sequence[str | None]] |
| 877 | """ |
| 878 | A normalized [`MultiModalUUIDDict`][vllm.inputs.MultiModalUUIDDict] |
| 879 | such that each entry corresponds to a list. |
| 880 | """ |
| 881 | |
| 882 | |
| 883 | def parse_mm_uuids(mm_uuids: MultiModalUUIDDict | None) -> MultiModalUUIDItems: |
| 884 | if mm_uuids is None: |
| 885 | return {} |
| 886 | |
| 887 | return { |
| 888 | modality: [uuids] if isinstance(uuids, str) else uuids |
| 889 | for modality, uuids in mm_uuids.items() |
| 890 | } |