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 AriaImagePixelInputs to TensorSchema for shape validation (#21620)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 26, 2025 at 06:08 UTC
9d197280fa60e5f313b15ecca723ef9bec9b5399
1 file changed
+21
-29
vllm/model_executor/models/aria.py
+21
-29
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Mapping, Sequence
4
-from typing import Optional, TypedDict, Union
4
+from typing import Annotated, Optional, Union
5
6
import torch
7
import torch.nn as nn
@@ -29,6 +29,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
29
PromptUpdate)
30
from vllm.multimodal.profiling import BaseDummyInputsBuilder
31
from vllm.sequence import IntermediateTensors
32
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
33
34
# yapf: disable
35
from .idefics2_vision_model import Idefics2VisionConfig
@@ -42,15 +43,26 @@ from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
43
merge_multimodal_embeddings)
44
45
45
-class AriaImagePixelInputs(TypedDict):
46
- pixel_values: torch.Tensor
47
- pixel_mask: Optional[torch.Tensor]
46
+class AriaImagePixelInputs(TensorSchema):
47
"""
49
- Shape:
50
- pixel_values: `(batch_size * num_images, num_channels, height, width)`
51
- pixel_mask: `(batch_size * num_images, height, width)`
48
+ Dimensions:
49
+ - b: Batch size
50
+ - n: Number of images
51
+ - c: Number of channels
52
+ - h: Height of each image
53
+ - w: Width of each image
54
"""
55
56
+ pixel_values: Annotated[
57
+ torch.Tensor,
58
+ TensorShape("bn", 3, "h", "w"),
59
+ ]
60
+
61
+ pixel_mask: Annotated[
62
+ Optional[torch.Tensor],
63
+ TensorShape("bn", "h", "w"),
64
+ ]
65
+
66
67
class AriaVisionTransformer(Idefics3VisionTransformer, SupportsQuant):
68
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
@@ -540,12 +552,6 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
552
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
553
self.vocab_size, logit_scale)
554
543
- def _validate_image_sizes(
544
- self, images: list[torch.Tensor]) -> list[torch.Tensor]:
545
- if not all(img.shape == images[0].shape for img in images):
546
- raise ValueError("All images must be the same size")
547
- return images
548
-
555
def _parse_and_validate_image_input(
556
self, **kwargs: object) -> Optional[AriaImagePixelInputs]:
557
pixel_values = kwargs.pop("pixel_values", None)
@@ -554,23 +560,9 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
560
if pixel_values is None:
561
return None
562
557
- if not isinstance(pixel_values, (torch.Tensor, list)):
558
- raise ValueError("Incorrect type of pixel values. "
559
- f"Got type: {type(pixel_values)}")
560
-
561
- pixel_values = self._validate_image_sizes(pixel_values)
562
- pixel_values = flatten_bn(pixel_values, concat=True)
563
-
564
- if pixel_mask is not None:
565
- if not isinstance(pixel_mask, (torch.Tensor, list)):
566
- raise ValueError("Incorrect type of pixel mask. "
567
- f"Got type: {type(pixel_mask)}")
568
-
569
- pixel_mask = flatten_bn(pixel_mask, concat=True)
570
-
563
return AriaImagePixelInputs(
572
- pixel_values=pixel_values,
573
- pixel_mask=pixel_mask,
564
+ pixel_values=flatten_bn(pixel_values, concat=True),
565
+ pixel_mask=flatten_bn(pixel_mask, concat=True),
566
)
567
568
def _create_patch_attention_mask(