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 ChameleonImagePixelInputs to TensorSchema (#21657)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 26, 2025 at 19:34 UTC
20950b29fb33465bcc822bb47a193ff0c0ba01bf
1 file changed
+21
-27
vllm/model_executor/models/chameleon.py
+21
-27
@@ -3,7 +3,7 @@
3
4
from collections.abc import Iterable, Mapping, Sequence
5
from functools import cached_property
6
-from typing import Any, Literal, Optional, TypedDict, Union
6
+from typing import Annotated, Any, Literal, Optional, Union
7
8
import torch
9
import torch.nn as nn
@@ -38,6 +38,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
38
PromptUpdate, PromptUpdateDetails)
39
from vllm.multimodal.profiling import BaseDummyInputsBuilder
40
from vllm.sequence import IntermediateTensors
41
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
42
43
from .interfaces import (MultiModalEmbeddings, SupportsMultiModal, SupportsPP,
44
SupportsQuant)
@@ -48,10 +49,16 @@ from .utils import (flatten_bn, is_pp_missing_parameter,
49
logger = init_logger(__name__)
50
51
51
-class ChameleonImagePixelInputs(TypedDict):
52
+class ChameleonImagePixelInputs(TensorSchema):
53
+ """
54
+ Dimensions:
55
+ - bn: Batch size * number of images
56
+ - c: Number of channels (3)
57
+ - h: Height of each image
58
+ - w: Width of each image
59
+ """
60
type: Literal["pixel_values"]
53
- data: torch.Tensor
54
- """Shape: `(batch_size * num_images, num_channels, height, width)`"""
61
+ data: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
62
63
64
class ChameleonProcessingInfo(BaseProcessingInfo):
@@ -962,19 +969,6 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,
969
self.make_empty_intermediate_tensors = (
970
self.model.make_empty_intermediate_tensors)
971
965
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
966
- vq_config: ChameleonVQVAEConfig = self.config.vq_config
967
- expected_dims = (3, vq_config.resolution, vq_config.resolution)
968
- actual_dims = tuple(data.shape[1:])
969
-
970
- if actual_dims != expected_dims:
971
- expected_expr = ("batch_size", *map(str, expected_dims))
972
- raise ValueError(
973
- f"The expected shape of pixel values is {expected_expr}. "
974
- f"You supplied {tuple(data.shape)}.")
975
-
976
- return data
977
-
972
def _parse_and_validate_image_input(
973
self, **kwargs: object) -> Optional[ChameleonImagePixelInputs]:
974
pixel_values = kwargs.pop("pixel_values", None)
@@ -982,16 +976,16 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,
976
if pixel_values is None:
977
return None
978
985
- if not isinstance(pixel_values, (torch.Tensor, list)):
986
- raise ValueError("Incorrect type of pixel values. "
987
- f"Got type: {type(pixel_values)}")
988
-
989
- pixel_values = flatten_bn(pixel_values, concat=True)
990
-
991
- return ChameleonImagePixelInputs(
992
- type="pixel_values",
993
- data=self._validate_pixel_values(pixel_values),
994
- )
979
+ vq_config: ChameleonVQVAEConfig = self.config.vq_config
980
+ expected_h = expected_w = vq_config.resolution
981
+
982
+ return ChameleonImagePixelInputs(type="pixel_values",
983
+ data=flatten_bn(pixel_values,
984
+ concat=True),
985
+ resolve_bindings={
986
+ "h": expected_h,
987
+ "w": expected_w
988
+ })
989
990
def get_language_model(self) -> torch.nn.Module:
991
return self.model