@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1c6158083a
Fork 0

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 3, 08:00 UTC

[Model] Openvla support (#42654)

Signed-off-by: Wang Yiwen <121547057+yiwen101@users.noreply.github.com>

Wang Yiwen committed May 19, 2026 at 23:17 UTC 1c6158083a6fc3aff408660d2defd7602f78f556
11 files changed +906
docs/models/supported_models.md
+1
@@ -605,6 +605,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
605 | `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
606 | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
607 | `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
608 +| `OpenVLAForActionPrediction` | OpenVLA | T + I | `openvla/openvla-7b` | | ✅︎ |
609 | `Ovis` | Ovis2, Ovis1.6 | T + I<sup>+</sup> | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ |
610 | `Ovis2_5` | Ovis2.5 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.5-9B`, etc. | | |
611 | `Ovis2_6ForCausalLM` | Ovis2.6 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.6-2B`, etc. | | |
tests/models/multimodal/processing/test_common.py
+5
@@ -416,6 +416,11 @@ def test_processing_correctness(
416 pytest.skip("Fix later")
417 if model_id == "OpenGVLab/InternVL2-2B":
418 pytest.skip("Fix later")
419 + if model_id == "openvla/openvla-7b":
420 + pytest.skip(
421 + "OpenVLA uses a custom vLLM processor because its HF remote "
422 + "processor is incompatible with current Transformers."
423 + )
424 if model_id == "jinaai/jina-reranker-m0":
425 pytest.skip("Fix later")
426 if model_id in {"Qwen/Qwen-VL", "Qwen/Qwen-VL-Chat"}:
tests/models/multimodal/processing/test_openvla.py new
+210
@@ -0,0 +1,210 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for OpenVLA multimodal preprocessing."""
4 +
5 +import numpy as np
6 +import pytest
7 +import torch
8 +from PIL import Image
9 +from transformers import LlamaConfig
10 +
11 +from vllm.model_executor.models.openvla import (
12 + OpenVLAForActionPrediction,
13 + OpenVLAMultiModalProcessor,
14 + OpenVLAProcessingInfo,
15 +)
16 +from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
17 +from vllm.transformers_utils.configs.openvla import OpenVLAConfig
18 +from vllm.transformers_utils.processors.openvla import (
19 + IMAGENET_MEAN,
20 + IMAGENET_STD,
21 + SIGLIP_MEAN,
22 + SIGLIP_STD,
23 + OpenVLAImageProcessor,
24 + OpenVLAProcessor,
25 + preprocess_openvla_image,
26 + to_rgb_image,
27 +)
28 +
29 +pytestmark = pytest.mark.cpu_test
30 +
31 +
32 +class _FakeTokenizer:
33 + bos_token_id = 1
34 + init_kwargs: dict[str, object] = {}
35 +
36 + def encode(self, prompt: str, **kwargs: object) -> list[int]:
37 + assert prompt == "In: test\nOut:"
38 + if kwargs == {"add_special_tokens": True}:
39 + return [self.bos_token_id, 10, 11]
40 + assert kwargs == {"add_special_tokens": False}
41 + return [10, 11]
42 +
43 + def __call__(self, text: str, **kwargs: object) -> dict[str, list[list[int]]]:
44 + return {"input_ids": [self.encode(text, **kwargs)]}
45 +
46 +
47 +class _FakeProcessingInfo:
48 + def __init__(self) -> None:
49 + self.config = OpenVLAConfig()
50 +
51 + def get_hf_config(self) -> OpenVLAConfig:
52 + return self.config
53 +
54 + def get_tokenizer(self) -> _FakeTokenizer:
55 + return _FakeTokenizer()
56 +
57 + def get_num_image_tokens(self, *, image_width: int, image_height: int) -> int:
58 + assert image_width > 0
59 + assert image_height > 0
60 + return 256
61 +
62 +
63 +class _FakeOpenVLAProcessingInfo(OpenVLAProcessingInfo):
64 + def get_hf_config(self) -> OpenVLAConfig:
65 + return OpenVLAConfig()
66 +
67 +
68 +def _make_processor() -> OpenVLAMultiModalProcessor:
69 + processor = OpenVLAMultiModalProcessor.__new__(OpenVLAMultiModalProcessor)
70 + processor.info = _FakeProcessingInfo()
71 + return processor
72 +
73 +
74 +def test_openvla_config_converts_text_config_dict() -> None:
75 + config = OpenVLAConfig(
76 + text_config={
77 + "vocab_size": 123,
78 + "hidden_size": 64,
79 + "intermediate_size": 128,
80 + "num_hidden_layers": 2,
81 + "num_attention_heads": 4,
82 + },
83 + )
84 +
85 + assert isinstance(config.text_config, LlamaConfig)
86 + assert config.text_config.vocab_size == 123
87 + assert config.text_config.hidden_size == 64
88 + assert config.text_config.architectures == ["LlamaForCausalLM"]
89 +
90 +
91 +@pytest.mark.parametrize(
92 + ("image", "expected_size", "expected_pixel"),
93 + [
94 + (
95 + torch.tensor(
96 + [
97 + [[1.0, 1.0], [1.0, 1.0]],
98 + [[0.0, 0.0], [0.0, 0.0]],
99 + [[0.0, 0.0], [0.0, 0.0]],
100 + ]
101 + ),
102 + (2, 2),
103 + (255, 0, 0),
104 + ),
105 + (
106 + np.full((4, 5, 1), 128, dtype=np.uint8),
107 + (5, 4),
108 + (128, 128, 128),
109 + ),
110 + ],
111 +)
112 +def test_openvla_to_rgb_image(
113 + image: torch.Tensor | np.ndarray,
114 + expected_size: tuple[int, int],
115 + expected_pixel: tuple[int, int, int],
116 +) -> None:
117 + rgb_image = to_rgb_image(image)
118 +
119 + assert rgb_image.mode == "RGB"
120 + assert rgb_image.size == expected_size
121 + assert rgb_image.getpixel((0, 0)) == expected_pixel
122 +
123 +
124 +def test_openvla_preprocess_image_matches_expected_normalization() -> None:
125 + image = Image.fromarray(
126 + np.arange(12 * 10 * 3, dtype=np.uint8).reshape(10, 12, 3),
127 + mode="RGB",
128 + )
129 +
130 + pixel_values = preprocess_openvla_image(image, image_size=224)
131 +
132 + resized = image.resize((224, 224), Image.Resampling.BICUBIC)
133 + raw = np.asarray(resized, dtype=np.float32) / 255.0
134 + expected_dinov2 = ((raw - IMAGENET_MEAN) / IMAGENET_STD).transpose(2, 0, 1)
135 + expected_siglip = ((raw - SIGLIP_MEAN) / SIGLIP_STD).transpose(2, 0, 1)
136 + expected = np.concatenate([expected_dinov2, expected_siglip], axis=0)
137 +
138 + assert pixel_values.shape == (6, 224, 224)
139 + assert pixel_values.dtype == torch.float32
140 + torch.testing.assert_close(pixel_values, torch.from_numpy(expected))
141 +
142 +
143 +def test_openvla_processor_outputs_pixel_values() -> None:
144 + processor = OpenVLAProcessor(
145 + image_processor=OpenVLAImageProcessor(image_size=224),
146 + tokenizer=_FakeTokenizer(),
147 + )
148 + image = Image.new("RGB", (8, 8), color=(255, 0, 0))
149 +
150 + batch = processor(
151 + text="In: test\nOut:",
152 + images=image,
153 + text_kwargs={"add_special_tokens": True},
154 + )
155 +
156 + assert batch["input_ids"] == [[1, 10, 11]]
157 + assert batch["pixel_values"].shape == (1, 6, 224, 224)
158 + assert batch["pixel_values"].dtype == torch.float32
159 +
160 +
161 +def test_openvla_image_processor_outputs_pixel_values() -> None:
162 + processor = OpenVLAImageProcessor(image_size=224)
163 + image = Image.new("RGB", (8, 8), color=(255, 0, 0))
164 +
165 + output = processor(images=image)
166 +
167 + assert output["pixel_values"].shape == (1, 6, 224, 224)
168 + assert output["pixel_values"].dtype == torch.float32
169 +
170 +
171 +def test_openvla_processing_info_token_counts() -> None:
172 + info = _FakeOpenVLAProcessingInfo.__new__(_FakeOpenVLAProcessingInfo)
173 +
174 + assert info.get_supported_mm_limits() == {"image": 1}
175 + assert info.get_num_image_tokens(image_width=640, image_height=480) == 256
176 + assert info.get_image_size_with_most_features().width == 224
177 + assert info.get_image_size_with_most_features().height == 224
178 + assert info.get_mm_max_tokens_per_item(seq_len=2048, mm_counts={"image": 1}) == {
179 + "image": 256
180 + }
181 +
182 +
183 +def test_openvla_prompt_update_inserts_image_tokens_after_bos() -> None:
184 + processor = _make_processor()
185 + image = Image.new("RGB", (640, 480), color=(255, 255, 255))
186 + mm_items = MultiModalDataItems({"image": ImageProcessorItems([image])})
187 +
188 + assert (
189 + processor._hf_processor_applies_updates("In: test\nOut:", mm_items, {}, {})
190 + is False
191 + )
192 +
193 + prompt_update = processor._get_prompt_updates(mm_items, {}, {})[0]
194 + resolved = prompt_update.resolve(0)
195 + content = resolved.content
196 +
197 + assert resolved.modality == "image"
198 + assert [
199 + (match.start_idx, match.end_idx)
200 + for match in resolved.iter_matches([1, 10, 11], _FakeTokenizer())
201 + ] == [(1, 1)]
202 + assert content.full == [32000] * 256
203 +
204 + is_embed = content.is_embed(None, content.full)
205 + assert is_embed.dtype == torch.bool
206 + assert is_embed.tolist() == [True] * 256
207 +
208 +
209 +def test_openvla_placeholder_string() -> None:
210 + assert OpenVLAForActionPrediction.get_placeholder_str("image", 0) is None
tests/models/registry.py
+4
@@ -1203,6 +1203,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1203 )
1204 },
1205 ),
1206 + "OpenVLAForActionPrediction": _HfExamplesInfo(
1207 + "openvla/openvla-7b",
1208 + trust_remote_code=True,
1209 + ),
1210 "Ovis": _HfExamplesInfo(
1211 "AIDC-AI/Ovis2-1B",
1212 trust_remote_code=True,
vllm/model_executor/models/openvla.py new
+528
@@ -0,0 +1,528 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterable, Mapping, Sequence
5 +from typing import Annotated, Literal
6 +
7 +import torch
8 +import torch.nn as nn
9 +from transformers import BatchFeature
10 +
11 +from vllm.config import VllmConfig
12 +from vllm.config.multimodal import BaseDummyOptions
13 +from vllm.inputs import MultiModalDataDict
14 +from vllm.model_executor.layers.activation import get_act_fn
15 +from vllm.model_executor.layers.linear import (
16 + ColumnParallelLinear,
17 + ReplicatedLinear,
18 + RowParallelLinear,
19 +)
20 +from vllm.model_executor.layers.quantization import QuantizationConfig
21 +from vllm.model_executor.models.interfaces import (
22 + MultiModalEmbeddings,
23 + SupportsMultiModal,
24 + SupportsPP,
25 +)
26 +from vllm.multimodal import MULTIMODAL_REGISTRY
27 +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
28 +from vllm.multimodal.parse import (
29 + ImageEmbeddingItems,
30 + ImageProcessorItems,
31 + ImageSize,
32 + MultiModalDataItems,
33 +)
34 +from vllm.multimodal.processing import (
35 + BaseDummyInputsBuilder,
36 + BaseMultiModalProcessor,
37 + BaseProcessingInfo,
38 + InputProcessingContext,
39 + PromptIndexTargets,
40 + PromptInsertion,
41 + PromptUpdate,
42 + PromptUpdateDetails,
43 +)
44 +from vllm.sequence import IntermediateTensors
45 +from vllm.transformers_utils.configs import OpenVLAConfig
46 +from vllm.transformers_utils.processors.openvla import (
47 + OpenVLAImageProcessor,
48 + OpenVLAProcessor,
49 +)
50 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
51 +
52 +from .module_mapping import MultiModelKeys
53 +from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
54 +
55 +# openvla/openvla-7b uses 224x224 images with ViT patch size 14, yielding a
56 +# 16x16 image-token grid.
57 +_OPENVLA_IMAGE_SIZE = 224
58 +_OPENVLA_PATCH_SIZE = 14
59 +_OPENVLA_TIMM_MODEL_IDS = (
60 + "vit_large_patch14_reg4_dinov2.lvd142m",
61 + "vit_so400m_patch14_siglip_224",
62 +)
63 +_OPENVLA_TIMM_OVERRIDE_ACT_LAYERS = (None, None)
64 +_OPENVLA_IMAGE_SIZES = (_OPENVLA_IMAGE_SIZE, _OPENVLA_IMAGE_SIZE)
65 +
66 +
67 +def _get_num_image_tokens(image_size: int) -> int:
68 + return (image_size // _OPENVLA_PATCH_SIZE) ** 2
69 +
70 +
71 +class OpenVLAImagePixelInputs(TensorSchema):
72 + """
73 + Dimensions:
74 + - bn: Batch size * number of images
75 + - c: Number of channels (6)
76 + - h: Height
77 + - w: Width
78 + """
79 +
80 + type: Literal["pixel_values"] = "pixel_values"
81 + data: Annotated[torch.Tensor, TensorShape("bn", 6, "h", "w")]
82 +
83 +
84 +class PrismaticVisionBackbone(nn.Module):
85 + """OpenVLA's fused DINOv2 + SigLIP vision backbone."""
86 +
87 + def __init__(
88 + self,
89 + *,
90 + image_sizes: Sequence[int],
91 + timm_model_ids: Sequence[str],
92 + timm_override_act_layers: Sequence[str | None],
93 + use_fused_vision_backbone: bool,
94 + ) -> None:
95 + super().__init__()
96 + if not use_fused_vision_backbone:
97 + raise ValueError(
98 + "OpenVLA currently supports only the fused DINOv2 + SigLIP "
99 + "vision backbone."
100 + )
101 + if tuple(image_sizes) != _OPENVLA_IMAGE_SIZES:
102 + raise ValueError(
103 + "OpenVLA currently supports only 224x224 image inputs, "
104 + f"got image_sizes={list(image_sizes)}."
105 + )
106 + if tuple(timm_model_ids) != _OPENVLA_TIMM_MODEL_IDS:
107 + raise ValueError(
108 + "OpenVLA currently supports only the dinosiglip-vit-so-224px "
109 + "vision backbone, got "
110 + f"timm_model_ids={list(timm_model_ids)}."
111 + )
112 + if tuple(timm_override_act_layers) != _OPENVLA_TIMM_OVERRIDE_ACT_LAYERS:
113 + raise ValueError(
114 + "OpenVLA currently supports only the default timm activation "
115 + "layers, got "
116 + f"timm_override_act_layers={list(timm_override_act_layers)}."
117 + )
118 +
119 + self.image_size = image_sizes[0]
120 + self.use_fused_vision_backbone = use_fused_vision_backbone
121 +
122 + self.embed_dim = 2176 if use_fused_vision_backbone else 1024
123 +
124 + try:
125 + import timm
126 + except ImportError as e:
127 + raise ImportError(
128 + "Please install timm to use OpenVLA. OpenVLA verification "
129 + "used timm==0.9.10."
130 + ) from e
131 +
132 + self.dinov2_featurizer = timm.create_model(
133 + timm_model_ids[0],
134 + pretrained=False,
135 + num_classes=0,
136 + img_size=self.image_size,
137 + act_layer=timm_override_act_layers[0],
138 + )
139 + self.siglip_featurizer = (
140 + timm.create_model(
141 + timm_model_ids[1],
142 + pretrained=False,
143 + num_classes=0,
144 + img_size=self.image_size,
145 + act_layer=timm_override_act_layers[1],
146 + )
147 + if use_fused_vision_backbone
148 + else None
149 + )
150 +
151 + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
152 + if self.dinov2_featurizer is None:
153 + raise RuntimeError("OpenVLA vision backbone is not initialized.")
154 +
155 + if self.use_fused_vision_backbone and pixel_values.shape[1] != 6:
156 + raise ValueError(
157 + "OpenVLA fused DINOv2 + SigLIP backbone expects 6-channel "
158 + "image inputs: 3 DINOv2-normalized channels followed by 3 "
159 + "SigLIP-normalized channels, "
160 + f"got {pixel_values.shape[1]} channels."
161 + )
162 +
163 + dinov2_pixels = pixel_values[:, :3]
164 +
165 + num_dinov2_blocks = len(self.dinov2_featurizer.blocks)
166 + dinov2_features = self.dinov2_featurizer.get_intermediate_layers(
167 + dinov2_pixels, n={num_dinov2_blocks - 2}
168 + )[0]
169 +
170 + if self.siglip_featurizer is not None:
171 + siglip_pixels = pixel_values[:, 3:]
172 + num_siglip_blocks = len(self.siglip_featurizer.blocks)
173 + siglip_features = self.siglip_featurizer.get_intermediate_layers(
174 + siglip_pixels, n={num_siglip_blocks - 2}
175 + )[0]
176 + return torch.cat([dinov2_features, siglip_features], dim=-1)
177 +
178 + return dinov2_features
179 +
180 +
181 +class PrismaticProjector(nn.Module):
182 + """Project Prismatic vision features into the language-model hidden size."""
183 +
184 + def __init__(
185 + self,
186 + *,
187 + vision_dim: int,
188 + text_dim: int,
189 + use_fused_vision_backbone: bool,
190 + quant_config: QuantizationConfig | None = None,
191 + prefix: str = "",
192 + ) -> None:
193 + super().__init__()
194 + self.use_fused_vision_backbone = use_fused_vision_backbone
195 +
196 + if use_fused_vision_backbone:
197 + intermediate_dim = 4 * vision_dim
198 + self.fc1 = ColumnParallelLinear(
199 + vision_dim,
200 + intermediate_dim,
201 + bias=True,
202 + quant_config=quant_config,
203 + prefix=f"{prefix}.fc1",
204 + )
205 + self.act_fn1 = get_act_fn("gelu")
206 + self.fc2 = RowParallelLinear(
207 + intermediate_dim,
208 + text_dim,
209 + bias=True,
210 + quant_config=quant_config,
211 + prefix=f"{prefix}.fc2",
212 + )
213 + self.act_fn2 = get_act_fn("gelu")
214 + self.fc3 = ReplicatedLinear(
215 + text_dim,
216 + text_dim,
217 + bias=True,
218 + quant_config=quant_config,
219 + prefix=f"{prefix}.fc3",
220 + )
221 + else:
222 + self.fc1 = ColumnParallelLinear(
223 + vision_dim,
224 + text_dim,
225 + bias=True,
226 + quant_config=quant_config,
227 + prefix=f"{prefix}.fc1",
228 + )
229 + self.act_fn1 = get_act_fn("gelu")
230 + self.fc2 = RowParallelLinear(
231 + text_dim,
232 + text_dim,
233 + bias=True,
234 + quant_config=quant_config,
235 + prefix=f"{prefix}.fc2",
236 + )
237 +
238 + def forward(self, image_features: torch.Tensor) -> torch.Tensor:
239 + hidden_states, _ = self.fc1(image_features)
240 + hidden_states = self.act_fn1(hidden_states)
241 + hidden_states, _ = self.fc2(hidden_states)
242 +
243 + if self.use_fused_vision_backbone:
244 + hidden_states = self.act_fn2(hidden_states)
245 + hidden_states, _ = self.fc3(hidden_states)
246 +
247 + return hidden_states
248 +
249 +
250 +class OpenVLAProcessingInfo(BaseProcessingInfo):
251 + def __init__(self, ctx: InputProcessingContext) -> None:
252 + super().__init__(ctx)
253 + self.hf_processor = OpenVLAProcessor(
254 + image_processor=OpenVLAImageProcessor(
255 + image_size=self.get_hf_config().image_sizes[0],
256 + ),
257 + tokenizer=self.get_tokenizer(),
258 + )
259 +
260 + def get_hf_config(self) -> OpenVLAConfig:
261 + return self.ctx.get_hf_config(OpenVLAConfig)
262 +
263 + def get_hf_processor(self, **kwargs: object) -> OpenVLAProcessor:
264 + return self.hf_processor
265 +
266 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
267 + return {"image": 1}
268 +
269 + def get_num_image_tokens(
270 + self,
271 + *,
272 + image_width: int,
273 + image_height: int,
274 + ) -> int:
275 + image_size = self.get_hf_config().image_sizes[0]
276 + return _get_num_image_tokens(image_size)
277 +
278 + def get_image_size_with_most_features(self) -> ImageSize:
279 + image_size = self.get_hf_config().image_sizes[0]
280 + return ImageSize(width=image_size, height=image_size)
281 +
282 + def get_mm_max_tokens_per_item(
283 + self,
284 + seq_len: int,
285 + mm_counts: Mapping[str, int],
286 + ) -> Mapping[str, int] | None:
287 + image_size = self.get_hf_config().image_sizes[0]
288 + return {"image": _get_num_image_tokens(image_size)}
289 +
290 +
291 +class OpenVLADummyInputsBuilder(BaseDummyInputsBuilder[OpenVLAProcessingInfo]):
292 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
293 + return ""
294 +
295 + def get_dummy_mm_data(
296 + self,
297 + seq_len: int,
298 + mm_counts: Mapping[str, int],
299 + mm_options: Mapping[str, BaseDummyOptions],
300 + ) -> MultiModalDataDict:
301 + num_images = mm_counts.get("image", 0)
302 + image_overrides = mm_options.get("image")
303 + image_size = self.info.get_image_size_with_most_features()
304 +
305 + return {
306 + "image": self._get_dummy_images(
307 + width=image_size.width,
308 + height=image_size.height,
309 + num_images=num_images,
310 + overrides=image_overrides,
311 + )
312 + }
313 +
314 +
315 +class OpenVLAMultiModalProcessor(BaseMultiModalProcessor[OpenVLAProcessingInfo]):
316 + """Processor contract for OpenVLA image inputs.
317 +
318 + OpenVLA feeds the same RGB image to DINOv2 and SigLIP after different
319 + normalizations. The processor exposes this as one 6-channel tensor:
320 + channels 0-2 are DINOv2-normalized and channels 3-5 are SigLIP-normalized.
321 + """
322 +
323 + def _get_mm_fields_config(
324 + self,
325 + hf_inputs: BatchFeature,
326 + hf_processor_mm_kwargs: Mapping[str, object],
327 + ) -> Mapping[str, MultiModalFieldConfig]:
328 + return dict(pixel_values=MultiModalFieldConfig.batched("image"))
329 +
330 + def _hf_processor_applies_updates(
331 + self,
332 + prompt_text: str,
333 + mm_items: MultiModalDataItems,
334 + hf_processor_mm_kwargs: Mapping[str, object],
335 + tokenization_kwargs: Mapping[str, object],
336 + ) -> bool:
337 + return False
338 +
339 + def _get_prompt_updates(
340 + self,
341 + mm_items: MultiModalDataItems,
342 + hf_processor_mm_kwargs: Mapping[str, object],
343 + out_mm_kwargs: MultiModalKwargsItems,
344 + ) -> Sequence[PromptUpdate]:
345 + hf_config = self.info.get_hf_config()
346 + image_token_id = hf_config.image_token_index
347 +
348 + tokenizer = self.info.get_tokenizer()
349 + bos_token_id = tokenizer.bos_token_id
350 +
351 + def get_insertion(item_idx: int) -> PromptUpdateDetails[list[int]]:
352 + images = mm_items.get_items(
353 + "image", (ImageEmbeddingItems, ImageProcessorItems)
354 + )
355 + if isinstance(images, ImageEmbeddingItems):
356 + num_image_tokens = images.get_feature_size(item_idx)
357 + else:
358 + image_size = images.get_image_size(item_idx)
359 + num_image_tokens = self.info.get_num_image_tokens(
360 + image_width=image_size.width,
361 + image_height=image_size.height,
362 + )
363 +
364 + image_tokens = [image_token_id] * num_image_tokens
365 + return PromptUpdateDetails.select_token_id(
366 + image_tokens,
367 + embed_token_id=image_token_id,
368 + )
369 +
370 + return [
371 + PromptInsertion(
372 + modality="image",
373 + target=PromptIndexTargets.prefix(
374 + [bos_token_id] if bos_token_id is not None else []
375 + ),
376 + insertion=get_insertion,
377 + )
378 + ]
379 +
380 +
381 +@MULTIMODAL_REGISTRY.register_processor(
382 + OpenVLAMultiModalProcessor,
383 + info=OpenVLAProcessingInfo,
384 + dummy_inputs=OpenVLADummyInputsBuilder,
385 +)
386 +class OpenVLAForActionPrediction(nn.Module, SupportsMultiModal, SupportsPP):
387 + """OpenVLA wrapper with vLLM language-model execution wired in."""
388 +
389 + @classmethod
390 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
391 + if modality.startswith("image"):
392 + return None
393 + raise ValueError("Only image modality is supported")
394 +
395 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
396 + super().__init__()
397 + config = vllm_config.model_config.hf_config
398 + quant_config = vllm_config.quant_config
399 + self.config = config
400 + self.multimodal_config = vllm_config.model_config.multimodal_config
401 + self.image_token_id = config.image_token_index
402 + self.n_action_bins = config.n_action_bins
403 + self.num_patches = _get_num_image_tokens(config.image_sizes[0])
404 +
405 + with self._mark_tower_model(vllm_config, "image"):
406 + self.vision_backbone = PrismaticVisionBackbone(
407 + image_sizes=config.image_sizes,
408 + timm_model_ids=config.timm_model_ids,
409 + timm_override_act_layers=config.timm_override_act_layers,
410 + use_fused_vision_backbone=config.use_fused_vision_backbone,
411 + )
412 + self.projector = PrismaticProjector(
413 + vision_dim=self.vision_backbone.embed_dim,
414 + text_dim=config.text_config.hidden_size,
415 + use_fused_vision_backbone=config.use_fused_vision_backbone,
416 + quant_config=quant_config,
417 + prefix=maybe_prefix(prefix, "projector"),
418 + )
419 +
420 + with self._mark_language_model(vllm_config):
421 + self.language_model = init_vllm_registered_model(
422 + vllm_config=vllm_config,
423 + hf_config=config.text_config,
424 + prefix=maybe_prefix(prefix, "language_model"),
425 + )
426 +
427 + self.make_empty_intermediate_tensors = (
428 + self.language_model.make_empty_intermediate_tensors
429 + )
430 +
431 + def get_language_model(self) -> nn.Module:
432 + return self.language_model
433 +
434 + def _parse_and_validate_image_input(
435 + self,
436 + **kwargs: object,
437 + ) -> OpenVLAImagePixelInputs | None:
438 + pixel_values = kwargs.pop("pixel_values", None)
439 + if pixel_values is None:
440 + return None
441 +
442 + return OpenVLAImagePixelInputs(
443 + type="pixel_values",
444 + data=pixel_values,
445 + resolve_bindings={
446 + "h": self.config.image_sizes[0],
447 + "w": self.config.image_sizes[0],
448 + },
449 + )
450 +
451 + def _process_image_input(
452 + self,
453 + image_input: OpenVLAImagePixelInputs,
454 + ) -> torch.Tensor:
455 + if self.vision_backbone.dinov2_featurizer is None:
456 + raise RuntimeError("OpenVLA vision backbone is not initialized.")
457 +
458 + pixel_values = image_input["data"].to(
459 + dtype=self.vision_backbone.dinov2_featurizer.patch_embed.proj.weight.dtype
460 + )
461 + vision_features = self.vision_backbone(pixel_values)
462 + return self.projector(vision_features)
463 +
464 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
465 + image_input = self._parse_and_validate_image_input(**kwargs)
466 + if image_input is None:
467 + return []
468 +
469 + return self._process_image_input(image_input)
470 +
471 + def forward(
472 + self,
473 + input_ids: torch.Tensor | None,
474 + positions: torch.Tensor,
475 + intermediate_tensors: IntermediateTensors | None = None,
476 + inputs_embeds: torch.Tensor | None = None,
477 + **kwargs: object,
478 + ) -> torch.Tensor | IntermediateTensors:
479 + if intermediate_tensors is not None:
480 + inputs_embeds = None
481 +
482 + return self.language_model.model(
483 + input_ids,
484 + positions,
485 + intermediate_tensors,
486 + inputs_embeds=inputs_embeds,
487 + )
488 +
489 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
490 + return self.language_model.compute_logits(hidden_states)
491 +
492 + def get_mm_mapping(self) -> MultiModelKeys:
493 + return MultiModelKeys.from_string_field(
494 + language_model="language_model",
495 + connector="projector",
496 + tower_model="vision_backbone",
497 + )
498 +
499 + def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
500 + return num_image_tokens
501 +
502 + def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
503 + return num_vision_tokens
504 +
505 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
506 + def maybe_rename_vision_weights(
507 + weights: Iterable[tuple[str, torch.Tensor]],
508 + ) -> Iterable[tuple[str, torch.Tensor]]:
509 + for name, weight in weights:
510 + if name.startswith("vision_backbone.featurizer."):
511 + name = name.replace(
512 + "vision_backbone.featurizer.",
513 + "vision_backbone.dinov2_featurizer.",
514 + 1,
515 + )
516 + elif name.startswith("vision_backbone.fused_featurizer."):
517 + name = name.replace(
518 + "vision_backbone.fused_featurizer.",
519 + "vision_backbone.siglip_featurizer.",
520 + 1,
521 + )
522 + # HF uses .scale_factor, timm uses .gamma
523 + if ".ls1.scale_factor" in name or ".ls2.scale_factor" in name:
524 + name = name.replace(".scale_factor", ".gamma")
525 + yield name, weight
526 +
527 + loader = AutoWeightsLoader(self)
528 + return loader.load_weights(maybe_rename_vision_weights(weights))
vllm/model_executor/models/registry.py
+1
@@ -504,6 +504,7 @@ _MULTIMODAL_MODELS = {
504 "openpangu_vl",
505 "OpenPanguVLForConditionalGeneration",
506 ),
507 + "OpenVLAForActionPrediction": ("openvla", "OpenVLAForActionPrediction"),
508 "Ovis": ("ovis", "Ovis"),
509 "Ovis2_5": ("ovis2_5", "Ovis2_5"),
510 "Ovis2_6ForCausalLM": ("ovis2_5", "Ovis2_5"),
vllm/transformers_utils/config.py
+1
@@ -120,6 +120,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
120 speculators="SpeculatorsConfig",
121 nemotron="NemotronConfig",
122 olmo_hybrid="OlmoHybridConfig",
123 + openvla="OpenVLAConfig",
124 ovis="OvisConfig",
125 ultravox="UltravoxConfig",
126 step3_vl="Step3VLConfig",
vllm/transformers_utils/configs/__init__.py
+2
@@ -61,6 +61,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
61 "NemotronConfig": "vllm.transformers_utils.configs.nemotron",
62 "NemotronHConfig": "vllm.transformers_utils.configs.nemotron_h",
63 "OlmoHybridConfig": "vllm.transformers_utils.configs.olmo_hybrid",
64 + "OpenVLAConfig": "vllm.transformers_utils.configs.openvla",
65 "OvisConfig": "vllm.transformers_utils.configs.ovis",
66 "PixelShuffleSiglip2VisionConfig": "vllm.transformers_utils.configs.isaac",
67 "RadioConfig": "vllm.transformers_utils.configs.radio",
@@ -128,6 +129,7 @@ __all__ = [
129 "NemotronConfig",
130 "NemotronHConfig",
131 "OlmoHybridConfig",
132 + "OpenVLAConfig",
133 "OvisConfig",
134 "PixelShuffleSiglip2VisionConfig",
135 "RadioConfig",
vllm/transformers_utils/configs/openvla.py new
+50
@@ -0,0 +1,50 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""OpenVLA configuration support.
4 +
5 +OpenVLA checkpoints use a custom ``model_type`` and nest the language model
6 +configuration under ``text_config``. This shim lets vLLM load the checkpoint
7 +configuration without executing Hugging Face remote code.
8 +"""
9 +
10 +from typing import Any
11 +
12 +from transformers import LlamaConfig, PretrainedConfig
13 +
14 +
15 +class OpenVLAConfig(PretrainedConfig):
16 + """Configuration class for OpenVLA models."""
17 +
18 + model_type = "openvla"
19 +
20 + def __init__(
21 + self,
22 + timm_model_ids: list[str] | None = None,
23 + timm_override_act_layers: list[str | None] | None = None,
24 + image_sizes: list[int] | None = None,
25 + use_fused_vision_backbone: bool = True,
26 + image_token_index: int = 32000,
27 + n_action_bins: int = 256,
28 + text_config: dict[str, Any] | LlamaConfig | None = None,
29 + **kwargs: Any,
30 + ) -> None:
31 + kwargs.setdefault("architectures", ["OpenVLAForActionPrediction"])
32 + super().__init__(**kwargs)
33 +
34 + self.timm_model_ids = timm_model_ids or [
35 + "vit_large_patch14_reg4_dinov2.lvd142m",
36 + "vit_so400m_patch14_siglip_224",
37 + ]
38 + self.timm_override_act_layers = timm_override_act_layers or [None, None]
39 + self.image_sizes = image_sizes or [224, 224]
40 + self.use_fused_vision_backbone = use_fused_vision_backbone
41 + self.image_token_index = image_token_index
42 + self.n_action_bins = n_action_bins
43 +
44 + if text_config is None:
45 + text_config = LlamaConfig(architectures=["LlamaForCausalLM"])
46 + elif isinstance(text_config, dict):
47 + text_config = text_config.copy()
48 + text_config.setdefault("architectures", ["LlamaForCausalLM"])
49 + text_config = LlamaConfig(**text_config)
50 + self.text_config = text_config
vllm/transformers_utils/processors/__init__.py
+2
@@ -35,6 +35,7 @@ __all__ = [
35 "NemotronVLProcessor",
36 "LlamaNemotronVLEmbedProcessor",
37 "NVLMProcessor",
38 + "OpenVLAProcessor",
39 "OvisProcessor",
40 "Ovis2_5Processor",
41 "QwenVLProcessor",
@@ -67,6 +68,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
68 "NemotronVLProcessor": "vllm.transformers_utils.processors.nemotron_vl",
69 "LlamaNemotronVLEmbedProcessor": "vllm.transformers_utils.processors.nemotron_vl",
70 "NVLMProcessor": "vllm.transformers_utils.processors.nvlm_d",
71 + "OpenVLAProcessor": "vllm.transformers_utils.processors.openvla",
72 "OvisProcessor": "vllm.transformers_utils.processors.ovis",
73 "Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5",
74 "QwenVLProcessor": "vllm.transformers_utils.processors.qwen_vl",
vllm/transformers_utils/processors/openvla.py new
+102
@@ -0,0 +1,102 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Sequence
5 +from typing import Any
6 +
7 +import numpy as np
8 +import torch
9 +from PIL import Image
10 +from transformers.processing_utils import ProcessorMixin
11 +
12 +IMAGENET_MEAN = np.array([0.484375, 0.455078125, 0.40625], dtype=np.float32)
13 +IMAGENET_STD = np.array([0.228515625, 0.2236328125, 0.224609375], dtype=np.float32)
14 +SIGLIP_MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32)
15 +SIGLIP_STD = np.array([0.5, 0.5, 0.5], dtype=np.float32)
16 +
17 +
18 +def to_rgb_image(image: Any) -> Image.Image:
19 + if isinstance(image, Image.Image):
20 + return image.convert("RGB")
21 +
22 + if isinstance(image, torch.Tensor):
23 + image = image.detach().cpu().numpy()
24 + if not isinstance(image, np.ndarray):
25 + raise TypeError(
26 + "OpenVLA image input must be a PIL image, numpy array, or torch tensor; "
27 + f"got {type(image)}"
28 + )
29 +
30 + if image.ndim != 3:
31 + raise ValueError(
32 + f"OpenVLA image input must have 3 dimensions, got shape {image.shape}"
33 + )
34 +
35 + if image.shape[0] in (1, 3):
36 + image = np.moveaxis(image, 0, -1)
37 +
38 + if image.shape[-1] == 1:
39 + image = np.repeat(image, 3, axis=-1)
40 + elif image.shape[-1] != 3:
41 + raise ValueError(
42 + f"OpenVLA image input must have 1 or 3 channels, got shape {image.shape}"
43 + )
44 +
45 + if image.dtype != np.uint8:
46 + image = image.astype(np.float32)
47 + if image.max(initial=0.0) <= 1.0:
48 + image = image * 255.0
49 + image = np.clip(image, 0, 255).astype(np.uint8)
50 +
51 + return Image.fromarray(image).convert("RGB")
52 +
53 +
54 +def preprocess_openvla_image(image: Any, image_size: int) -> torch.Tensor:
55 + rgb_image = to_rgb_image(image)
56 + rgb_image = rgb_image.resize(
57 + (image_size, image_size),
58 + Image.Resampling.BICUBIC,
59 + )
60 +
61 + raw = np.asarray(rgb_image, dtype=np.float32) / 255.0
62 + dinov2_pixels = ((raw - IMAGENET_MEAN) / IMAGENET_STD).transpose(2, 0, 1)
63 + siglip_pixels = ((raw - SIGLIP_MEAN) / SIGLIP_STD).transpose(2, 0, 1)
64 + pixel_values = np.concatenate([dinov2_pixels, siglip_pixels], axis=0)
65 + return torch.from_numpy(pixel_values)
66 +
67 +
68 +class OpenVLAImageProcessor:
69 + def __init__(self, *, image_size: int) -> None:
70 + self.image_size = image_size
71 +
72 + def __call__(
73 + self,
74 + images: Any | None = None,
75 + **kwargs: object,
76 + ) -> dict[str, object]:
77 + if images is None:
78 + return {}
79 + if not isinstance(images, Sequence) or isinstance(images, (str, bytes)):
80 + images = [images]
81 + if len(images) == 0:
82 + return {}
83 +
84 + pixel_values = torch.stack(
85 + [
86 + preprocess_openvla_image(image, image_size=self.image_size)
87 + for image in images
88 + ],
89 + dim=0,
90 + )
91 + return {"pixel_values": pixel_values}
92 +
93 +
94 +class OpenVLAProcessor(ProcessorMixin):
95 + def __init__(
96 + self,
97 + *,
98 + image_processor: OpenVLAImageProcessor,
99 + tokenizer: Any,
100 + ) -> None:
101 + self.image_processor = image_processor
102 + self.tokenizer = tokenizer