| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | import typing |
| 5 | from collections.abc import Iterable, Mapping, Sequence |
| 6 | from typing import Annotated, Any, Literal, TypeAlias |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | import torch.nn as nn |
| 11 | from transformers import ( |
| 12 | BatchFeature, |
| 13 | PreTrainedConfig, |
| 14 | ProcessorMixin, |
| 15 | SequenceFeatureExtractor, |
| 16 | SiglipVisionConfig, |
| 17 | ) |
| 18 | |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.config.multimodal import MultiModalDummyOptions |
| 21 | from vllm.distributed import get_pp_group |
| 22 | from vllm.inputs import MultiModalDataDict |
| 23 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 24 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 25 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 26 | ParallelLMHead, |
| 27 | ) |
| 28 | from vllm.model_executor.models.llama import LlamaModel |
| 29 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 30 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 31 | from vllm.multimodal.inputs import ( |
| 32 | MultiModalFieldConfig, |
| 33 | MultiModalKwargsItems, |
| 34 | NestedTensors, |
| 35 | ) |
| 36 | from vllm.multimodal.parse import ( |
| 37 | AudioProcessorItems, |
| 38 | ImageEmbeddingItems, |
| 39 | ImageProcessorItems, |
| 40 | ImageSize, |
| 41 | MultiModalDataItems, |
| 42 | MultiModalDataParser, |
| 43 | ) |
| 44 | from vllm.multimodal.processing import BaseDummyInputsBuilder |
| 45 | from vllm.multimodal.processing.processor import ( |
| 46 | BaseMultiModalProcessor, |
| 47 | BaseProcessingInfo, |
| 48 | HFMultiModalInputs, |
| 49 | PromptReplacement, |
| 50 | PromptUpdate, |
| 51 | ResolvedPromptUpdate, |
| 52 | cached_encode, |
| 53 | ) |
| 54 | from vllm.sequence import IntermediateTensors |
| 55 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 56 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 57 | |
| 58 | from .idefics2_vision_model import Idefics2VisionTransformer |
| 59 | from .interfaces import MultiModalEmbeddings, SupportsLoRA, SupportsMultiModal |
| 60 | from .phi4mm_audio import AudioEmbedding |
| 61 | from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix |
| 62 | |
| 63 | # <|endoftext10|> (see vocab.json in hf model) |
| 64 | _IMAGE_PLACEHOLDER_TOKEN_ID = 200010 |
| 65 | # <|endoftext11|> |
| 66 | _AUDIO_PLACEHOLDER_TOKEN_ID = 200011 |
| 67 | |
| 68 | _AUDIO_MAX_SOUNDFILE_SIZE = 241_000 |
| 69 | |
| 70 | SIGLIP_NAME = "siglip-so400m-patch14-448" |
| 71 | VISION_ENCODER_TO_PROCESSING_CONFIG = { |
| 72 | "siglip-so400m-patch14-448": { |
| 73 | "vit_image_size": 448, |
| 74 | "vit_patch_size": 14, |
| 75 | "token_compression_factor": 2, |
| 76 | }, |
| 77 | } |
| 78 | |
| 79 | |
| 80 | def _get_padding_size( |
| 81 | orig_width: int, orig_height: int, target_height: int, target_width: int |
| 82 | ): |
| 83 | ratio_width = target_width / orig_width |
| 84 | ratio_height = target_height / orig_height |
| 85 | |
| 86 | if ratio_width < ratio_height: |
| 87 | padding_width = 0 |
| 88 | padding_height = target_height - int(orig_height * ratio_width) |
| 89 | else: |
| 90 | padding_width = target_width - int(orig_width * ratio_height) |
| 91 | padding_height = 0 |
| 92 | return padding_height, padding_width |
| 93 | |
| 94 | |
| 95 | def get_navit_vision_model(layer_idx: int = -1, **kwargs): |
| 96 | vision_config = { |
| 97 | "hidden_size": 1152, |
| 98 | "image_size": 448, |
| 99 | "intermediate_size": 4304, |
| 100 | "model_type": "siglip_vision_model", |
| 101 | "num_attention_heads": 16, |
| 102 | "num_hidden_layers": 27, |
| 103 | "patch_size": 14, |
| 104 | } |
| 105 | |
| 106 | model_config = SiglipVisionConfig(**vision_config, **kwargs) |
| 107 | if layer_idx < 0: |
| 108 | num_hidden_layers = model_config.num_hidden_layers + layer_idx + 1 |
| 109 | else: |
| 110 | num_hidden_layers = layer_idx + 1 |
| 111 | |
| 112 | vision_model = Idefics2VisionTransformer( |
| 113 | config=model_config, |
| 114 | require_post_norm=False, |
| 115 | num_hidden_layers_override=num_hidden_layers, |
| 116 | ) |
| 117 | |
| 118 | return vision_model |
| 119 | |
| 120 | |
| 121 | class Phi4MMImageEncoder(nn.Module): |
| 122 | """Image embedding.""" |
| 123 | |
| 124 | def __init__( |
| 125 | self, |
| 126 | config: PreTrainedConfig, |
| 127 | quant_config: QuantizationConfig | None, |
| 128 | prefix: str = "", |
| 129 | model_dir: str = "", |
| 130 | ) -> None: |
| 131 | super().__init__() |
| 132 | |
| 133 | # n_embed or hidden_size |
| 134 | hidden_size = config.n_embd if hasattr(config, "n_embd") else config.hidden_size |
| 135 | |
| 136 | # layer_idx to output the img features |
| 137 | if isinstance(config.img_processor, dict): |
| 138 | self.layer_idx = config.img_processor.get("layer_idx", -2) |
| 139 | self.type_feature = config.img_processor.get("type_feature", "patch") |
| 140 | else: |
| 141 | self.layer_idx = -2 |
| 142 | self.type_feature = "patch" |
| 143 | |
| 144 | self.img_processor = get_navit_vision_model(layer_idx=self.layer_idx) |
| 145 | |
| 146 | pe_weight = self.img_processor.embeddings.position_embedding.weight |
| 147 | L, D = pe_weight.size() |
| 148 | H = int(math.sqrt(L)) |
| 149 | assert H**2 == L, f"position embedding size {L} is not square" |
| 150 | if H % 2 != 0: |
| 151 | self.img_processor_padding = nn.ReflectionPad2d((0, 1, 0, 1)) |
| 152 | H += 1 |
| 153 | image_dim_out = D |
| 154 | # ((448/14)//2)**2 |
| 155 | self.num_img_tokens = (H // 2) ** 2 |
| 156 | self.base_feat_height_target = H |
| 157 | |
| 158 | self.image_dim_out = image_dim_out |
| 159 | self.img_sizes = None |
| 160 | self.image_attention_mask = None |
| 161 | |
| 162 | # global_gn and sub_gn for hd transform, serves as line separator |
| 163 | self.use_hd_transform = True |
| 164 | self.with_learnable_separator = True |
| 165 | self.hd_transform_order = "sub_glb" |
| 166 | self.freeze_img_processor = False |
| 167 | self.crop_size = 448 |
| 168 | |
| 169 | # image token compression |
| 170 | self.image_token_compression_cls = "avg_pool_2d" |
| 171 | self.image_token_compression = nn.AvgPool2d(kernel_size=2, stride=2) |
| 172 | self.base_feat_height_reduction = 1 |
| 173 | self.base_feat_height_target = self.base_feat_height_target // 2 |
| 174 | |
| 175 | # with_hd_transform and with_learnable_separator should have same value |
| 176 | assert self.use_hd_transform == self.with_learnable_separator, ( |
| 177 | "use_hd_transform and with_learnable_separator should have same value" |
| 178 | ) |
| 179 | assert self.use_hd_transform, "learnable separator is only for hd transform" |
| 180 | # 1024 * 4, merge spatial to channel dimension |
| 181 | self.glb_GN = nn.Parameter( |
| 182 | torch.zeros([1, 1, self.image_dim_out * self.base_feat_height_reduction**2]) |
| 183 | ) |
| 184 | self.sub_GN = nn.Parameter( |
| 185 | torch.zeros( |
| 186 | [1, 1, 1, self.image_dim_out * self.base_feat_height_reduction**2] |
| 187 | ) |
| 188 | ) |
| 189 | |
| 190 | dim_projection = hidden_size |
| 191 | depth = 2 |
| 192 | layers = [ |
| 193 | nn.Linear( |
| 194 | image_dim_out * self.base_feat_height_reduction**2, dim_projection |
| 195 | ) |
| 196 | ] |
| 197 | for _ in range(1, depth): |
| 198 | layers.extend([nn.GELU(), nn.Linear(dim_projection, dim_projection)]) |
| 199 | self.img_projection = nn.Sequential(*layers) |
| 200 | |
| 201 | self.vocab_size = config.vocab_size |
| 202 | self.img_features = None |
| 203 | |
| 204 | self.use_out_place_operations = False |
| 205 | |
| 206 | def get_img_features( |
| 207 | self, img_embeds: torch.FloatTensor, attention_mask=None |
| 208 | ) -> torch.FloatTensor: |
| 209 | img_feature = self.img_processor( |
| 210 | img_embeds, patch_attention_mask=attention_mask |
| 211 | ) |
| 212 | |
| 213 | if self.type_feature == "patch": |
| 214 | patch_feature = img_feature |
| 215 | |
| 216 | use_token_compression = self.image_token_compression is not None |
| 217 | use_padding = getattr(self, "img_processor_padding", None) is not None |
| 218 | if use_token_compression or use_padding: |
| 219 | # reshape to 2D tensor |
| 220 | width = int(math.sqrt(patch_feature.size(1))) |
| 221 | patch_feature = patch_feature.view( |
| 222 | -1, width, width, patch_feature.size(-1) |
| 223 | ) |
| 224 | # convert to NCHW |
| 225 | patch_feature = patch_feature.permute(0, 3, 1, 2) |
| 226 | |
| 227 | if use_padding: |
| 228 | patch_feature = self.img_processor_padding(patch_feature) |
| 229 | if use_token_compression: |
| 230 | patch_feature = self.image_token_compression(patch_feature) |
| 231 | |
| 232 | # convert to NHWC |
| 233 | patch_feature = patch_feature.permute(0, 2, 3, 1) |
| 234 | patch_feature = patch_feature.view( |
| 235 | -1, |
| 236 | patch_feature.size(1) * patch_feature.size(2), |
| 237 | patch_feature.size(-1), |
| 238 | ) |
| 239 | |
| 240 | return patch_feature |
| 241 | |
| 242 | raise NotImplementedError |
| 243 | |
| 244 | def forward( |
| 245 | self, |
| 246 | pixel_values: torch.FloatTensor, |
| 247 | image_sizes: torch.Tensor, |
| 248 | image_attention_mask: torch.Tensor, |
| 249 | ) -> list[torch.FloatTensor]: |
| 250 | """Process image and return vision embeddings. |
| 251 | |
| 252 | pixel_values: (num_images, num_crops, c, h, w) |
| 253 | image_sizes: [[h1, w1], [h2, w2]] |
| 254 | image_attention_mask: num_images x num_crops x 32 x 32 |
| 255 | output: (num_images, num_img_tokens, hidden_size) |
| 256 | """ |
| 257 | # eg |
| 258 | # pixel_values: torch.Size([1, 7, 3, 448, 448]) |
| 259 | # image_sizes: tensor([[ 896, 1344]], device='cuda:0') |
| 260 | # output: torch.Size([1, 1841, 3072]) |
| 261 | |
| 262 | if isinstance(self.img_projection, nn.Sequential): |
| 263 | target_device = self.img_projection[0].bias.device |
| 264 | target_dtype = self.img_projection[0].bias.dtype |
| 265 | else: # It's a single nn.Linear layer |
| 266 | target_device = self.img_projection.bias.device |
| 267 | target_dtype = self.img_projection.bias.dtype |
| 268 | |
| 269 | img_sizes = image_sizes |
| 270 | num_images, num_crops, c, h, w = pixel_values.shape |
| 271 | bs = num_images |
| 272 | pixel_values = pixel_values.flatten(0, 1) |
| 273 | |
| 274 | img_features = self.get_img_features( |
| 275 | pixel_values, image_attention_mask.bool().flatten(0, 1) |
| 276 | ) |
| 277 | |
| 278 | base_feat_height_target = self.base_feat_height_target |
| 279 | base_resolution = self.crop_size |
| 280 | base_feat_height_reduction = self.base_feat_height_reduction |
| 281 | |
| 282 | base_feat_height = base_feat_width = int(np.sqrt(img_features.shape[1])) |
| 283 | assert ( |
| 284 | base_feat_height == base_feat_height_target |
| 285 | and base_feat_width == base_feat_height_target |
| 286 | ), ( |
| 287 | f"base_feat_height: {base_feat_height}, " |
| 288 | f"base_feat_width: {base_feat_width}, " |
| 289 | f"expect {base_feat_height_target} features for hd transform" |
| 290 | ) |
| 291 | |
| 292 | # bs x max_num_crops x (24x24) x C |
| 293 | img_features = img_features.view( |
| 294 | bs, -1, base_feat_height * base_feat_width, self.image_dim_out |
| 295 | ) |
| 296 | C = self.image_dim_out |
| 297 | H = base_feat_height |
| 298 | |
| 299 | output_imgs = [] |
| 300 | output_len = [] |
| 301 | # training is tensor, inference is list |
| 302 | if isinstance(img_sizes, torch.Tensor): |
| 303 | img_sizes = img_sizes.view(-1, 2) |
| 304 | for _bs in range(bs): |
| 305 | h, w = img_sizes[_bs] |
| 306 | h = h // base_resolution |
| 307 | w = w // base_resolution |
| 308 | B_ = h * w |
| 309 | |
| 310 | # 1 x (24x24) x 1024 |
| 311 | global_img_feature = img_features[_bs, :1] |
| 312 | |
| 313 | # 1 x 12 x 12 x 4096 |
| 314 | glb_img = ( |
| 315 | global_img_feature.reshape(1, H, H, C) |
| 316 | .reshape( |
| 317 | 1, |
| 318 | H // base_feat_height_reduction, |
| 319 | base_feat_height_reduction, |
| 320 | H // base_feat_height_reduction, |
| 321 | base_feat_height_reduction, |
| 322 | C, |
| 323 | ) |
| 324 | .contiguous() |
| 325 | .permute(0, 1, 3, 2, 4, 5) |
| 326 | .reshape( |
| 327 | 1, |
| 328 | H // base_feat_height_reduction, |
| 329 | H // base_feat_height_reduction, |
| 330 | base_feat_height_reduction * base_feat_height_reduction * C, |
| 331 | ) |
| 332 | .contiguous() |
| 333 | ) |
| 334 | temp_glb_GN = self.sub_GN.repeat(1, H // base_feat_height_reduction, 1, 1) |
| 335 | |
| 336 | # 1 x 156 x 4096 |
| 337 | glb_img = torch.cat([glb_img, temp_glb_GN], dim=2).reshape( |
| 338 | 1, -1, base_feat_height_reduction * base_feat_height_reduction * C |
| 339 | ) |
| 340 | |
| 341 | # (max_num_crops-1) x (12x12) x C |
| 342 | sub_img = img_features[_bs, 1:] |
| 343 | # 16x574x1024 |
| 344 | # get rid of padding sub_img |
| 345 | sub_img = sub_img[:B_] |
| 346 | |
| 347 | # (num_crops, 12, 2, 12, 2, 1024) -> |
| 348 | # (num_crops, 12, 12, 2, 2, 1024) -> (num_crops, 12*12, 4*1024) |
| 349 | sub_img = ( |
| 350 | sub_img.reshape(B_, H, H, C) |
| 351 | .reshape( |
| 352 | B_, |
| 353 | H // base_feat_height_reduction, |
| 354 | base_feat_height_reduction, |
| 355 | H // base_feat_height_reduction, |
| 356 | base_feat_height_reduction, |
| 357 | C, |
| 358 | ) |
| 359 | .contiguous() |
| 360 | .permute(0, 1, 3, 2, 4, 5) |
| 361 | .reshape( |
| 362 | B_, -1, base_feat_height_reduction * base_feat_height_reduction * C |
| 363 | ) |
| 364 | .contiguous() |
| 365 | ) |
| 366 | sub_img = ( |
| 367 | sub_img.reshape( |
| 368 | 1, |
| 369 | h, |
| 370 | w, |
| 371 | base_feat_height // base_feat_height_reduction, |
| 372 | base_feat_width // base_feat_height_reduction, |
| 373 | -1, |
| 374 | ) |
| 375 | .permute(0, 1, 3, 2, 4, 5) |
| 376 | .reshape( |
| 377 | 1, |
| 378 | h * base_feat_height // base_feat_height_reduction, |
| 379 | w * base_feat_width // base_feat_height_reduction, |
| 380 | base_feat_height_reduction * base_feat_height_reduction * C, |
| 381 | ) |
| 382 | ) |
| 383 | |
| 384 | if image_attention_mask is not None and len(image_attention_mask) > 0: |
| 385 | reshaped_image_attention_mask = ( |
| 386 | image_attention_mask[_bs, 1 : B_ + 1, 0::2, 0::2] |
| 387 | .reshape( |
| 388 | 1, |
| 389 | h, |
| 390 | w, |
| 391 | base_feat_height // base_feat_height_reduction, |
| 392 | base_feat_width // base_feat_height_reduction, |
| 393 | ) |
| 394 | .permute(0, 1, 3, 2, 4) |
| 395 | .reshape( |
| 396 | 1, |
| 397 | h * base_feat_height // base_feat_height_reduction, |
| 398 | w * base_feat_width // base_feat_height_reduction, |
| 399 | ) |
| 400 | ) |
| 401 | # The mask stays on device for the encoder above, so these |
| 402 | # per-image counts have to come back to the host to drive the |
| 403 | # slicing and the Python-level length arithmetic. |
| 404 | with gpu_sync_allowed(): |
| 405 | useful_height = int( |
| 406 | reshaped_image_attention_mask[0, :, 0].sum().item() |
| 407 | ) |
| 408 | useful_width = int( |
| 409 | reshaped_image_attention_mask[0, 0, :].sum().item() |
| 410 | ) |
| 411 | mask_token_count = int( |
| 412 | image_attention_mask[_bs, : B_ + 1, 0::2, 0::2].sum().item() |
| 413 | ) |
| 414 | sub_img = sub_img[:, :useful_height, :useful_width] |
| 415 | temp_sub_GN = self.sub_GN.repeat(1, useful_height, 1, 1) |
| 416 | temp_len = ( |
| 417 | mask_token_count |
| 418 | + (useful_height + 1) |
| 419 | + base_feat_height // base_feat_height_reduction |
| 420 | ) |
| 421 | else: |
| 422 | temp_sub_GN = self.sub_GN.repeat( |
| 423 | 1, h * base_feat_height // base_feat_height_reduction, 1, 1 |
| 424 | ) |
| 425 | temp_len = int( |
| 426 | (h * w + 1) * self.num_img_tokens |
| 427 | + 1 |
| 428 | + (h + 1) * base_feat_height // base_feat_height_reduction |
| 429 | ) |
| 430 | |
| 431 | sub_img = torch.cat([sub_img, temp_sub_GN], dim=2).reshape( |
| 432 | 1, -1, base_feat_height_reduction * base_feat_height_reduction * C |
| 433 | ) |
| 434 | # (1, num_img_tokens, 1024*4) |
| 435 | |
| 436 | # glb + sub |
| 437 | if self.hd_transform_order == "glb_sub": |
| 438 | output_imgs.append(torch.cat([glb_img, self.glb_GN, sub_img], dim=1)) |
| 439 | elif self.hd_transform_order == "sub_glb": |
| 440 | output_imgs.append(torch.cat([sub_img, self.glb_GN, glb_img], dim=1)) |
| 441 | else: |
| 442 | raise NotImplementedError( |
| 443 | f"hd_transform_order = {self.hd_transform_order}, not implemented" |
| 444 | ) |
| 445 | |
| 446 | # temp_len = int((h*w+1)*144 + 1 + (h+1)*12) |
| 447 | assert temp_len == output_imgs[-1].shape[1], ( |
| 448 | f"temp_len: {temp_len}, output_imgs[-1].shape[1]: " |
| 449 | f"{output_imgs[-1].shape[1]}" |
| 450 | ) |
| 451 | |
| 452 | output_len.append(temp_len) |
| 453 | |
| 454 | img_set_tensor = [] |
| 455 | for _output_img in output_imgs: |
| 456 | img_feature_proj = self.img_projection( |
| 457 | _output_img.to(target_device).to(target_dtype) |
| 458 | ) |
| 459 | img_set_tensor.append(img_feature_proj.squeeze(0)) |
| 460 | |
| 461 | return img_set_tensor |
| 462 | |
| 463 | |
| 464 | class Phi4MMImagePixelInputs(TensorSchema): |
| 465 | """Dimensions: |
| 466 | - bn: Batch size * number of images |
| 467 | - p: Number of patches (1 + num_patches) |
| 468 | - c: Number of channels (3) |
| 469 | - h: Height of each image patch |
| 470 | - w: Width of each image patch |
| 471 | - nc: Number of crops |
| 472 | - H_mask: Height of attention mask |
| 473 | - W_mask: Width of attention mask |
| 474 | """ |
| 475 | |
| 476 | type: Literal["pixel_values"] |
| 477 | |
| 478 | pixel_values: Annotated[ |
| 479 | torch.Tensor | list[torch.Tensor], |
| 480 | TensorShape( |
| 481 | "bn", "p", 3, "h", "w", dynamic_dims={"p"} |
| 482 | ), # may be different per batch and image |
| 483 | ] |
| 484 | |
| 485 | image_sizes: Annotated[ |
| 486 | torch.Tensor, |
| 487 | TensorShape("bn", 2), # (height, width) |
| 488 | ] |
| 489 | |
| 490 | num_img_tokens: Annotated[ |
| 491 | list[int], |
| 492 | TensorShape("bn"), |
| 493 | ] |
| 494 | |
| 495 | image_attention_mask: Annotated[ |
| 496 | torch.Tensor | list[torch.Tensor], |
| 497 | TensorShape("bn", "nc", 32, 32, dynamic_dims={"nc"}), # H_mask, W_mask |
| 498 | ] |
| 499 | |
| 500 | |
| 501 | class Phi4MMAudioFeatureInputs(TensorSchema): |
| 502 | """Dimensions: |
| 503 | - bn: Batch size * number of audios |
| 504 | - t: Time frames (M) |
| 505 | """ |
| 506 | |
| 507 | type: Literal["audio_features"] |
| 508 | |
| 509 | audio_features: Annotated[ |
| 510 | torch.Tensor | list[torch.Tensor], |
| 511 | TensorShape("bn", "t", 80, dynamic_dims={"t"}), |
| 512 | ] |
| 513 | |
| 514 | |
| 515 | class Phi4MMAudioEmbeddingInputs(TensorSchema): |
| 516 | """Dimensions: |
| 517 | - b: Batch size |
| 518 | - n: Number of audios |
| 519 | - f: Audio feature size |
| 520 | - h: Hidden size (must match language model backbone) |
| 521 | """ |
| 522 | |
| 523 | type: Literal["audio_embeds"] |
| 524 | data: Annotated[ |
| 525 | NestedTensors, |
| 526 | TensorShape("b", "n", "f", "h"), |
| 527 | ] |
| 528 | |
| 529 | |
| 530 | Phi4MMAudioInputs: TypeAlias = Phi4MMAudioFeatureInputs | Phi4MMAudioEmbeddingInputs |
| 531 | |
| 532 | |
| 533 | class Phi4MMParsedModalities(typing.TypedDict, total=False): |
| 534 | images: Phi4MMImagePixelInputs |
| 535 | audios: Phi4MMAudioInputs |
| 536 | |
| 537 | |
| 538 | def stack_with_pad( |
| 539 | tensors: torch.Tensor | list[torch.Tensor], |
| 540 | padding_value: int | float = 0, |
| 541 | ) -> torch.Tensor: |
| 542 | """Stack tensors, padding dimensions that differ across items.""" |
| 543 | if isinstance(tensors, torch.Tensor): |
| 544 | return tensors |
| 545 | |
| 546 | assert len(tensors) > 0, "Cannot stack an empty tensor list" |
| 547 | first_shape = tensors[0].shape |
| 548 | if all(t.shape == first_shape for t in tensors): |
| 549 | return torch.stack(tensors) |
| 550 | |
| 551 | ndim = tensors[0].dim() |
| 552 | assert all(t.dim() == ndim for t in tensors[1:]), ( |
| 553 | "All tensors must have the same number of dimensions" |
| 554 | ) |
| 555 | |
| 556 | out_size = [ |
| 557 | len(tensors), |
| 558 | *(max(t.shape[i] for t in tensors) for i in range(ndim)), |
| 559 | ] |
| 560 | output = tensors[0].new_full(out_size, padding_value) |
| 561 | |
| 562 | for idx, tensor in enumerate(tensors): |
| 563 | slices = [idx, *(slice(0, size) for size in tensor.shape)] |
| 564 | output[tuple(slices)] = tensor |
| 565 | |
| 566 | return output |
| 567 | |
| 568 | |
| 569 | class Phi4MMProcessingInfo(BaseProcessingInfo): |
| 570 | @property |
| 571 | def image_tokens(self) -> list[str]: |
| 572 | return [f"<|image_{i + 1}|>" for i in range(100)] |
| 573 | |
| 574 | @property |
| 575 | def audio_tokens(self) -> list[str]: |
| 576 | return [f"<|audio_{i + 1}|>" for i in range(100)] |
| 577 | |
| 578 | def get_dynamic_hd( |
| 579 | self, |
| 580 | processor: ProcessorMixin, |
| 581 | ) -> int: |
| 582 | image_processor = processor.image_processor |
| 583 | return image_processor.dynamic_hd |
| 584 | |
| 585 | def get_feature_extractor(self, **kwargs: object) -> SequenceFeatureExtractor: |
| 586 | return self.get_hf_processor(**kwargs).audio_processor |
| 587 | |
| 588 | def get_data_parser(self): |
| 589 | feature_extractor = self.get_feature_extractor() |
| 590 | |
| 591 | return MultiModalDataParser( |
| 592 | target_sr=feature_extractor.sampling_rate, |
| 593 | audio_resample_method="scipy", |
| 594 | expected_hidden_size=self._get_expected_hidden_size(), |
| 595 | ) |
| 596 | |
| 597 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 598 | return {"audio": None, "image": None} |
| 599 | |
| 600 | def _find_target_aspect_ratio( |
| 601 | self, |
| 602 | orig_width: int, |
| 603 | orig_height: int, |
| 604 | image_size: int, |
| 605 | max_num: int, |
| 606 | min_num: int, |
| 607 | ): |
| 608 | w_crop_num = math.ceil(orig_width / float(image_size)) |
| 609 | h_crop_num = math.ceil(orig_height / float(image_size)) |
| 610 | if w_crop_num * h_crop_num > max_num: |
| 611 | aspect_ratio = orig_width / orig_height |
| 612 | |
| 613 | # calculate the existing image aspect ratio |
| 614 | target_ratio_candidates = set( |
| 615 | (i, j) |
| 616 | for i in range(1, max_num + 1) |
| 617 | for j in range(1, max_num + 1) |
| 618 | if i * j <= max_num and i * j >= min_num |
| 619 | ) |
| 620 | target_ratios = sorted(target_ratio_candidates, key=lambda x: x[0] * x[1]) |
| 621 | |
| 622 | # find the closest aspect ratio to the target |
| 623 | image_processor = self.get_hf_processor().image_processor |
| 624 | target_aspect_ratio = image_processor.find_closest_aspect_ratio( |
| 625 | aspect_ratio, |
| 626 | target_ratios, |
| 627 | orig_width, |
| 628 | orig_height, |
| 629 | image_size, |
| 630 | ) |
| 631 | |
| 632 | # calculate the target width and height |
| 633 | target_width = image_size * target_aspect_ratio[0] |
| 634 | target_height = image_size * target_aspect_ratio[1] |
| 635 | else: |
| 636 | target_width = image_size * w_crop_num |
| 637 | target_height = image_size * h_crop_num |
| 638 | target_aspect_ratio = (w_crop_num, h_crop_num) |
| 639 | return target_aspect_ratio, target_height, target_width |
| 640 | |
| 641 | def _compute_num_image_tokens( |
| 642 | self, |
| 643 | orig_width: int, |
| 644 | orig_height: int, |
| 645 | dynamic_hd_size: int, |
| 646 | vit_image_size: int, |
| 647 | vit_patch_size: int, |
| 648 | token_compression_factor: int = 2, |
| 649 | ): |
| 650 | """Compute the number of tokens an image is expected to take up considering |
| 651 | the image encoder architecture and exclude output features containing |
| 652 | only padding pixels |
| 653 | |
| 654 | for siglip, vit_image_size=448, vit_patch_size=14, so output will be |
| 655 | 32x32 feature map |
| 656 | NOTE right now, Phi4MM uses hard-coded token_compression_factor=2 |
| 657 | """ |
| 658 | assert vit_image_size % vit_patch_size == 0, ( |
| 659 | "vit_image_size must be divisible by vit_patch_size" |
| 660 | ) |
| 661 | assert vit_image_size // vit_patch_size % token_compression_factor == 0, ( |
| 662 | "vit_image_size // vit_patch_size must be divisible by " |
| 663 | "token_compression_factor" |
| 664 | ) |
| 665 | |
| 666 | target_aspect_ratio, target_height, target_width = ( |
| 667 | self._find_target_aspect_ratio( |
| 668 | orig_width, orig_height, vit_image_size, dynamic_hd_size, min_num=1 |
| 669 | ) |
| 670 | ) |
| 671 | assert target_aspect_ratio[0] * vit_image_size == target_width, ( |
| 672 | f"{target_aspect_ratio[0]} * {vit_image_size} != {target_width}" |
| 673 | ) |
| 674 | assert target_aspect_ratio[1] * vit_image_size == target_height, ( |
| 675 | f"{target_aspect_ratio[1]} * {vit_image_size} != {target_height}" |
| 676 | ) |
| 677 | assert ( |
| 678 | target_height % vit_image_size == 0 and target_width % vit_image_size == 0 |
| 679 | ) |
| 680 | |
| 681 | padding_height, padding_width = _get_padding_size( |
| 682 | orig_width, orig_height, target_height, target_width |
| 683 | ) |
| 684 | assert padding_width == 0 or padding_height == 0, ( |
| 685 | "padding_width or padding_height must be 0" |
| 686 | ) |
| 687 | |
| 688 | target_feat_width = target_width // vit_patch_size |
| 689 | target_feat_height = target_height // vit_patch_size |
| 690 | if padding_width >= vit_patch_size: |
| 691 | assert padding_height == 0, "padding_height not 0" |
| 692 | non_pad_feat_width = target_feat_width - math.floor( |
| 693 | padding_width / vit_patch_size |
| 694 | ) |
| 695 | non_pad_feat_height = target_feat_height |
| 696 | elif padding_height >= vit_patch_size: |
| 697 | assert padding_width == 0, "padding_width not 0" |
| 698 | non_pad_feat_height = target_feat_height - math.floor( |
| 699 | padding_height / vit_patch_size |
| 700 | ) |
| 701 | non_pad_feat_width = target_feat_width |
| 702 | else: |
| 703 | # small padding shorter than a vit patch |
| 704 | non_pad_feat_width = target_feat_width |
| 705 | non_pad_feat_height = target_feat_height |
| 706 | |
| 707 | feat_width = non_pad_feat_width // token_compression_factor |
| 708 | feat_height = non_pad_feat_height // token_compression_factor |
| 709 | # NOTE it's possible that the non-padding feature is not divisible |
| 710 | if non_pad_feat_width % token_compression_factor != 0: |
| 711 | feat_width += 1 |
| 712 | if non_pad_feat_height % token_compression_factor != 0: |
| 713 | feat_height += 1 |
| 714 | num_hd_patch_tokens = feat_width * feat_height |
| 715 | num_hd_newline_tokens = feat_height |
| 716 | vit_feature_size = vit_image_size // vit_patch_size |
| 717 | num_global_image_tokens = (vit_feature_size // token_compression_factor) ** 2 |
| 718 | num_sep_tokens = 1 |
| 719 | num_global_image_newline_tokens = vit_feature_size // token_compression_factor |
| 720 | |
| 721 | return ( |
| 722 | num_global_image_tokens |
| 723 | + num_sep_tokens |
| 724 | + num_hd_patch_tokens |
| 725 | + num_hd_newline_tokens |
| 726 | + num_global_image_newline_tokens |
| 727 | ) |
| 728 | |
| 729 | def get_num_image_tokens( |
| 730 | self, |
| 731 | *, |
| 732 | image_width: int, |
| 733 | image_height: int, |
| 734 | processor: ProcessorMixin, |
| 735 | ) -> int: |
| 736 | hf_config = self.get_hf_config() |
| 737 | vision_encoder_name = hf_config.img_processor |
| 738 | if vision_encoder_name is None: |
| 739 | vision_encoder_name = SIGLIP_NAME |
| 740 | prepro_config = VISION_ENCODER_TO_PROCESSING_CONFIG[vision_encoder_name] |
| 741 | vit_image_size = prepro_config["vit_image_size"] |
| 742 | vit_patch_size = prepro_config["vit_patch_size"] |
| 743 | token_compression_factor = prepro_config["token_compression_factor"] |
| 744 | |
| 745 | dynamic_hd_size = self.get_dynamic_hd(processor=processor) |
| 746 | |
| 747 | image_num_tokens = self._compute_num_image_tokens( |
| 748 | image_width, |
| 749 | image_height, |
| 750 | dynamic_hd_size=dynamic_hd_size, |
| 751 | vit_image_size=vit_image_size, |
| 752 | vit_patch_size=vit_patch_size, |
| 753 | token_compression_factor=token_compression_factor, |
| 754 | ) |
| 755 | |
| 756 | return image_num_tokens |
| 757 | |
| 758 | def get_image_size_with_most_features(self) -> ImageSize: |
| 759 | processor = self.get_hf_processor() |
| 760 | |
| 761 | hf_config = self.get_hf_config() |
| 762 | vision_encoder_name = hf_config.img_processor |
| 763 | if vision_encoder_name is None: |
| 764 | vision_encoder_name = SIGLIP_NAME |
| 765 | prepro_config = VISION_ENCODER_TO_PROCESSING_CONFIG[vision_encoder_name] |
| 766 | vit_image_size = prepro_config["vit_image_size"] |
| 767 | |
| 768 | max_side = vit_image_size * self.get_dynamic_hd(processor=processor) |
| 769 | return ImageSize(height=max_side, width=vit_image_size) |
| 770 | |
| 771 | def get_audio_num_frames(self, audio_len: int, sr: float) -> int: |
| 772 | """Compute the output size of the `extract_features` method. |
| 773 | |
| 774 | Args: |
| 775 | audio_len (int): Length of the input waveform in samples. |
| 776 | sr (float): Sampling rate of the waveform, either 16000 or 8000. |
| 777 | |
| 778 | Returns: |
| 779 | tuple (int, int): Output size as (T, D), where: |
| 780 | T: Number of time frames. |
| 781 | D: Number of Mel filterbank bins (80). |
| 782 | |
| 783 | """ |
| 784 | # Resample to 16000 or 8000 if needed |
| 785 | if sr > 16000: |
| 786 | resample_ratio = int(sr // 16000) |
| 787 | audio_len //= resample_ratio |
| 788 | elif 8000 <= sr < 16000: |
| 789 | # We'll resample to 16K from 8K |
| 790 | audio_len *= 2 |
| 791 | elif sr < 8000: |
| 792 | raise RuntimeError(f"Unsupported sample rate {sr}") |
| 793 | |
| 794 | # Spectrogram parameters for 16 kHz |
| 795 | win_length = 400 # Frame length in samples |
| 796 | hop_length = 160 # Frame shift in samples |
| 797 | |
| 798 | # Calculate number of frames (T) |
| 799 | num_frames = (audio_len - win_length) // hop_length + 1 |
| 800 | if num_frames < 1: |
| 801 | raise ValueError("Waveform too short for given parameters.") |
| 802 | |
| 803 | # Return time frames (T) |
| 804 | return num_frames |
| 805 | |
| 806 | def _compute_audio_embed_size(self, audio_frames: int) -> int: |
| 807 | """Compute the audio embedding size based on the audio frames and |
| 808 | compression rate. |
| 809 | """ |
| 810 | hf_config = self.get_hf_config() |
| 811 | compression_rate = hf_config.embd_layer["audio_embd_layer"]["compression_rate"] |
| 812 | # NOTE: this is a hard-coded value but might be configurable |
| 813 | # in the future |
| 814 | qformer_compression_rate = 1 |
| 815 | integer = audio_frames // compression_rate |
| 816 | remainder = audio_frames % compression_rate |
| 817 | |
| 818 | result = integer if remainder == 0 else integer + 1 |
| 819 | |
| 820 | integer = result // qformer_compression_rate |
| 821 | remainder = result % qformer_compression_rate |
| 822 | # qformer compression |
| 823 | result = integer if remainder == 0 else integer + 1 |
| 824 | |
| 825 | return result |
| 826 | |
| 827 | |
| 828 | class Phi4MMDummyInputsBuilder(BaseDummyInputsBuilder[Phi4MMProcessingInfo]): |
| 829 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 830 | num_audios = mm_counts.get("audio", 0) |
| 831 | num_images = mm_counts.get("image", 0) |
| 832 | |
| 833 | image_tokens: list[str] = self.info.image_tokens[:num_images] |
| 834 | audio_tokens: list[str] = self.info.audio_tokens[:num_audios] |
| 835 | |
| 836 | return "".join(image_tokens + audio_tokens) |
| 837 | |
| 838 | def get_dummy_mm_data( |
| 839 | self, |
| 840 | seq_len: int, |
| 841 | mm_counts: Mapping[str, int], |
| 842 | mm_options: MultiModalDummyOptions, |
| 843 | ) -> MultiModalDataDict: |
| 844 | target_width, target_height = self.info.get_image_size_with_most_features() |
| 845 | |
| 846 | mm_data = { |
| 847 | "image": self._get_dummy_images( |
| 848 | width=target_width, |
| 849 | height=target_height, |
| 850 | num_images=mm_counts.get("image", 0), |
| 851 | overrides=mm_options.get("image"), |
| 852 | ), |
| 853 | "audio": self._get_dummy_audios( |
| 854 | length=_AUDIO_MAX_SOUNDFILE_SIZE, |
| 855 | num_audios=mm_counts.get("audio", 0), |
| 856 | overrides=mm_options.get("audio"), |
| 857 | ), |
| 858 | } |
| 859 | |
| 860 | return mm_data |
| 861 | |
| 862 | |
| 863 | class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]): |
| 864 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 865 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 866 | |
| 867 | def _get_hf_mm_inputs( |
| 868 | self, |
| 869 | mm_items: MultiModalDataItems, |
| 870 | hf_kwargs: Mapping[str, object], |
| 871 | ) -> HFMultiModalInputs: |
| 872 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 873 | |
| 874 | # The Phi4MM processor expects "audios" instead of "audio" |
| 875 | hf_data = hf_inputs.hf_data |
| 876 | if "audio" in hf_data: |
| 877 | hf_data["audios"] = hf_data.pop("audio") |
| 878 | |
| 879 | sr = self.info.get_feature_extractor(**hf_inputs.hf_kwargs).sampling_rate |
| 880 | if audio_data := hf_data.get("audios", []): |
| 881 | assert isinstance(audio_data, Sequence) |
| 882 | hf_data["audios"] = [(data, sr) for data in audio_data] |
| 883 | |
| 884 | return hf_inputs |
| 885 | |
| 886 | def _postprocess_hf_mm_data( |
| 887 | self, |
| 888 | hf_data: Mapping[str, object], |
| 889 | hf_kwargs: Mapping[str, object], |
| 890 | processed_data: BatchFeature, |
| 891 | ) -> BatchFeature: |
| 892 | if not hf_data: |
| 893 | return processed_data |
| 894 | |
| 895 | hf_processor = self.info.get_hf_processor(**hf_kwargs) |
| 896 | num_img_tokens = [ |
| 897 | self.info.get_num_image_tokens( |
| 898 | image_width=img_size[0], |
| 899 | image_height=img_size[1], |
| 900 | processor=hf_processor, |
| 901 | ) |
| 902 | for img_size in processed_data["image_sizes"] |
| 903 | ] |
| 904 | processed_data["num_img_tokens"] = num_img_tokens |
| 905 | |
| 906 | audios = hf_data.get("audios", []) |
| 907 | assert isinstance(audios, list) |
| 908 | audio_features = processed_data.get("input_audio_embeds") |
| 909 | if audio_features is None: |
| 910 | return processed_data |
| 911 | |
| 912 | sr = self.info.get_feature_extractor(**hf_kwargs).sampling_rate |
| 913 | feature_sizes = [ |
| 914 | self.info.get_audio_num_frames(len(audio), sr) for audio, _ in audios |
| 915 | ] |
| 916 | processed_data["input_audio_embeds"] = [ |
| 917 | audio_features[idx, :size] for idx, size in enumerate(feature_sizes) |
| 918 | ] |
| 919 | |
| 920 | return processed_data |
| 921 | |
| 922 | def _get_mm_fields_config( |
| 923 | self, |
| 924 | hf_inputs: BatchFeature, |
| 925 | hf_kwargs: Mapping[str, object], |
| 926 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 927 | return dict( |
| 928 | input_image_embeds=MultiModalFieldConfig.batched("image"), |
| 929 | image_attention_mask=MultiModalFieldConfig.batched("image"), |
| 930 | image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 931 | num_img_tokens=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 932 | input_audio_embeds=MultiModalFieldConfig.batched("audio"), |
| 933 | ) |
| 934 | |
| 935 | def _get_prompt_updates( |
| 936 | self, |
| 937 | mm_items: MultiModalDataItems, |
| 938 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 939 | out_mm_kwargs: MultiModalKwargsItems, |
| 940 | ) -> Sequence[PromptUpdate]: |
| 941 | tokenizer = self.info.get_tokenizer() |
| 942 | image_tokens: list[str] = self.info.image_tokens # type: ignore |
| 943 | audio_tokens: list[str] = self.info.audio_tokens # type: ignore |
| 944 | |
| 945 | def get_image_token_ids(item_idx: int) -> list[int]: |
| 946 | return cached_encode( |
| 947 | tokenizer, image_tokens[item_idx], add_special_tokens=False |
| 948 | ) |
| 949 | |
| 950 | def get_audio_token_ids(item_idx: int) -> list[int]: |
| 951 | return cached_encode( |
| 952 | tokenizer, audio_tokens[item_idx], add_special_tokens=False |
| 953 | ) |
| 954 | |
| 955 | feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) |
| 956 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 957 | |
| 958 | def get_image_replacement_phi4mm(item_idx: int): |
| 959 | images = mm_items.get_items( |
| 960 | "image", (ImageEmbeddingItems, ImageProcessorItems) |
| 961 | ) |
| 962 | |
| 963 | if isinstance(images, ImageEmbeddingItems): |
| 964 | num_image_tokens = images.get_feature_size(item_idx) |
| 965 | else: |
| 966 | assert isinstance(images, ImageProcessorItems) |
| 967 | image_size = images.get_image_size(item_idx) |
| 968 | num_image_tokens = self.info.get_num_image_tokens( |
| 969 | image_width=image_size.width, |
| 970 | image_height=image_size.height, |
| 971 | processor=hf_processor, |
| 972 | ) |
| 973 | |
| 974 | return [_IMAGE_PLACEHOLDER_TOKEN_ID] * num_image_tokens |
| 975 | |
| 976 | def get_audio_replacement_phi4mm(item_idx: int): |
| 977 | audios = mm_items.get_items("audio", AudioProcessorItems) |
| 978 | # TODO(Isotr0py): support embedding inputs |
| 979 | audio_len = audios.get_audio_length(item_idx) |
| 980 | audio_frames = self.info.get_audio_num_frames( |
| 981 | audio_len, feature_extractor.sampling_rate |
| 982 | ) |
| 983 | audio_embed_size = self.info._compute_audio_embed_size(audio_frames) |
| 984 | |
| 985 | return [_AUDIO_PLACEHOLDER_TOKEN_ID] * audio_embed_size |
| 986 | |
| 987 | return [ |
| 988 | PromptReplacement( |
| 989 | modality="image", |
| 990 | target=get_image_token_ids, |
| 991 | replacement=get_image_replacement_phi4mm, |
| 992 | ), |
| 993 | PromptReplacement( |
| 994 | modality="audio", |
| 995 | target=get_audio_token_ids, |
| 996 | replacement=get_audio_replacement_phi4mm, |
| 997 | ), |
| 998 | ] |
| 999 | |
| 1000 | def _recompute_cached_prompt_update( |
| 1001 | self, |
| 1002 | cached_update: ResolvedPromptUpdate, |
| 1003 | new_item_idx: int, |
| 1004 | ) -> ResolvedPromptUpdate: |
| 1005 | new_update = super()._recompute_cached_prompt_update( |
| 1006 | cached_update, |
| 1007 | new_item_idx, |
| 1008 | ) |
| 1009 | |
| 1010 | tokenizer = self.info.get_tokenizer() |
| 1011 | |
| 1012 | if cached_update.modality == "image": |
| 1013 | image_tokens: list[str] = self.info.image_tokens # type: ignore |
| 1014 | new_update = new_update.with_target( |
| 1015 | cached_encode( |
| 1016 | tokenizer, image_tokens[new_item_idx], add_special_tokens=False |
| 1017 | ) |
| 1018 | ) |
| 1019 | elif cached_update.modality == "audio": |
| 1020 | audio_tokens: list[str] = self.info.audio_tokens # type: ignore |
| 1021 | new_update = new_update.with_target( |
| 1022 | cached_encode( |
| 1023 | tokenizer, audio_tokens[new_item_idx], add_special_tokens=False |
| 1024 | ) |
| 1025 | ) |
| 1026 | |
| 1027 | return new_update |
| 1028 | |
| 1029 | |
| 1030 | @MULTIMODAL_REGISTRY.register_processor( |
| 1031 | Phi4MMMultiModalProcessor, |
| 1032 | info=Phi4MMProcessingInfo, |
| 1033 | dummy_inputs=Phi4MMDummyInputsBuilder, |
| 1034 | ) |
| 1035 | class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal): |
| 1036 | """Implements the Phi-4-multimodal-instruct model in vLLM.""" |
| 1037 | |
| 1038 | packed_modules_mapping = { |
| 1039 | "qkv_proj": [ |
| 1040 | "qkv_proj", |
| 1041 | ], |
| 1042 | "gate_up_proj": [ |
| 1043 | "gate_up_proj", |
| 1044 | ], |
| 1045 | } |
| 1046 | |
| 1047 | hf_to_vllm_mapper = WeightsMapper( |
| 1048 | orig_to_new_substr={ |
| 1049 | "base_layer.": "", |
| 1050 | "lora": None, |
| 1051 | }, |
| 1052 | orig_to_new_prefix={ |
| 1053 | "model.embed_tokens_extend.audio_embed.audio_projection.vision.": "embed_tokens_extend.audio_projection_for_vision.", # noqa: E501 |
| 1054 | "model.embed_tokens_extend.audio_embed.audio_projection.speech.": "embed_tokens_extend.audio_projection.", # noqa: E501 |
| 1055 | "model.embed_tokens_extend.audio_embed.": "embed_tokens_extend.", |
| 1056 | "model.embed_tokens_extend.image_embed.": "vision_encoder.", |
| 1057 | }, |
| 1058 | ) |
| 1059 | |
| 1060 | @classmethod |
| 1061 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 1062 | if modality.startswith("image"): |
| 1063 | return f"<|image_{i}|>" |
| 1064 | if modality.startswith("audio"): |
| 1065 | return f"<|audio_{i}|>" |
| 1066 | |
| 1067 | raise ValueError("Only image or audio modality is supported") |
| 1068 | |
| 1069 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1070 | super().__init__() |
| 1071 | config = vllm_config.model_config.hf_config |
| 1072 | multimodal_config = vllm_config.model_config.multimodal_config |
| 1073 | assert multimodal_config, "multimodal_config is required" |
| 1074 | quant_config = vllm_config.quant_config |
| 1075 | |
| 1076 | self.config = config |
| 1077 | self.multimodal_config = multimodal_config |
| 1078 | self.quant_config = quant_config |
| 1079 | |
| 1080 | # Tensor/Pipeline parallel not supported for now. |
| 1081 | assert get_pp_group().world_size == 1, "pipeline parallel is not supported" |
| 1082 | |
| 1083 | with self._mark_tower_model(vllm_config, {"image", "video"}): |
| 1084 | self.vision_encoder = Phi4MMImageEncoder( |
| 1085 | config, |
| 1086 | quant_config, |
| 1087 | prefix=maybe_prefix(prefix, "model.vision_embed_tokens"), |
| 1088 | model_dir=config._name_or_path, |
| 1089 | ) |
| 1090 | |
| 1091 | if isinstance(config.embd_layer["audio_embd_layer"], dict): |
| 1092 | embedding_config = { |
| 1093 | "embedding_cls": config.embd_layer["audio_embd_layer"]["embedding_cls"], |
| 1094 | **config.embd_layer["audio_embd_layer"], |
| 1095 | } |
| 1096 | else: |
| 1097 | embedding_config = { |
| 1098 | "embedding_cls": self.config.embd_layer["embedding_cls"] |
| 1099 | } |
| 1100 | |
| 1101 | with self._mark_tower_model(vllm_config, "audio"): |
| 1102 | self.embed_tokens_extend = AudioEmbedding(config, **embedding_config) |
| 1103 | |
| 1104 | with self._mark_language_model(vllm_config): |
| 1105 | self.model = LlamaModel( |
| 1106 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 1107 | ) |
| 1108 | |
| 1109 | self.lm_head = ParallelLMHead( |
| 1110 | config.vocab_size, |
| 1111 | config.hidden_size, |
| 1112 | quant_config=quant_config, |
| 1113 | prefix=maybe_prefix(prefix, "lm_head"), |
| 1114 | ) |
| 1115 | if config.tie_word_embeddings: |
| 1116 | self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) |
| 1117 | logit_scale = getattr(config, "logit_scale", 1.0) |
| 1118 | self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale) |
| 1119 | |
| 1120 | def _parse_and_validate_audio_input( |
| 1121 | self, **kwargs: object |
| 1122 | ) -> Phi4MMAudioInputs | None: |
| 1123 | """Parse and validate the audio input to the model. This handles both |
| 1124 | audio features and audio embeddings, but only the former is used for |
| 1125 | now. |
| 1126 | |
| 1127 | Args: |
| 1128 | kwargs (object): Keyword arguments. |
| 1129 | |
| 1130 | Returns: |
| 1131 | Optional[Phi4MMAudioInputs]: Parsed and validated audio inputs. |
| 1132 | |
| 1133 | """ |
| 1134 | audio_features = kwargs.pop("input_audio_embeds", None) |
| 1135 | audio_embeds = kwargs.pop("audio_embeds", None) |
| 1136 | |
| 1137 | if audio_features is None and audio_embeds is None: |
| 1138 | return None |
| 1139 | |
| 1140 | if audio_features is not None: |
| 1141 | return Phi4MMAudioFeatureInputs( |
| 1142 | type="audio_features", |
| 1143 | audio_features=audio_features, |
| 1144 | ) |
| 1145 | |
| 1146 | if audio_embeds is not None: |
| 1147 | return Phi4MMAudioEmbeddingInputs(type="audio_embeds", data=audio_embeds) |
| 1148 | |
| 1149 | raise AssertionError("This line should be unreachable.") |
| 1150 | |
| 1151 | def _process_audio_input( |
| 1152 | self, audio_input: Phi4MMAudioInputs, audio_projection_mode: str |
| 1153 | ) -> NestedTensors: |
| 1154 | """Create the audio embeddings from the audio input, where the audio input |
| 1155 | is pairs of audio features and audio embed lengths. The audio input is |
| 1156 | created by `input_mapper_for_phi4mm_audio`. |
| 1157 | |
| 1158 | Args: |
| 1159 | audio_input (Phi4MMAudioInputs): Audio input. |
| 1160 | audio_projection_mode (str): Projection mode to use, selecting |
| 1161 | which audio projector weights are applied. |
| 1162 | |
| 1163 | Returns: |
| 1164 | NestedTensors: Audio embeddings |
| 1165 | |
| 1166 | """ |
| 1167 | if audio_input["type"] == "audio_embeds": |
| 1168 | return audio_input["data"] |
| 1169 | |
| 1170 | audio_features = audio_input["audio_features"] |
| 1171 | # (e.g. multiple examples) and the second dim is the multi-audio dim |
| 1172 | # (e.g. multiple audios in the same example) |
| 1173 | |
| 1174 | dtype = next(self.embed_tokens_extend.parameters()).dtype |
| 1175 | audio_embeds = [ |
| 1176 | self.embed_tokens_extend( |
| 1177 | features.to(dtype), |
| 1178 | audio_projection_mode=audio_projection_mode, |
| 1179 | ) |
| 1180 | for features in audio_features |
| 1181 | ] |
| 1182 | return audio_embeds |
| 1183 | |
| 1184 | def _parse_and_validate_image_input( |
| 1185 | self, **kwargs: object |
| 1186 | ) -> Phi4MMImagePixelInputs | None: |
| 1187 | pixel_values = kwargs.get("input_image_embeds") |
| 1188 | if pixel_values is None: |
| 1189 | return None |
| 1190 | |
| 1191 | image_sizes = kwargs.get("image_sizes") |
| 1192 | image_attention_mask = kwargs.get("image_attention_mask") |
| 1193 | num_img_tokens = kwargs.get("num_img_tokens") |
| 1194 | assert ( |
| 1195 | image_sizes is not None |
| 1196 | and image_attention_mask is not None |
| 1197 | and num_img_tokens is not None |
| 1198 | ), "Missing image inputs" |
| 1199 | |
| 1200 | return Phi4MMImagePixelInputs( |
| 1201 | type="pixel_values", |
| 1202 | pixel_values=pixel_values, |
| 1203 | image_sizes=image_sizes, |
| 1204 | image_attention_mask=image_attention_mask, |
| 1205 | num_img_tokens=num_img_tokens, |
| 1206 | ) |
| 1207 | |
| 1208 | def _parse_and_validate_multimodal_inputs( |
| 1209 | self, **kwargs: object |
| 1210 | ) -> Phi4MMParsedModalities: |
| 1211 | modalities = Phi4MMParsedModalities() |
| 1212 | |
| 1213 | # Preserve the order of modalities if there are multiple of them |
| 1214 | # from the order of kwargs. |
| 1215 | for input_key in kwargs: |
| 1216 | if ( |
| 1217 | input_key in ("input_image_embeds", "image_embeds") |
| 1218 | and "images" not in modalities |
| 1219 | ): |
| 1220 | image_input = self._parse_and_validate_image_input(**kwargs) |
| 1221 | assert image_input is not None |
| 1222 | modalities["images"] = image_input |
| 1223 | if ( |
| 1224 | input_key in ("input_audio_embeds", "audio_embeds") |
| 1225 | and "audios" not in modalities |
| 1226 | ): |
| 1227 | audio_input = self._parse_and_validate_audio_input(**kwargs) |
| 1228 | assert audio_input is not None |
| 1229 | modalities["audios"] = audio_input |
| 1230 | |
| 1231 | return modalities |
| 1232 | |
| 1233 | def _process_image_input( |
| 1234 | self, image_input: Phi4MMImagePixelInputs |
| 1235 | ) -> list[torch.Tensor]: |
| 1236 | dtype = next(self.vision_encoder.parameters()).dtype |
| 1237 | pixel_values = stack_with_pad(image_input["pixel_values"]).to(dtype) |
| 1238 | image_sizes = image_input["image_sizes"] |
| 1239 | image_attention_mask = stack_with_pad(image_input["image_attention_mask"]) |
| 1240 | image_embeds = self.vision_encoder( |
| 1241 | pixel_values, image_sizes, image_attention_mask |
| 1242 | ) |
| 1243 | return image_embeds |
| 1244 | |
| 1245 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 1246 | modalities = self._parse_and_validate_multimodal_inputs(**kwargs) |
| 1247 | if not modalities: |
| 1248 | return [] |
| 1249 | |
| 1250 | # The result multimodal_embeddings is tuple of tensors, with each |
| 1251 | # tensor corresponding to a multimodal data item (image or video). |
| 1252 | multimodal_embeddings: tuple[torch.Tensor, ...] = () |
| 1253 | |
| 1254 | # NOTE: It is important to iterate over the keys in this dictionary |
| 1255 | # to preserve the order of the modalities. |
| 1256 | audio_projection_mode = "speech" |
| 1257 | for modality in modalities: |
| 1258 | # make sure process images first |
| 1259 | if modality == "images": |
| 1260 | audio_projection_mode = "vision" |
| 1261 | image_input = modalities["images"] |
| 1262 | image_embeddings = self._process_image_input(image_input) |
| 1263 | multimodal_embeddings += tuple(image_embeddings) |
| 1264 | if modality == "audios": |
| 1265 | audio_input = modalities["audios"] |
| 1266 | audio_embeddings = self._process_audio_input( |
| 1267 | audio_input, audio_projection_mode=audio_projection_mode |
| 1268 | ) |
| 1269 | multimodal_embeddings += tuple(audio_embeddings) |
| 1270 | |
| 1271 | return multimodal_embeddings |
| 1272 | |
| 1273 | def forward( |
| 1274 | self, |
| 1275 | input_ids: torch.Tensor | None, |
| 1276 | positions: torch.Tensor, |
| 1277 | intermediate_tensors: IntermediateTensors | None = None, |
| 1278 | inputs_embeds: torch.Tensor | None = None, |
| 1279 | **kwargs: object, |
| 1280 | ) -> torch.Tensor: |
| 1281 | if intermediate_tensors is not None: |
| 1282 | inputs_embeds = None |
| 1283 | |
| 1284 | hidden_states = self.model( |
| 1285 | input_ids, |
| 1286 | positions, |
| 1287 | intermediate_tensors, |
| 1288 | inputs_embeds=inputs_embeds, |
| 1289 | ) |
| 1290 | |
| 1291 | return hidden_states |
| 1292 | |
| 1293 | def compute_logits( |
| 1294 | self, |
| 1295 | hidden_states: torch.Tensor, |
| 1296 | ) -> torch.Tensor | None: |
| 1297 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 1298 | return logits |
| 1299 | |
| 1300 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1301 | loader = AutoWeightsLoader(self) |
| 1302 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 1303 | |
| 1304 | def get_mm_mapping(self) -> MultiModelKeys: |
| 1305 | """Get the module prefix in multimodal models.""" |
| 1306 | return MultiModelKeys.from_string_field( |
| 1307 | language_model="model.", |
| 1308 | connector=["audio_projection_for_vision", "audio_projection"], |
| 1309 | tower_model=["vision_encoder", "embed_tokens_extend"], |
| 1310 | ) |