A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 23:45 UTC
Migrate AyaVisionImagePixelInputs to TensorSchema for shape validation (#21622)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 26, 2025 at 06:08 UTC
de10ff0b7cc757af3d0374d82c1a2130196af496
1 file changed
+29
-38
vllm/model_executor/models/aya_vision.py
+29
-38
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# Adapted from https://github.com/huggingface/transformers/tree/main/src/transformers/models/aya_vision
4
from collections.abc import Iterable, Mapping, Sequence
5
-from typing import Literal, Optional, TypedDict, Union, cast
5
+from typing import Annotated, Literal, Optional, Union, cast
6
7
import torch
8
from torch import nn
@@ -29,6 +29,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
29
PromptUpdateDetails)
30
from vllm.multimodal.profiling import BaseDummyInputsBuilder
31
from vllm.sequence import IntermediateTensors
32
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
33
34
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
35
from .siglip import SiglipVisionModel
@@ -37,18 +38,28 @@ from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
38
merge_multimodal_embeddings)
39
40
40
-class AyaVisionImagePixelInputs(TypedDict):
41
- type: Literal["pixel_values"]
42
- pixel_values: torch.Tensor
41
+class AyaVisionImagePixelInputs(TensorSchema):
42
"""
44
- Shape: `(num_patches_total, num_channels, height, width)`
45
-
46
- `num_patches_total` is the total number of patches over each image over each
47
- prompt in the batch.
43
+ Dimensions:
44
+ - np: The total number of patches over each image over each prompt in
45
+ the batch
46
+ - c: Number of channels
47
+ - h: Height of each image patch
48
+ - w: Width of each image patch
49
+ - bn: Batch size * number of images
50
"""
51
50
- num_patches: torch.Tensor
51
- """Shape: `(batch_size * num_images)`"""
52
+ type: Literal["pixel_values"]
53
+
54
+ pixel_values: Annotated[
55
+ torch.Tensor,
56
+ TensorShape("np", 3, "h", "w"),
57
+ ]
58
+
59
+ num_patches: Annotated[
60
+ torch.Tensor,
61
+ TensorShape("bn"),
62
+ ]
63
64
65
class AyaVisionMultiModalProjector(nn.Module):
@@ -383,21 +394,6 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,
394
e.flatten(0, 2) for e in image_embeds.split(num_patches.tolist())
395
]
396
386
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
387
- h = w = self.config.vision_config.image_size
388
- expected_dims = (3, h, w)
389
-
390
- def _validate_shape(d: torch.Tensor):
391
- if d.shape != expected_dims:
392
- raise ValueError(
393
- "The expected shape of pixel values per image per batch "
394
- f"is {expected_dims}. You supplied {tuple(d.shape)}.")
395
-
396
- for d in data:
397
- _validate_shape(d)
398
-
399
- return data
400
-
397
def _parse_and_validate_image_input(
398
self, **kwargs: object) -> Optional[AyaVisionImagePixelInputs]:
399
pixel_values = kwargs.pop("pixel_values", None)
@@ -405,22 +401,17 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,
401
image_embeds = kwargs.pop("image_embeds", None)
402
assert image_embeds is None, "Aya Vision does not support image_embeds."
403
408
- if not isinstance(pixel_values, (torch.Tensor, list)):
409
- raise ValueError("Incorrect type of pixel values. "
410
- f"Got type: {type(pixel_values)}")
411
- if num_patches is not None and not isinstance(num_patches,
412
- (torch.Tensor, list)):
413
- raise ValueError("Incorrect type of num_patches. "
414
- f"Got type: {type(num_patches)}")
415
-
416
- pixel_values = flatten_bn(pixel_values, concat=True)
417
- num_patches = flatten_bn(num_patches, concat=True)
404
+ if pixel_values is None:
405
+ return None
406
407
return AyaVisionImagePixelInputs(
408
type="pixel_values",
421
- pixel_values=self._validate_pixel_values(pixel_values),
422
- num_patches=num_patches,
423
- )
409
+ pixel_values=flatten_bn(pixel_values, concat=True),
410
+ num_patches=flatten_bn(num_patches, concat=True),
411
+ resolve_bindings={
412
+ "h": self.config.vision_config.image_size,
413
+ "w": self.config.vision_config.image_size,
414
+ })
415
416
def get_language_model(self) -> torch.nn.Module:
417
return self.language_model