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, 05:30 UTC
Migrate DonutImagePixelInputs to TensorSchema (#23509)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 24, 2025 at 22:02 UTC
787cdb3829676504da2f612fad041db4b8acc271
1 file changed
+19
-30
vllm/model_executor/models/donut.py
+19
-30
@@ -3,7 +3,7 @@
3
4
import math
5
from collections.abc import Iterable, Mapping, Sequence
6
-from typing import Literal, Optional, TypedDict, Union
6
+from typing import Annotated, Literal, Optional, Union
7
8
import torch
9
import torch.nn as nn
@@ -29,6 +29,7 @@ from vllm.multimodal.processing import (BaseProcessingInfo,
29
PromptIndexTargets, PromptInsertion,
30
PromptUpdate)
31
from vllm.multimodal.profiling import BaseDummyInputsBuilder
32
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
33
34
35
class MBartDecoderWrapper(nn.Module):
@@ -132,10 +133,16 @@ class DonutLanguageForConditionalGeneration(nn.Module, SupportsV0Only):
133
return loaded_params
134
135
135
-class DonutImagePixelInputs(TypedDict):
136
+class DonutImagePixelInputs(TensorSchema):
137
+ """
138
+ Dimensions:
139
+ - b: Batch size
140
+ - c: Number of channels (3)
141
+ - h: Height
142
+ - w: Width
143
+ """
144
type: Literal["pixel_values"]
137
- data: torch.Tensor
138
- """Shape: (batch_size, num_channel, height, width)"""
145
+ data: Annotated[torch.Tensor, TensorShape("b", 3, "h", "w")]
146
147
148
class DonutProcessingInfo(BaseProcessingInfo):
@@ -275,27 +282,6 @@ class DonutForConditionalGeneration(nn.Module, SupportsMultiModal,
282
)
283
self.pad_token_id = config.pad_token_id
284
278
- def _validate_pixel_values(
279
- self, data: Union[torch.Tensor, list[torch.Tensor]]
280
- ) -> Union[torch.Tensor, list[torch.Tensor]]:
281
-
282
- # size = self.processor_config["size"]
283
- h, w = self.config.encoder.image_size
284
- expected_dims = (3, h, w)
285
-
286
- def _validate_shape(d: torch.Tensor):
287
- actual_dims = tuple(d.shape)
288
-
289
- if actual_dims != expected_dims:
290
- raise ValueError(
291
- "The expected shape of pixel values per batch "
292
- f"is {expected_dims}. You supplied {actual_dims}.")
293
-
294
- for d in data:
295
- _validate_shape(d)
296
-
297
- return data
298
-
285
def _parse_and_validate_image_input(self, **kwargs: object):
286
pixel_values: Optional[Union[list[list[torch.Tensor]],
287
list[torch.Tensor],
@@ -314,11 +300,14 @@ class DonutForConditionalGeneration(nn.Module, SupportsMultiModal,
300
"Both pixel values and image embeds are provided.")
301
302
if pixel_values is not None:
317
- return DonutImagePixelInputs(
318
- type="pixel_values",
319
- data=self._validate_pixel_values(
320
- flatten_bn(pixel_values, concat=True)),
321
- )
303
+ h, w = self.config.encoder.image_size
304
+ return DonutImagePixelInputs(type="pixel_values",
305
+ data=flatten_bn(pixel_values,
306
+ concat=True),
307
+ resolve_bindings={
308
+ "h": h,
309
+ "w": w,
310
+ })
311
312
if image_embeds is not None:
313
raise NotImplementedError