| 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 MultiModalDummyOptions |
| 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 ( |
| 28 | MultiModalFieldConfig, |
| 29 | MultiModalKwargsItem, |
| 30 | MultiModalKwargsItems, |
| 31 | ) |
| 32 | from vllm.multimodal.parse import ( |
| 33 | ImageEmbeddingItems, |
| 34 | ImageProcessorItems, |
| 35 | ImageSize, |
| 36 | MultiModalDataItems, |
| 37 | ) |
| 38 | from vllm.multimodal.processing import ( |
| 39 | BaseDummyInputsBuilder, |
| 40 | BaseMultiModalProcessor, |
| 41 | BaseProcessingInfo, |
| 42 | InputProcessingContext, |
| 43 | PromptIndexTargets, |
| 44 | PromptInsertion, |
| 45 | PromptUpdate, |
| 46 | PromptUpdateDetails, |
| 47 | ) |
| 48 | from vllm.sequence import IntermediateTensors |
| 49 | from vllm.transformers_utils.configs import OpenVLAConfig |
| 50 | from vllm.transformers_utils.processors.openvla import ( |
| 51 | OpenVLAImageProcessor, |
| 52 | OpenVLAProcessor, |
| 53 | ) |
| 54 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 55 | |
| 56 | from .module_mapping import MultiModelKeys |
| 57 | from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix |
| 58 | |
| 59 | # openvla/openvla-7b uses 224x224 images with ViT patch size 14, yielding a |
| 60 | # 16x16 image-token grid. |
| 61 | _OPENVLA_IMAGE_SIZE = 224 |
| 62 | _OPENVLA_PATCH_SIZE = 14 |
| 63 | _OPENVLA_TIMM_MODEL_IDS = ( |
| 64 | "vit_large_patch14_reg4_dinov2.lvd142m", |
| 65 | "vit_so400m_patch14_siglip_224", |
| 66 | ) |
| 67 | _OPENVLA_TIMM_OVERRIDE_ACT_LAYERS = (None, None) |
| 68 | _OPENVLA_IMAGE_SIZES = (_OPENVLA_IMAGE_SIZE, _OPENVLA_IMAGE_SIZE) |
| 69 | |
| 70 | |
| 71 | def _get_num_image_tokens(image_size: int) -> int: |
| 72 | return (image_size // _OPENVLA_PATCH_SIZE) ** 2 |
| 73 | |
| 74 | |
| 75 | class OpenVLAImagePixelInputs(TensorSchema): |
| 76 | """Dimensions: |
| 77 | - bn: Batch size * number of images |
| 78 | - c: Number of channels (6) |
| 79 | - h: Height |
| 80 | - w: Width |
| 81 | """ |
| 82 | |
| 83 | type: Literal["pixel_values"] = "pixel_values" |
| 84 | data: Annotated[torch.Tensor, TensorShape("bn", 6, "h", "w")] |
| 85 | |
| 86 | |
| 87 | class PrismaticVisionBackbone(nn.Module): |
| 88 | """OpenVLA's fused DINOv2 + SigLIP vision backbone.""" |
| 89 | |
| 90 | def __init__( |
| 91 | self, |
| 92 | *, |
| 93 | image_sizes: Sequence[int], |
| 94 | timm_model_ids: Sequence[str], |
| 95 | timm_override_act_layers: Sequence[str | None], |
| 96 | use_fused_vision_backbone: bool, |
| 97 | ) -> None: |
| 98 | super().__init__() |
| 99 | if not use_fused_vision_backbone: |
| 100 | raise ValueError( |
| 101 | "OpenVLA currently supports only the fused DINOv2 + SigLIP " |
| 102 | "vision backbone." |
| 103 | ) |
| 104 | if tuple(image_sizes) != _OPENVLA_IMAGE_SIZES: |
| 105 | raise ValueError( |
| 106 | "OpenVLA currently supports only 224x224 image inputs, " |
| 107 | f"got image_sizes={list(image_sizes)}." |
| 108 | ) |
| 109 | if tuple(timm_model_ids) != _OPENVLA_TIMM_MODEL_IDS: |
| 110 | raise ValueError( |
| 111 | "OpenVLA currently supports only the dinosiglip-vit-so-224px " |
| 112 | "vision backbone, got " |
| 113 | f"timm_model_ids={list(timm_model_ids)}." |
| 114 | ) |
| 115 | if tuple(timm_override_act_layers) != _OPENVLA_TIMM_OVERRIDE_ACT_LAYERS: |
| 116 | raise ValueError( |
| 117 | "OpenVLA currently supports only the default timm activation " |
| 118 | "layers, got " |
| 119 | f"timm_override_act_layers={list(timm_override_act_layers)}." |
| 120 | ) |
| 121 | |
| 122 | self.image_size = image_sizes[0] |
| 123 | self.use_fused_vision_backbone = use_fused_vision_backbone |
| 124 | |
| 125 | self.embed_dim = 2176 if use_fused_vision_backbone else 1024 |
| 126 | |
| 127 | try: |
| 128 | import timm |
| 129 | except ImportError as e: |
| 130 | raise ImportError( |
| 131 | "Please install timm to use OpenVLA. OpenVLA verification " |
| 132 | "used timm==0.9.10." |
| 133 | ) from e |
| 134 | |
| 135 | self.dinov2_featurizer = timm.create_model( |
| 136 | timm_model_ids[0], |
| 137 | pretrained=False, |
| 138 | num_classes=0, |
| 139 | img_size=self.image_size, |
| 140 | act_layer=timm_override_act_layers[0], |
| 141 | ) |
| 142 | self.siglip_featurizer = ( |
| 143 | timm.create_model( |
| 144 | timm_model_ids[1], |
| 145 | pretrained=False, |
| 146 | num_classes=0, |
| 147 | img_size=self.image_size, |
| 148 | act_layer=timm_override_act_layers[1], |
| 149 | ) |
| 150 | if use_fused_vision_backbone |
| 151 | else None |
| 152 | ) |
| 153 | |
| 154 | def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: |
| 155 | if self.dinov2_featurizer is None: |
| 156 | raise RuntimeError("OpenVLA vision backbone is not initialized.") |
| 157 | |
| 158 | if self.use_fused_vision_backbone and pixel_values.shape[1] != 6: |
| 159 | raise ValueError( |
| 160 | "OpenVLA fused DINOv2 + SigLIP backbone expects 6-channel " |
| 161 | "image inputs: 3 DINOv2-normalized channels followed by 3 " |
| 162 | "SigLIP-normalized channels, " |
| 163 | f"got {pixel_values.shape[1]} channels." |
| 164 | ) |
| 165 | |
| 166 | dinov2_pixels = pixel_values[:, :3] |
| 167 | |
| 168 | num_dinov2_blocks = len(self.dinov2_featurizer.blocks) |
| 169 | dinov2_features = self.dinov2_featurizer.get_intermediate_layers( |
| 170 | dinov2_pixels, n={num_dinov2_blocks - 2} |
| 171 | )[0] |
| 172 | |
| 173 | if self.siglip_featurizer is not None: |
| 174 | siglip_pixels = pixel_values[:, 3:] |
| 175 | num_siglip_blocks = len(self.siglip_featurizer.blocks) |
| 176 | siglip_features = self.siglip_featurizer.get_intermediate_layers( |
| 177 | siglip_pixels, n={num_siglip_blocks - 2} |
| 178 | )[0] |
| 179 | return torch.cat([dinov2_features, siglip_features], dim=-1) |
| 180 | |
| 181 | return dinov2_features |
| 182 | |
| 183 | |
| 184 | class PrismaticProjector(nn.Module): |
| 185 | """Project Prismatic vision features into the language-model hidden size.""" |
| 186 | |
| 187 | def __init__( |
| 188 | self, |
| 189 | *, |
| 190 | vision_dim: int, |
| 191 | text_dim: int, |
| 192 | use_fused_vision_backbone: bool, |
| 193 | quant_config: QuantizationConfig | None = None, |
| 194 | prefix: str = "", |
| 195 | ) -> None: |
| 196 | super().__init__() |
| 197 | self.use_fused_vision_backbone = use_fused_vision_backbone |
| 198 | |
| 199 | if use_fused_vision_backbone: |
| 200 | intermediate_dim = 4 * vision_dim |
| 201 | self.fc1 = ColumnParallelLinear( |
| 202 | vision_dim, |
| 203 | intermediate_dim, |
| 204 | bias=True, |
| 205 | quant_config=quant_config, |
| 206 | prefix=f"{prefix}.fc1", |
| 207 | ) |
| 208 | self.act_fn1 = get_act_fn("gelu") |
| 209 | self.fc2 = RowParallelLinear( |
| 210 | intermediate_dim, |
| 211 | text_dim, |
| 212 | bias=True, |
| 213 | quant_config=quant_config, |
| 214 | prefix=f"{prefix}.fc2", |
| 215 | ) |
| 216 | self.act_fn2 = get_act_fn("gelu") |
| 217 | self.fc3 = ReplicatedLinear( |
| 218 | text_dim, |
| 219 | text_dim, |
| 220 | bias=True, |
| 221 | quant_config=quant_config, |
| 222 | prefix=f"{prefix}.fc3", |
| 223 | ) |
| 224 | else: |
| 225 | self.fc1 = ColumnParallelLinear( |
| 226 | vision_dim, |
| 227 | text_dim, |
| 228 | bias=True, |
| 229 | quant_config=quant_config, |
| 230 | prefix=f"{prefix}.fc1", |
| 231 | ) |
| 232 | self.act_fn1 = get_act_fn("gelu") |
| 233 | self.fc2 = RowParallelLinear( |
| 234 | text_dim, |
| 235 | text_dim, |
| 236 | bias=True, |
| 237 | quant_config=quant_config, |
| 238 | prefix=f"{prefix}.fc2", |
| 239 | ) |
| 240 | |
| 241 | def forward(self, image_features: torch.Tensor) -> torch.Tensor: |
| 242 | hidden_states, _ = self.fc1(image_features) |
| 243 | hidden_states = self.act_fn1(hidden_states) |
| 244 | hidden_states, _ = self.fc2(hidden_states) |
| 245 | |
| 246 | if self.use_fused_vision_backbone: |
| 247 | hidden_states = self.act_fn2(hidden_states) |
| 248 | hidden_states, _ = self.fc3(hidden_states) |
| 249 | |
| 250 | return hidden_states |
| 251 | |
| 252 | |
| 253 | class OpenVLAProcessingInfo(BaseProcessingInfo): |
| 254 | def __init__(self, ctx: InputProcessingContext) -> None: |
| 255 | super().__init__(ctx) |
| 256 | self.hf_processor = OpenVLAProcessor( |
| 257 | image_processor=OpenVLAImageProcessor( |
| 258 | image_size=self.get_hf_config().image_sizes[0], |
| 259 | ), |
| 260 | tokenizer=self.get_tokenizer(), |
| 261 | ) |
| 262 | |
| 263 | def get_hf_config(self) -> OpenVLAConfig: |
| 264 | return self.ctx.get_hf_config(OpenVLAConfig) |
| 265 | |
| 266 | def get_hf_processor(self, **kwargs: object) -> OpenVLAProcessor: |
| 267 | return self.hf_processor |
| 268 | |
| 269 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 270 | return {"image": 1} |
| 271 | |
| 272 | def get_num_image_tokens( |
| 273 | self, |
| 274 | *, |
| 275 | image_width: int, |
| 276 | image_height: int, |
| 277 | ) -> int: |
| 278 | image_size = self.get_hf_config().image_sizes[0] |
| 279 | return _get_num_image_tokens(image_size) |
| 280 | |
| 281 | def get_image_size_with_most_features(self) -> ImageSize: |
| 282 | image_size = self.get_hf_config().image_sizes[0] |
| 283 | return ImageSize(width=image_size, height=image_size) |
| 284 | |
| 285 | def get_mm_max_tokens_per_item( |
| 286 | self, |
| 287 | seq_len: int, |
| 288 | mm_counts: Mapping[str, int], |
| 289 | ) -> Mapping[str, int] | None: |
| 290 | image_size = self.get_hf_config().image_sizes[0] |
| 291 | return {"image": _get_num_image_tokens(image_size)} |
| 292 | |
| 293 | |
| 294 | class OpenVLADummyInputsBuilder(BaseDummyInputsBuilder[OpenVLAProcessingInfo]): |
| 295 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 296 | return "" |
| 297 | |
| 298 | def get_dummy_mm_data( |
| 299 | self, |
| 300 | seq_len: int, |
| 301 | mm_counts: Mapping[str, int], |
| 302 | mm_options: MultiModalDummyOptions, |
| 303 | ) -> MultiModalDataDict: |
| 304 | image_size = self.info.get_image_size_with_most_features() |
| 305 | |
| 306 | return { |
| 307 | "image": self._get_dummy_images( |
| 308 | width=image_size.width, |
| 309 | height=image_size.height, |
| 310 | num_images=mm_counts.get("image", 0), |
| 311 | overrides=mm_options.get("image"), |
| 312 | ) |
| 313 | } |
| 314 | |
| 315 | |
| 316 | class OpenVLAMultiModalProcessor(BaseMultiModalProcessor[OpenVLAProcessingInfo]): |
| 317 | """Processor contract for OpenVLA image inputs. |
| 318 | |
| 319 | OpenVLA feeds the same RGB image to DINOv2 and SigLIP after different |
| 320 | normalizations. The processor exposes this as one 6-channel tensor: |
| 321 | channels 0-2 are DINOv2-normalized and channels 3-5 are SigLIP-normalized. |
| 322 | """ |
| 323 | |
| 324 | def _get_mm_fields_config( |
| 325 | self, |
| 326 | hf_inputs: BatchFeature, |
| 327 | hf_processor_mm_kwargs: Mapping[str, object], |
| 328 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 329 | return dict(pixel_values=MultiModalFieldConfig.batched("image")) |
| 330 | |
| 331 | def _get_prompt_updates( |
| 332 | self, |
| 333 | mm_items: MultiModalDataItems, |
| 334 | hf_processor_mm_kwargs: Mapping[str, object], |
| 335 | out_mm_kwargs: MultiModalKwargsItems, |
| 336 | ) -> Sequence[PromptUpdate]: |
| 337 | hf_config = self.info.get_hf_config() |
| 338 | image_token_id = hf_config.image_token_index |
| 339 | |
| 340 | tokenizer = self.info.get_tokenizer() |
| 341 | bos_token_id = tokenizer.bos_token_id |
| 342 | |
| 343 | def get_insertion(item_idx: int) -> PromptUpdateDetails: |
| 344 | images = mm_items.get_items( |
| 345 | "image", (ImageEmbeddingItems, ImageProcessorItems) |
| 346 | ) |
| 347 | if isinstance(images, ImageEmbeddingItems): |
| 348 | num_image_tokens = images.get_feature_size(item_idx) |
| 349 | else: |
| 350 | assert isinstance(images, ImageProcessorItems) |
| 351 | image_size = images.get_image_size(item_idx) |
| 352 | num_image_tokens = self.info.get_num_image_tokens( |
| 353 | image_width=image_size.width, |
| 354 | image_height=image_size.height, |
| 355 | ) |
| 356 | |
| 357 | image_tokens = [image_token_id] * num_image_tokens |
| 358 | return PromptUpdateDetails.select_token_id( |
| 359 | image_tokens, |
| 360 | embed_token_id=image_token_id, |
| 361 | ) |
| 362 | |
| 363 | return [ |
| 364 | PromptInsertion( |
| 365 | modality="image", |
| 366 | target=PromptIndexTargets.prefix( |
| 367 | [bos_token_id] if bos_token_id is not None else [] |
| 368 | ), |
| 369 | insertion=get_insertion, |
| 370 | ) |
| 371 | ] |
| 372 | |
| 373 | |
| 374 | @MULTIMODAL_REGISTRY.register_processor( |
| 375 | OpenVLAMultiModalProcessor, |
| 376 | info=OpenVLAProcessingInfo, |
| 377 | dummy_inputs=OpenVLADummyInputsBuilder, |
| 378 | ) |
| 379 | class OpenVLAForActionPrediction(nn.Module, SupportsMultiModal, SupportsPP): |
| 380 | """OpenVLA wrapper with vLLM language-model execution wired in.""" |
| 381 | |
| 382 | @classmethod |
| 383 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 384 | if modality.startswith("image"): |
| 385 | return None |
| 386 | raise ValueError("Only image modality is supported") |
| 387 | |
| 388 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 389 | super().__init__() |
| 390 | config = vllm_config.model_config.hf_config |
| 391 | quant_config = vllm_config.quant_config |
| 392 | self.config = config |
| 393 | self.multimodal_config = vllm_config.model_config.multimodal_config |
| 394 | self.image_token_id = config.image_token_index |
| 395 | self.n_action_bins = config.n_action_bins |
| 396 | self.num_patches = _get_num_image_tokens(config.image_sizes[0]) |
| 397 | |
| 398 | with self._mark_tower_model(vllm_config, "image"): |
| 399 | self.vision_backbone = PrismaticVisionBackbone( |
| 400 | image_sizes=config.image_sizes, |
| 401 | timm_model_ids=config.timm_model_ids, |
| 402 | timm_override_act_layers=config.timm_override_act_layers, |
| 403 | use_fused_vision_backbone=config.use_fused_vision_backbone, |
| 404 | ) |
| 405 | self.projector = PrismaticProjector( |
| 406 | vision_dim=self.vision_backbone.embed_dim, |
| 407 | text_dim=config.text_config.hidden_size, |
| 408 | use_fused_vision_backbone=config.use_fused_vision_backbone, |
| 409 | quant_config=quant_config, |
| 410 | prefix=maybe_prefix(prefix, "projector"), |
| 411 | ) |
| 412 | |
| 413 | with self._mark_language_model(vllm_config): |
| 414 | self.language_model = init_vllm_registered_model( |
| 415 | vllm_config=vllm_config, |
| 416 | hf_config=config.text_config, |
| 417 | prefix=maybe_prefix(prefix, "language_model"), |
| 418 | ) |
| 419 | |
| 420 | self.make_empty_intermediate_tensors = ( |
| 421 | self.language_model.make_empty_intermediate_tensors |
| 422 | ) |
| 423 | |
| 424 | def get_language_model(self) -> nn.Module: |
| 425 | return self.language_model |
| 426 | |
| 427 | def _parse_and_validate_image_input( |
| 428 | self, |
| 429 | **kwargs: object, |
| 430 | ) -> OpenVLAImagePixelInputs | None: |
| 431 | pixel_values = kwargs.pop("pixel_values", None) |
| 432 | if pixel_values is None: |
| 433 | return None |
| 434 | |
| 435 | return OpenVLAImagePixelInputs( |
| 436 | type="pixel_values", |
| 437 | data=pixel_values, |
| 438 | resolve_bindings={ |
| 439 | "h": self.config.image_sizes[0], |
| 440 | "w": self.config.image_sizes[0], |
| 441 | }, |
| 442 | ) |
| 443 | |
| 444 | def _process_image_input( |
| 445 | self, |
| 446 | image_input: OpenVLAImagePixelInputs, |
| 447 | ) -> torch.Tensor: |
| 448 | if self.vision_backbone.dinov2_featurizer is None: |
| 449 | raise RuntimeError("OpenVLA vision backbone is not initialized.") |
| 450 | |
| 451 | pixel_values = image_input["data"].to( |
| 452 | dtype=self.vision_backbone.dinov2_featurizer.patch_embed.proj.weight.dtype |
| 453 | ) |
| 454 | vision_features = self.vision_backbone(pixel_values) |
| 455 | return self.projector(vision_features) |
| 456 | |
| 457 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 458 | image_input = self._parse_and_validate_image_input(**kwargs) |
| 459 | if image_input is None: |
| 460 | return [] |
| 461 | |
| 462 | return self._process_image_input(image_input) |
| 463 | |
| 464 | def forward( |
| 465 | self, |
| 466 | input_ids: torch.Tensor | None, |
| 467 | positions: torch.Tensor, |
| 468 | intermediate_tensors: IntermediateTensors | None = None, |
| 469 | inputs_embeds: torch.Tensor | None = None, |
| 470 | **kwargs: object, |
| 471 | ) -> torch.Tensor | IntermediateTensors: |
| 472 | if intermediate_tensors is not None: |
| 473 | inputs_embeds = None |
| 474 | |
| 475 | return self.language_model.model( |
| 476 | input_ids, |
| 477 | positions, |
| 478 | intermediate_tensors, |
| 479 | inputs_embeds=inputs_embeds, |
| 480 | ) |
| 481 | |
| 482 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: |
| 483 | return self.language_model.compute_logits(hidden_states) |
| 484 | |
| 485 | def get_mm_mapping(self) -> MultiModelKeys: |
| 486 | return MultiModelKeys.from_string_field( |
| 487 | language_model="language_model", |
| 488 | connector="projector", |
| 489 | tower_model="vision_backbone", |
| 490 | ) |
| 491 | |
| 492 | def get_mm_lora_token_counts( |
| 493 | self, |
| 494 | *, |
| 495 | modality: str, |
| 496 | mm_kwargs: MultiModalKwargsItem | None, |
| 497 | num_mm_embeds: int, |
| 498 | ) -> tuple[int, int | None]: |
| 499 | del modality, mm_kwargs |
| 500 | return num_mm_embeds, num_mm_embeds |
| 501 | |
| 502 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 503 | def maybe_rename_vision_weights( |
| 504 | weights: Iterable[tuple[str, torch.Tensor]], |
| 505 | ) -> Iterable[tuple[str, torch.Tensor]]: |
| 506 | for name, weight in weights: |
| 507 | if name.startswith("vision_backbone.featurizer."): |
| 508 | name = name.replace( |
| 509 | "vision_backbone.featurizer.", |
| 510 | "vision_backbone.dinov2_featurizer.", |
| 511 | 1, |
| 512 | ) |
| 513 | elif name.startswith("vision_backbone.fused_featurizer."): |
| 514 | name = name.replace( |
| 515 | "vision_backbone.fused_featurizer.", |
| 516 | "vision_backbone.siglip_featurizer.", |
| 517 | 1, |
| 518 | ) |
| 519 | # HF uses .scale_factor, timm uses .gamma |
| 520 | if ".ls1.scale_factor" in name or ".ls2.scale_factor" in name: |
| 521 | name = name.replace(".scale_factor", ".gamma") |
| 522 | yield name, weight |
| 523 | |
| 524 | loader = AutoWeightsLoader(self) |
| 525 | return loader.load_weights(maybe_rename_vision_weights(weights)) |