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 Florence2ImagePixelInputs to TensorSchema (#21663)
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Benji Beck committed
Jul 27, 2025 at 02:43 UTC
5f8c9a425ec99fcc2d7e547a7b9a0493490efd52
1 file changed
+24
-28
vllm/model_executor/models/florence2.py
+24
-28
@@ -4,7 +4,7 @@
4
import math
5
from collections import OrderedDict
6
from collections.abc import Iterable, Mapping, Sequence
7
-from typing import Literal, Optional, TypedDict, Union
7
+from typing import Annotated, Literal, Optional, Union
8
9
import torch
10
import torch.nn as nn
@@ -29,16 +29,28 @@ from vllm.multimodal.processing import (BaseProcessingInfo,
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
from .interfaces import (MultiModalEmbeddings, SupportsMultiModal,
35
SupportsV0Only)
36
from .utils import AutoWeightsLoader, flatten_bn, merge_multimodal_embeddings
37
38
38
-class Florence2ImagePixelInputs(TypedDict):
39
+class Florence2ImagePixelInputs(TensorSchema):
40
+ """
41
+ Dimensions:
42
+ - b: Batch size
43
+ - c: Number of channels (3)
44
+ - h: Height of the image
45
+ - w: Width of the image
46
+ """
47
+
48
type: Literal["pixel_values"]
40
- data: torch.Tensor
41
- """Shape: (batch_size, num_channel, height, width)"""
49
+
50
+ data: Annotated[
51
+ torch.Tensor,
52
+ TensorShape("b", 3, "h", "w"),
53
+ ]
54
55
56
# ViT implementation are all copied from
@@ -931,28 +943,6 @@ class Florence2ForConditionalGeneration(nn.Module, SupportsMultiModal,
943
raise NotImplementedError(
944
'Florence2 only supports COSINE as temporal embedding.')
945
934
- def _validate_pixel_values(
935
- self, data: Union[torch.Tensor, list[torch.Tensor]]
936
- ) -> Union[torch.Tensor, list[torch.Tensor]]:
937
-
938
- size = self.processor_config["size"]
939
- h, w = size["height"], size["width"]
940
- expected_dims = (3, h, w)
941
-
942
- def _validate_shape(d: torch.Tensor):
943
- actual_dims = tuple(d.shape)
944
-
945
- if actual_dims != expected_dims:
946
- expected_expr = tuple(*map(str, expected_dims))
947
- raise ValueError(
948
- "The expected shape of pixel values per batch "
949
- f"is {expected_expr}. You supplied {tuple(d.shape)}.")
950
-
951
- for d in data:
952
- _validate_shape(d)
953
-
954
- return data
955
-
946
def _parse_and_validate_image_input(self, **kwargs: object):
947
pixel_values: Optional[Union[list[list[torch.Tensor]],
948
list[torch.Tensor],
@@ -971,10 +961,16 @@ class Florence2ForConditionalGeneration(nn.Module, SupportsMultiModal,
961
"Both pixel values and image embeds are provided.")
962
963
if pixel_values is not None:
964
+ size = self.processor_config["size"]
965
+ expected_h, expected_w = size["height"], size["width"]
966
+
967
return Florence2ImagePixelInputs(
968
type="pixel_values",
976
- data=self._validate_pixel_values(
977
- flatten_bn(pixel_values, concat=True)),
969
+ data=flatten_bn(pixel_values, concat=True),
970
+ resolve_bindings={
971
+ "h": expected_h,
972
+ "w": expected_w
973
+ },
974
)
975
976
if image_embeds is not None: