| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import abstractmethod |
| 5 | from collections.abc import Iterable, Mapping |
| 6 | from typing import Annotated, Literal, Protocol, TypeAlias, TypeVar |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | from transformers import BatchFeature, LlavaNextConfig, LlavaNextProcessor |
| 11 | from transformers.models.llava_next.modeling_llava_next import ( |
| 12 | get_anyres_image_grid_shape, |
| 13 | unpad_image, |
| 14 | ) |
| 15 | |
| 16 | from vllm.config import VllmConfig |
| 17 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 18 | from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItem |
| 19 | from vllm.multimodal.parse import ImageSize |
| 20 | from vllm.sequence import IntermediateTensors |
| 21 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 22 | |
| 23 | from .clip import CLIPVisionModel |
| 24 | from .interfaces import ( |
| 25 | MultiModalEmbeddings, |
| 26 | SupportsLoRA, |
| 27 | SupportsMultiModal, |
| 28 | SupportsPP, |
| 29 | ) |
| 30 | from .llava import ( |
| 31 | BaseLlavaMultiModalProcessor, |
| 32 | BaseLlavaProcessingInfo, |
| 33 | LlavaDummyInputsBuilder, |
| 34 | LlavaLikeConfig, |
| 35 | LlavaMultiModalProjector, |
| 36 | init_vision_tower_for_llava, |
| 37 | ) |
| 38 | from .module_mapping import MultiModelKeys |
| 39 | from .siglip import SiglipVisionModel |
| 40 | from .utils import ( |
| 41 | AutoWeightsLoader, |
| 42 | WeightsMapper, |
| 43 | init_vllm_registered_model, |
| 44 | maybe_prefix, |
| 45 | ) |
| 46 | from .vision import get_num_selected_vision_tokens, get_vision_encoder_info |
| 47 | |
| 48 | |
| 49 | class LlavaNextImagePixelInputs(TensorSchema): |
| 50 | """Dimensions: |
| 51 | - bn: Batch size * number of images |
| 52 | - np: Number of patches + 1 |
| 53 | - c: Number of channels (3) |
| 54 | - h: Height |
| 55 | - w: Width |
| 56 | |
| 57 | Note that `num_patches` may be different per batch and image, |
| 58 | in which case the data is passed as a list instead of a batched tensor. |
| 59 | """ |
| 60 | |
| 61 | type: Literal["pixel_values"] = "pixel_values" |
| 62 | pixel_values: Annotated[ |
| 63 | torch.Tensor | list[torch.Tensor], |
| 64 | TensorShape("bn", "np", 3, "h", "w", dynamic_dims={"np"}), |
| 65 | ] |
| 66 | |
| 67 | image_sizes: Annotated[torch.Tensor | None, TensorShape("bn", 2)] |
| 68 | # This should be in `(height, width)` format. |
| 69 | |
| 70 | |
| 71 | class LlavaNextImageEmbeddingInputs(TensorSchema): |
| 72 | """Dimensions: |
| 73 | - bn: Batch size * number of images |
| 74 | - ifs: Image feature size |
| 75 | - hs: Hidden size (must match language model backbone) |
| 76 | """ |
| 77 | |
| 78 | type: Literal["image_embeds"] = "image_embeds" |
| 79 | data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")] |
| 80 | |
| 81 | |
| 82 | LlavaNextImageInputs: TypeAlias = ( |
| 83 | LlavaNextImagePixelInputs | LlavaNextImageEmbeddingInputs |
| 84 | ) |
| 85 | """Alias for supported LLaVA-NeXT image input types.""" |
| 86 | |
| 87 | |
| 88 | class LlavaNextLikeConfig(LlavaLikeConfig, Protocol): |
| 89 | image_grid_pinpoints: list[list[int]] |
| 90 | |
| 91 | |
| 92 | class LlavaNextProcessingInfo(BaseLlavaProcessingInfo): |
| 93 | def get_hf_config(self) -> LlavaNextLikeConfig: |
| 94 | return self.ctx.get_hf_config(LlavaNextConfig) |
| 95 | |
| 96 | def get_hf_processor(self, **kwargs: object): |
| 97 | hf_processor = self.ctx.get_hf_processor(LlavaNextProcessor, **kwargs) |
| 98 | |
| 99 | # In case patch_size is omitted from `processor_config.json` |
| 100 | # e.g. for E5-V: https://huggingface.co/royokong/e5-v |
| 101 | if hf_processor.patch_size is None: |
| 102 | patch_size = self.get_vision_encoder_info().get_patch_size() |
| 103 | hf_processor.patch_size = patch_size |
| 104 | |
| 105 | return hf_processor |
| 106 | |
| 107 | # Based on: https://github.com/huggingface/text-generation-inference/blob/v3.0.1/server/text_generation_server/models/vlm_causal_lm.py#L113 |
| 108 | def get_num_image_tokens( |
| 109 | self, |
| 110 | *, |
| 111 | image_width: int, |
| 112 | image_height: int, |
| 113 | ) -> int: |
| 114 | """Get the number of image tokens for the given image dimensions.""" |
| 115 | hf_config = self.get_hf_config() |
| 116 | vision_encoder_info = self.get_vision_encoder_info() |
| 117 | |
| 118 | base_feature_size = get_num_selected_vision_tokens( |
| 119 | vision_encoder_info.get_num_image_tokens( |
| 120 | image_width=image_width, |
| 121 | image_height=image_height, |
| 122 | ), |
| 123 | hf_config.vision_feature_select_strategy, |
| 124 | ) |
| 125 | |
| 126 | num_patch_height, num_patch_width = get_anyres_image_grid_shape( |
| 127 | image_size=(image_height, image_width), |
| 128 | grid_pinpoints=hf_config.image_grid_pinpoints, |
| 129 | patch_size=vision_encoder_info.get_image_size(), |
| 130 | ) |
| 131 | |
| 132 | ( |
| 133 | unpadded_feature_size, |
| 134 | newline_feature_size, |
| 135 | ) = self._get_num_unpadded_features( |
| 136 | original_height=image_height, |
| 137 | original_width=image_width, |
| 138 | npatches=vision_encoder_info.get_patch_grid_length(), |
| 139 | num_patch_height=num_patch_height, |
| 140 | num_patch_width=num_patch_width, |
| 141 | ) |
| 142 | |
| 143 | return unpadded_feature_size + newline_feature_size + base_feature_size |
| 144 | |
| 145 | # Based on: https://github.com/huggingface/text-generation-inference/blob/v3.0.1/server/text_generation_server/models/vlm_causal_lm.py#L86 |
| 146 | def _get_num_unpadded_features( |
| 147 | self, |
| 148 | *, |
| 149 | original_height: int, |
| 150 | original_width: int, |
| 151 | npatches: int, |
| 152 | num_patch_height: int, |
| 153 | num_patch_width: int, |
| 154 | ) -> tuple[int, int]: |
| 155 | current_height = npatches * num_patch_height |
| 156 | current_width = npatches * num_patch_width |
| 157 | |
| 158 | aspect_ratio = original_width / original_height |
| 159 | current_aspect_ratio = current_width / current_height |
| 160 | |
| 161 | if aspect_ratio > current_aspect_ratio: |
| 162 | new_height = int( |
| 163 | round(original_height * (current_width / original_width), 7) |
| 164 | ) |
| 165 | padding = (current_height - new_height) // 2 |
| 166 | current_height = current_height - (2 * padding) |
| 167 | else: |
| 168 | new_width = int( |
| 169 | round(original_width * (current_height / original_height), 7) |
| 170 | ) |
| 171 | padding = (current_width - new_width) // 2 |
| 172 | current_width = current_width - (2 * padding) |
| 173 | |
| 174 | unpadded_features = current_height * current_width |
| 175 | newline_features = current_height |
| 176 | |
| 177 | return (unpadded_features, newline_features) |
| 178 | |
| 179 | def get_image_size_with_most_features(self) -> ImageSize: |
| 180 | hf_config = self.get_hf_config() |
| 181 | |
| 182 | largest_feature_size, largest_feature_pinpoint = 0, None |
| 183 | for height, width in hf_config.image_grid_pinpoints: |
| 184 | feat_size = self.get_num_image_tokens( |
| 185 | image_width=width, image_height=height |
| 186 | ) |
| 187 | if feat_size > largest_feature_size: |
| 188 | largest_feature_size = feat_size |
| 189 | largest_feature_pinpoint = ImageSize(width=width, height=height) |
| 190 | |
| 191 | if largest_feature_size == 0 or largest_feature_pinpoint is None: |
| 192 | raise ValueError("Cannot have a largest feature size of 0!") |
| 193 | |
| 194 | return largest_feature_pinpoint |
| 195 | |
| 196 | |
| 197 | _I = TypeVar("_I", bound=LlavaNextProcessingInfo) |
| 198 | |
| 199 | |
| 200 | class BaseLlavaNextMultiModalProcessor(BaseLlavaMultiModalProcessor[_I]): |
| 201 | # Copied from BaseMultiModalProcessor |
| 202 | @abstractmethod |
| 203 | def _get_mm_fields_config( |
| 204 | self, |
| 205 | hf_inputs: BatchFeature, |
| 206 | hf_processor_mm_kwargs: Mapping[str, object], |
| 207 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 208 | raise NotImplementedError |
| 209 | |
| 210 | |
| 211 | class LlavaNextMultiModalProcessor(BaseLlavaNextMultiModalProcessor[_I]): |
| 212 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 213 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 214 | |
| 215 | def _get_mm_fields_config( |
| 216 | self, |
| 217 | hf_inputs: BatchFeature, |
| 218 | hf_processor_mm_kwargs: Mapping[str, object], |
| 219 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 220 | return dict( |
| 221 | pixel_values=MultiModalFieldConfig.batched("image"), |
| 222 | image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 223 | image_embeds=MultiModalFieldConfig.batched("image"), |
| 224 | ) |
| 225 | |
| 226 | |
| 227 | @MULTIMODAL_REGISTRY.register_processor( |
| 228 | LlavaNextMultiModalProcessor, |
| 229 | info=LlavaNextProcessingInfo, |
| 230 | dummy_inputs=LlavaDummyInputsBuilder, |
| 231 | ) |
| 232 | class LlavaNextForConditionalGeneration( |
| 233 | nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP |
| 234 | ): |
| 235 | supports_tower_connector_lora = True |
| 236 | packed_modules_mapping = { |
| 237 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 238 | "gate_up_proj": ["gate_proj", "up_proj"], |
| 239 | } |
| 240 | |
| 241 | hf_to_vllm_mapper = WeightsMapper( |
| 242 | orig_to_new_prefix={ |
| 243 | # mapping for new names in checkpoint saved after transformers v4.52 |
| 244 | "model.language_model.": "language_model.model.", |
| 245 | "model.vision_tower.": "vision_tower.", |
| 246 | "model.multi_modal_projector.": "multi_modal_projector.", |
| 247 | "model.image_newline": "image_newline", |
| 248 | "lm_head.": "language_model.lm_head.", |
| 249 | } |
| 250 | ) |
| 251 | |
| 252 | @classmethod |
| 253 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 254 | if modality.startswith("image"): |
| 255 | return "<image>" |
| 256 | |
| 257 | raise ValueError("Only image modality is supported") |
| 258 | |
| 259 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 260 | super().__init__() |
| 261 | |
| 262 | config = vllm_config.model_config.hf_config |
| 263 | quant_config = vllm_config.quant_config |
| 264 | multimodal_config = vllm_config.model_config.multimodal_config |
| 265 | |
| 266 | vision_feature_layer = config.vision_feature_layer |
| 267 | # Determine the layer up to which we will initialize the vision tower |
| 268 | if isinstance(vision_feature_layer, int): |
| 269 | vision_hidden_size = config.vision_config.hidden_size |
| 270 | self.select_layers = None |
| 271 | # Used for multimodal granite models to control encoder outputs |
| 272 | elif isinstance(vision_feature_layer, (list, tuple)): |
| 273 | vision_hidden_size = config.vision_config.hidden_size * len( |
| 274 | vision_feature_layer |
| 275 | ) |
| 276 | self.select_layers = vision_feature_layer |
| 277 | else: |
| 278 | raise TypeError( |
| 279 | f"vision_layer_feature type: {type(vision_feature_layer)}" |
| 280 | " is not supported" |
| 281 | ) |
| 282 | |
| 283 | self.config = config |
| 284 | self.multimodal_config = multimodal_config |
| 285 | |
| 286 | self.configure_mm_token_handling( |
| 287 | vocab_size=config.text_config.vocab_size, |
| 288 | mm_token_ids=[config.image_token_index], |
| 289 | ) |
| 290 | |
| 291 | with self._mark_tower_model(vllm_config, "image"): |
| 292 | self.vision_tower = init_vision_tower_for_llava( |
| 293 | config, |
| 294 | quant_config=quant_config, |
| 295 | require_post_norm=False, |
| 296 | prefix=maybe_prefix(prefix, "vision_tower"), |
| 297 | ) |
| 298 | self.image_newline = nn.Parameter( |
| 299 | torch.empty(config.text_config.hidden_size) |
| 300 | ) |
| 301 | self.multi_modal_projector = LlavaMultiModalProjector( |
| 302 | vision_hidden_size=vision_hidden_size, |
| 303 | text_hidden_size=config.text_config.hidden_size, |
| 304 | projector_hidden_act=config.projector_hidden_act, |
| 305 | multimodal_projector_bias=config.multimodal_projector_bias, |
| 306 | quant_config=quant_config, |
| 307 | prefix=maybe_prefix(prefix, "multi_modal_projector"), |
| 308 | ) |
| 309 | |
| 310 | with self._mark_language_model(vllm_config): |
| 311 | self.language_model = init_vllm_registered_model( |
| 312 | vllm_config=vllm_config, |
| 313 | hf_config=config.text_config, |
| 314 | prefix=maybe_prefix(prefix, "language_model"), |
| 315 | ) |
| 316 | |
| 317 | self.make_empty_intermediate_tensors = ( |
| 318 | self.language_model.make_empty_intermediate_tensors |
| 319 | ) |
| 320 | |
| 321 | def _parse_and_validate_image_input( |
| 322 | self, **kwargs: object |
| 323 | ) -> LlavaNextImageInputs | None: |
| 324 | pixel_values = kwargs.pop("pixel_values", None) |
| 325 | image_sizes = kwargs.pop("image_sizes", None) |
| 326 | image_embeds = kwargs.pop("image_embeds", None) |
| 327 | |
| 328 | if pixel_values is None and image_embeds is None: |
| 329 | return None |
| 330 | |
| 331 | if pixel_values is not None: |
| 332 | expected_h = expected_w = self.config.vision_config.image_size |
| 333 | return LlavaNextImagePixelInputs( |
| 334 | type="pixel_values", |
| 335 | pixel_values=pixel_values, |
| 336 | image_sizes=image_sizes, |
| 337 | resolve_bindings={ |
| 338 | "h": expected_h, |
| 339 | "w": expected_w, |
| 340 | }, |
| 341 | ) |
| 342 | |
| 343 | if image_embeds is not None: |
| 344 | return LlavaNextImageEmbeddingInputs( |
| 345 | type="image_embeds", |
| 346 | data=image_embeds, |
| 347 | ) |
| 348 | |
| 349 | raise AssertionError("This line should be unreachable.") |
| 350 | |
| 351 | def _image_pixels_to_features( |
| 352 | self, |
| 353 | vision_tower: CLIPVisionModel | SiglipVisionModel, |
| 354 | pixel_values: torch.Tensor, |
| 355 | ) -> torch.Tensor: |
| 356 | # NOTE: we skip the step to select the vision feature layer since |
| 357 | # this is already done inside the vision tower |
| 358 | return vision_tower( |
| 359 | pixel_values, |
| 360 | select_layers=self.select_layers, |
| 361 | feature_select_strategy=self.config.vision_feature_select_strategy, |
| 362 | ) |
| 363 | |
| 364 | # Based on: https://github.com/haotian-liu/LLaVA/blob/main/llava/model/llava_arch.py |
| 365 | def _merge_image_patch_embeddings( |
| 366 | self, image_size: torch.Tensor, patch_embeddings: torch.Tensor, *, strategy: str |
| 367 | ) -> torch.Tensor: |
| 368 | if strategy == "flat": |
| 369 | return patch_embeddings.flatten(0, 1) |
| 370 | |
| 371 | if strategy.startswith("spatial"): |
| 372 | height = width = ( |
| 373 | self.config.vision_config.image_size |
| 374 | // self.config.vision_config.patch_size |
| 375 | ) |
| 376 | |
| 377 | base_patch_embeds = patch_embeddings[0] |
| 378 | if height * width != base_patch_embeds.shape[0]: |
| 379 | raise ValueError( |
| 380 | "The number of patches is not consistent with the image size." |
| 381 | ) |
| 382 | |
| 383 | if patch_embeddings.shape[0] > 1: |
| 384 | other_patch_embeds = patch_embeddings[1:] |
| 385 | |
| 386 | # Move to CPU to avoid floating-point errors |
| 387 | orig_height, orig_width = image_size.tolist() |
| 388 | |
| 389 | # image_aspect_ratio == "anyres" |
| 390 | num_patch_height, num_patch_width = get_anyres_image_grid_shape( |
| 391 | (orig_height, orig_width), |
| 392 | self.config.image_grid_pinpoints, |
| 393 | self.config.vision_config.image_size, |
| 394 | ) |
| 395 | num_patches = num_patch_height * num_patch_width |
| 396 | |
| 397 | # Image patches might be padded for batch processing |
| 398 | other_patch_embeds = other_patch_embeds[:num_patches].view( |
| 399 | num_patch_height, num_patch_width, height, width, -1 |
| 400 | ) |
| 401 | |
| 402 | if "unpad" in strategy: |
| 403 | other_patch_embeds = ( |
| 404 | other_patch_embeds.permute(4, 0, 2, 1, 3) |
| 405 | .contiguous() |
| 406 | .flatten(1, 2) |
| 407 | .flatten(2, 3) |
| 408 | ) |
| 409 | other_patch_embeds = unpad_image( |
| 410 | other_patch_embeds, (orig_height, orig_width) |
| 411 | ) |
| 412 | other_patch_embeds = torch.cat( |
| 413 | ( |
| 414 | other_patch_embeds, |
| 415 | self.image_newline[:, None, None] |
| 416 | .expand(*other_patch_embeds.shape[:-1], 1) |
| 417 | .to(other_patch_embeds.device), |
| 418 | ), |
| 419 | dim=-1, |
| 420 | ) |
| 421 | other_patch_embeds = other_patch_embeds.flatten(1, 2).transpose( |
| 422 | 0, 1 |
| 423 | ) |
| 424 | else: |
| 425 | other_patch_embeds = ( |
| 426 | other_patch_embeds.permute(0, 2, 1, 3, 4) |
| 427 | .contiguous() |
| 428 | .flatten(0, 3) |
| 429 | ) |
| 430 | |
| 431 | merged_patch_embeddings = torch.cat( |
| 432 | (base_patch_embeds, other_patch_embeds), dim=0 |
| 433 | ) |
| 434 | else: |
| 435 | if "unpad" in strategy: |
| 436 | merged_patch_embeddings = torch.cat( |
| 437 | ( |
| 438 | base_patch_embeds, |
| 439 | self.image_newline[None].to(base_patch_embeds.device), |
| 440 | ), |
| 441 | dim=0, |
| 442 | ) |
| 443 | else: |
| 444 | merged_patch_embeddings = base_patch_embeds |
| 445 | |
| 446 | return merged_patch_embeddings |
| 447 | |
| 448 | raise ValueError(f"Unexpected patch merge strategy: {strategy}") |
| 449 | |
| 450 | def _process_image_pixels( |
| 451 | self, |
| 452 | inputs: LlavaNextImagePixelInputs, |
| 453 | ) -> torch.Tensor | tuple[torch.Tensor, ...]: |
| 454 | pixel_values = inputs["pixel_values"] |
| 455 | |
| 456 | if isinstance(pixel_values, torch.Tensor): |
| 457 | b, num_patches, c, h, w = pixel_values.shape |
| 458 | stacked_pixel_values = pixel_values.view(b * num_patches, c, h, w) |
| 459 | stacked_image_features = self._image_pixels_to_features( |
| 460 | self.vision_tower, stacked_pixel_values |
| 461 | ) |
| 462 | stacked_patch_embeddings = self.multi_modal_projector( |
| 463 | stacked_image_features |
| 464 | ) |
| 465 | |
| 466 | return stacked_patch_embeddings.view( |
| 467 | b, num_patches, *stacked_patch_embeddings.shape[1:] |
| 468 | ) |
| 469 | |
| 470 | num_patches_per_batch = [v.shape[0] for v in pixel_values] |
| 471 | stacked_pixel_values = torch.cat(pixel_values) |
| 472 | stacked_image_features = self._image_pixels_to_features( |
| 473 | self.vision_tower, stacked_pixel_values |
| 474 | ) |
| 475 | |
| 476 | return torch.split( |
| 477 | self.multi_modal_projector(stacked_image_features), num_patches_per_batch |
| 478 | ) |
| 479 | |
| 480 | def _process_image_input( |
| 481 | self, |
| 482 | image_input: LlavaNextImageInputs, |
| 483 | ) -> torch.Tensor | list[torch.Tensor]: |
| 484 | if isinstance(image_input, LlavaNextImageEmbeddingInputs): |
| 485 | return image_input["data"] |
| 486 | |
| 487 | patch_embeddings = self._process_image_pixels(image_input) |
| 488 | |
| 489 | image_sizes = image_input.get("image_sizes") |
| 490 | if image_sizes is None: |
| 491 | batch_size = len(image_input["data"]) |
| 492 | vision_config = self.config.vision_config |
| 493 | default_height = default_width = vision_config.image_size |
| 494 | image_sizes = torch.as_tensor( |
| 495 | [[default_height, default_width] for _ in range(batch_size)] |
| 496 | ) |
| 497 | |
| 498 | return [ |
| 499 | self._merge_image_patch_embeddings( |
| 500 | image_sizes[i], patch_features_batch, strategy="spatial_unpad" |
| 501 | ) |
| 502 | for i, patch_features_batch in enumerate(patch_embeddings) |
| 503 | ] |
| 504 | |
| 505 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 506 | image_input = self._parse_and_validate_image_input(**kwargs) |
| 507 | if image_input is None: |
| 508 | return [] |
| 509 | vision_embeddings = self._process_image_input(image_input) |
| 510 | return vision_embeddings |
| 511 | |
| 512 | def embed_input_ids( |
| 513 | self, |
| 514 | input_ids: torch.Tensor, |
| 515 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 516 | *, |
| 517 | is_multimodal: torch.Tensor | None = None, |
| 518 | ) -> torch.Tensor: |
| 519 | # This is to satisfy the type checker for each overload |
| 520 | if multimodal_embeddings is None or is_multimodal is None: |
| 521 | return super().embed_input_ids(input_ids) |
| 522 | |
| 523 | return super().embed_input_ids( |
| 524 | input_ids, |
| 525 | multimodal_embeddings=multimodal_embeddings, |
| 526 | is_multimodal=is_multimodal, |
| 527 | ) |
| 528 | |
| 529 | def forward( |
| 530 | self, |
| 531 | input_ids: torch.Tensor | None, |
| 532 | positions: torch.Tensor, |
| 533 | intermediate_tensors: IntermediateTensors | None = None, |
| 534 | inputs_embeds: torch.Tensor | None = None, |
| 535 | **kwargs: object, |
| 536 | ) -> torch.Tensor | IntermediateTensors: |
| 537 | """Run forward pass for LlaVA-NeXT. |
| 538 | |
| 539 | One key thing to understand is the `input_ids` already accounts for the |
| 540 | positions of the to-be-inserted image embeddings. |
| 541 | |
| 542 | Concretely, consider a text prompt: |
| 543 | `"A chat between a curious human and an artificial intelligence |
| 544 | assistant. The assistant gives helpful, detailed, and polite answers to |
| 545 | the human's questions. |
| 546 | USER: <image>\\nWhat is shown in this image? ASSISTANT:"`. |
| 547 | |
| 548 | Tokenizer outputs: |
| 549 | `[1, 319, 13563, 1546, 263, 12758, 5199, 322, 385, 23116, 21082, 20255, |
| 550 | 29889, 450, 20255, 4076, 8444, 29892, 13173, 29892, 322, 1248, 568, |
| 551 | 6089, 304, 278, 5199, 29915, 29879, 5155, 29889, 3148, 1001, 29901, |
| 552 | 29871, 32000, 13, 5618, 338, 4318, 297, 445, 1967, 29973, 319, 1799, |
| 553 | 9047, 13566, 29901]`. |
| 554 | |
| 555 | To reserve space in KV cache, we have to insert placeholder tokens |
| 556 | before they are inputted to the model, so the input processor prepends |
| 557 | additional image tokens (denoted as `32000`), resulting in: |
| 558 | `[1, 319, 13563, 1546, 263, 12758, 5199, 322, 385, 23116, 21082, 20255, |
| 559 | 29889, 450, 20255, 4076, 8444, 29892, 13173, 29892, 322, 1248, 568, |
| 560 | 6089, 304, 278, 5199, 29915, 29879, 5155, 29889, 3148, 1001, 29901, |
| 561 | 29871, 32000, ..., 32000, 13, 5618, 338, 4318, 297, 445, 1967, 29973, |
| 562 | 319, 1799, 9047, 13566, 29901]`. |
| 563 | |
| 564 | Unlike in LLaVA-1.5, the number of image tokens inputted to the language |
| 565 | model depends on the original size of the input image. Including the |
| 566 | original image token in the input, the required number of image tokens |
| 567 | is given by [`LlavaNextProcessingInfo.get_num_image_tokens`][vllm.\ |
| 568 | model_executor.models.llava_next.LlavaNextProcessingInfo.get_num_image_tokens]. |
| 569 | |
| 570 | This way, the `positions` and `attn_metadata` are consistent |
| 571 | with the `input_ids`. |
| 572 | |
| 573 | Args: |
| 574 | input_ids: Flattened (concatenated) input_ids corresponding to a |
| 575 | batch. |
| 576 | positions: Position indices for the input tokens. |
| 577 | intermediate_tensors: Intermediate tensors from prior forward pass. |
| 578 | inputs_embeds: Optional tensor of input embeddings. |
| 579 | **kwargs: Multimodal inputs for this batch, forwarded to the |
| 580 | multimodal embedding path. |
| 581 | |
| 582 | Info: |
| 583 | [`LlavaNextImageInputs`][vllm.model_executor.models.llava_next.LlavaNextImageInputs] |
| 584 | |
| 585 | """ |
| 586 | if intermediate_tensors is not None: |
| 587 | inputs_embeds = None |
| 588 | |
| 589 | hidden_states = self.language_model.model( |
| 590 | input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds |
| 591 | ) |
| 592 | return hidden_states |
| 593 | |
| 594 | def compute_logits( |
| 595 | self, |
| 596 | hidden_states: torch.Tensor, |
| 597 | ) -> torch.Tensor | None: |
| 598 | return self.language_model.compute_logits(hidden_states) |
| 599 | |
| 600 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 601 | loader = AutoWeightsLoader(self) |
| 602 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 603 | |
| 604 | def get_mm_mapping(self) -> MultiModelKeys: |
| 605 | """Get the module prefix in multimodal models""" |
| 606 | return MultiModelKeys.from_string_field( |
| 607 | language_model="language_model", |
| 608 | connector="multi_modal_projector", |
| 609 | tower_model="vision_tower", |
| 610 | ) |
| 611 | |
| 612 | def get_mm_lora_token_counts( |
| 613 | self, |
| 614 | *, |
| 615 | modality: str, |
| 616 | mm_kwargs: MultiModalKwargsItem | None, |
| 617 | num_mm_embeds: int, |
| 618 | ) -> tuple[int, int | None]: |
| 619 | del modality |
| 620 | |
| 621 | pixel_values = mm_kwargs.get("pixel_values") if mm_kwargs else None |
| 622 | if pixel_values is None or not isinstance(pixel_values.data, torch.Tensor): |
| 623 | return num_mm_embeds, num_mm_embeds |
| 624 | |
| 625 | # Unpad runs after the connector, so `num_mm_embeds` is not invertible. |
| 626 | num_tiles = pixel_values.data.shape[0] |
| 627 | encoder_info = get_vision_encoder_info(self.config) |
| 628 | tile_size = encoder_info.get_image_size() |
| 629 | tokens_per_tile = encoder_info.get_num_image_tokens( |
| 630 | image_width=tile_size, image_height=tile_size |
| 631 | ) |
| 632 | selected_per_tile = get_num_selected_vision_tokens( |
| 633 | tokens_per_tile, self.config.vision_feature_select_strategy |
| 634 | ) |
| 635 | return num_tiles * tokens_per_tile, num_tiles * selected_per_tile |