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 Blip2ImagePixelInputs and Blip2ImageEmbeddingInputs to TensorSchema (#21656)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 26, 2025 at 19:33 UTC
ccf27cc4d4645a1382c7742bb6bddedb7dc50fea
1 file changed
+30
-42
vllm/model_executor/models/blip2.py
+30
-42
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Iterable, Mapping, Sequence
5
-from typing import Literal, Optional, TypedDict, Union
5
+from typing import Annotated, Literal, Optional, Union
6
7
import torch
8
import torch.nn as nn
@@ -22,6 +22,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
22
PromptInsertion, PromptUpdate)
23
from vllm.multimodal.profiling import BaseDummyInputsBuilder
24
from vllm.sequence import IntermediateTensors
25
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
26
27
from .blip import BlipVisionModel
28
from .interfaces import (MultiModalEmbeddings, SupportsMultiModal, SupportsPP,
@@ -34,19 +35,27 @@ from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
35
_IMAGE_TOKEN_ID = 50265
36
37
37
-class Blip2ImagePixelInputs(TypedDict):
38
+class Blip2ImagePixelInputs(TensorSchema):
39
+ """
40
+ Dimensions:
41
+ - bn: Batch size * number of images
42
+ - c: Number of channels (3)
43
+ - h: Height of each image
44
+ - w: Width of each image
45
+ """
46
type: Literal["pixel_values"]
39
- data: torch.Tensor
40
- """Shape: `(batch_size * num_images, num_channels, height, width)`"""
47
+ data: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
48
49
43
-class Blip2ImageEmbeddingInputs(TypedDict):
44
- type: Literal["image_embeds"]
45
- data: torch.Tensor
46
- """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
47
-
48
- `hidden_size` must match the hidden size of language model backbone.
50
+class Blip2ImageEmbeddingInputs(TensorSchema):
51
"""
52
+ Dimensions:
53
+ - bn: Batch size * number of images
54
+ - f: Image feature size
55
+ - h: Hidden size (must match the hidden size of language model backbone)
56
+ """
57
+ type: Literal["image_embeds"]
58
+ data: Annotated[torch.Tensor, TensorShape("bn", "f", "h")]
59
60
61
Blip2ImageInputs = Union[Blip2ImagePixelInputs, Blip2ImageEmbeddingInputs]
@@ -551,21 +560,8 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
560
self.make_empty_intermediate_tensors = (
561
self.language_model.make_empty_intermediate_tensors)
562
554
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
555
- h = w = self.config.vision_config.image_size
556
- expected_dims = (3, h, w)
557
- actual_dims = tuple(data.shape[1:])
558
-
559
- if actual_dims != expected_dims:
560
- expected_expr = ("batch_size", *map(str, expected_dims))
561
- raise ValueError(
562
- f"The expected shape of pixel values is {expected_expr}. "
563
- f"You supplied {tuple(data.shape)}.")
564
-
565
- return data
566
-
567
- def _parse_and_validate_image_input(
568
- self, **kwargs: object) -> Optional[Blip2ImageInputs]:
563
+ def _create_image_input(self,
564
+ **kwargs: object) -> Optional[Blip2ImageInputs]:
565
pixel_values = kwargs.pop("pixel_values", None)
566
image_embeds = kwargs.pop("image_embeds", None)
567
@@ -573,27 +569,19 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
569
return None
570
571
if pixel_values is not None:
576
- if not isinstance(pixel_values, (torch.Tensor, list)):
577
- raise ValueError("Incorrect type of pixel values. "
578
- f"Got type: {type(pixel_values)}")
579
-
580
- pixel_values = flatten_bn(pixel_values, concat=True)
581
-
582
- return Blip2ImagePixelInputs(
583
- type="pixel_values",
584
- data=self._validate_pixel_values(pixel_values),
585
- )
572
+ expected_h = expected_w = self.config.vision_config.image_size
573
+ return Blip2ImagePixelInputs(type="pixel_values",
574
+ data=flatten_bn(pixel_values,
575
+ concat=True),
576
+ resolve_bindings={
577
+ "h": expected_h,
578
+ "w": expected_w
579
+ })
580
581
if image_embeds is not None:
588
- if not isinstance(image_embeds, (torch.Tensor, list)):
589
- raise ValueError("Incorrect type of image embeddings. "
590
- f"Got type: {type(image_embeds)}")
591
-
592
- image_embeds = flatten_bn(image_embeds, concat=True)
593
-
582
return Blip2ImageEmbeddingInputs(
583
type="image_embeds",
596
- data=image_embeds,
584
+ data=flatten_bn(image_embeds, concat=True),
585
)
586
587
raise AssertionError("This line should be unreachable.")