| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | from abc import abstractmethod |
| 5 | from collections.abc import Iterable, Iterator, Mapping, Sequence |
| 6 | from functools import partial |
| 7 | from typing import Annotated, Any, Literal, TypeAlias, TypeVar |
| 8 | |
| 9 | import numpy as np |
| 10 | import torch |
| 11 | import torch.nn as nn |
| 12 | from einops import rearrange |
| 13 | from transformers import BaseImageProcessor, PreTrainedConfig |
| 14 | from transformers.activations import GELUActivation |
| 15 | from transformers.feature_extraction_utils import BatchFeature |
| 16 | from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling |
| 17 | from transformers.utils import torch_int |
| 18 | |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.config.multimodal import MultiModalDummyOptions |
| 21 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 22 | from vllm.inputs import ModalityData, MultiModalDataDict |
| 23 | from vllm.logger import init_logger |
| 24 | from vllm.model_executor.layers.attention import ( |
| 25 | MMEncoderAttention, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.conv import Conv2dLayer |
| 28 | from vllm.model_executor.layers.linear import ( |
| 29 | ColumnParallelLinear, |
| 30 | QKVParallelLinear, |
| 31 | RowParallelLinear, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 34 | from vllm.model_executor.layers.rotary_embedding.common import ( |
| 35 | ApplyRotaryEmb, |
| 36 | ) |
| 37 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 38 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 39 | from vllm.multimodal.inputs import ( |
| 40 | ImageItem, |
| 41 | MultiModalFeatureSpec, |
| 42 | MultiModalFieldConfig, |
| 43 | MultiModalKwargsItem, |
| 44 | MultiModalKwargsItems, |
| 45 | VideoItem, |
| 46 | ) |
| 47 | from vllm.multimodal.parse import ( |
| 48 | DictEmbeddingItems, |
| 49 | ImageSize, |
| 50 | ModalityDataItems, |
| 51 | MultiModalDataItems, |
| 52 | MultiModalDataParser, |
| 53 | ) |
| 54 | from vllm.multimodal.processing import ( |
| 55 | BaseDummyInputsBuilder, |
| 56 | BaseMultiModalProcessor, |
| 57 | BaseProcessingInfo, |
| 58 | PromptReplacement, |
| 59 | PromptUpdate, |
| 60 | ) |
| 61 | from vllm.sequence import IntermediateTensors |
| 62 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 63 | from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d |
| 64 | |
| 65 | from .interfaces import ( |
| 66 | MultiModalEmbeddings, |
| 67 | SupportsLoRA, |
| 68 | SupportsMRoPE, |
| 69 | SupportsMultiModal, |
| 70 | SupportsPP, |
| 71 | ) |
| 72 | from .siglip import SiglipMLP |
| 73 | from .utils import ( |
| 74 | AutoWeightsLoader, |
| 75 | WeightsMapper, |
| 76 | init_vllm_registered_model, |
| 77 | maybe_prefix, |
| 78 | ) |
| 79 | from .vision import is_vit_use_data_parallel |
| 80 | |
| 81 | logger = init_logger(__name__) |
| 82 | |
| 83 | |
| 84 | def smart_resize( |
| 85 | height: int, |
| 86 | width: int, |
| 87 | factor: int, |
| 88 | min_pixels: int, |
| 89 | max_pixels: int, |
| 90 | ): |
| 91 | if height < factor: |
| 92 | logger.warning( |
| 93 | "smart_resize: height=%s < factor=%s, reset height=factor", |
| 94 | height, |
| 95 | factor, |
| 96 | ) |
| 97 | width = round((width * factor) / height) |
| 98 | height = factor |
| 99 | |
| 100 | if width < factor: |
| 101 | logger.warning( |
| 102 | "smart_resize: width=%s < factor=%s, reset width=factor", |
| 103 | width, |
| 104 | factor, |
| 105 | ) |
| 106 | height = round((height * factor) / width) |
| 107 | width = factor |
| 108 | |
| 109 | if max(height, width) / min(height, width) > 200: |
| 110 | raise ValueError( |
| 111 | "absolute aspect ratio must be smaller than 200, got " |
| 112 | "{max(height, width) / min(height, width)}" |
| 113 | ) |
| 114 | h_bar = round(height / factor) * factor |
| 115 | w_bar = round(width / factor) * factor |
| 116 | if h_bar * w_bar > max_pixels: |
| 117 | beta = math.sqrt((height * width) / max_pixels) |
| 118 | h_bar = math.floor(height / beta / factor) * factor |
| 119 | w_bar = math.floor(width / beta / factor) * factor |
| 120 | elif h_bar * w_bar < min_pixels: |
| 121 | beta = math.sqrt(min_pixels / (height * width)) |
| 122 | h_bar = math.ceil(height * beta / factor) * factor |
| 123 | w_bar = math.ceil(width * beta / factor) * factor |
| 124 | return h_bar, w_bar |
| 125 | |
| 126 | |
| 127 | class KeyeImagePixelInputs(TensorSchema): |
| 128 | """Dimensions: |
| 129 | - bnp: Batch size * Number of patches |
| 130 | - c: Number of channels |
| 131 | - ps: Patch size |
| 132 | - ni: Number of images |
| 133 | - g: Grid dimensions (3 for t, h, w) |
| 134 | """ |
| 135 | |
| 136 | type: Literal["pixel_values"] |
| 137 | pixel_values: Annotated[ |
| 138 | torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"}) |
| 139 | ] |
| 140 | image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)] |
| 141 | |
| 142 | |
| 143 | class KeyeImageEmbeddingInputs(TensorSchema): |
| 144 | """Dimensions: |
| 145 | - nf: Number of image features |
| 146 | - hs: Hidden size (must match the hidden size of language model |
| 147 | backbone) |
| 148 | - ni: Number of images |
| 149 | - g: Grid dimensions (3 for t, h, w) |
| 150 | """ |
| 151 | |
| 152 | type: Literal["image_embeds"] |
| 153 | image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")] |
| 154 | image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)] |
| 155 | |
| 156 | |
| 157 | KeyeImageInputs: TypeAlias = KeyeImagePixelInputs | KeyeImageEmbeddingInputs |
| 158 | |
| 159 | |
| 160 | class KeyeVideoPixelInputs(TensorSchema): |
| 161 | """Dimensions: |
| 162 | - bnp: Batch size * Number of patches |
| 163 | - c: Number of channels |
| 164 | - ps: Patch size |
| 165 | - ni: Number of images |
| 166 | - g: Grid dimensions (3 for t, h, w) |
| 167 | """ |
| 168 | |
| 169 | type: Literal["pixel_values_videos"] |
| 170 | pixel_values_videos: Annotated[ |
| 171 | torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"}) |
| 172 | ] |
| 173 | video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)] |
| 174 | |
| 175 | |
| 176 | class KeyeVideoEmbeddingInputs(TensorSchema): |
| 177 | """Dimensions: |
| 178 | - nf: Number of video features |
| 179 | - hs: Hidden size (must match the hidden size of language model |
| 180 | backbone) |
| 181 | - nv: Number of videos |
| 182 | - g: Grid dimensions (3 for t, h, w) |
| 183 | """ |
| 184 | |
| 185 | type: Literal["video_embeds"] |
| 186 | video_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")] |
| 187 | video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)] |
| 188 | |
| 189 | |
| 190 | KeyeVideoInputs: TypeAlias = KeyeVideoPixelInputs | KeyeVideoEmbeddingInputs |
| 191 | |
| 192 | |
| 193 | class KeyeVisionEmbeddings(nn.Module): |
| 194 | def __init__(self, config: PreTrainedConfig): |
| 195 | super().__init__() |
| 196 | self.config = config |
| 197 | self.embed_dim = config.hidden_size |
| 198 | self.image_size = config.image_size |
| 199 | self.patch_size = config.patch_size |
| 200 | |
| 201 | self.patch_embedding = Conv2dLayer( |
| 202 | in_channels=config.num_channels, |
| 203 | out_channels=self.embed_dim, |
| 204 | kernel_size=self.patch_size, |
| 205 | stride=self.patch_size, |
| 206 | padding="valid", |
| 207 | ) |
| 208 | |
| 209 | self.num_patches = (self.image_size // self.patch_size) ** 2 |
| 210 | self.num_positions = self.num_patches |
| 211 | self.cache_position_embedding: dict[tuple[int, int], torch.Tensor] = {} |
| 212 | self.cache_position_count: dict[tuple[int, int], int] = {} |
| 213 | self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) |
| 214 | self.packing_position_embedding = nn.Embedding(32768, self.embed_dim) |
| 215 | |
| 216 | self.register_buffer( |
| 217 | "position_ids", |
| 218 | torch.arange(self.num_positions).expand((1, -1)), |
| 219 | persistent=False, |
| 220 | ) |
| 221 | |
| 222 | def interpolate_pos_encoding( |
| 223 | self, |
| 224 | embeddings: torch.Tensor, |
| 225 | height: int, |
| 226 | width: int, |
| 227 | is_after_patchify: bool = False, |
| 228 | ) -> torch.Tensor: |
| 229 | num_positions = self.position_embedding.weight.shape[0] |
| 230 | |
| 231 | patch_pos_embed = self.position_embedding.weight.unsqueeze(0) |
| 232 | |
| 233 | dim = embeddings.shape[-1] |
| 234 | |
| 235 | if is_after_patchify: |
| 236 | new_height = height |
| 237 | new_width = width |
| 238 | else: |
| 239 | new_height = height // self.patch_size |
| 240 | new_width = width // self.patch_size |
| 241 | |
| 242 | sqrt_num_positions = torch_int(num_positions**0.5) |
| 243 | patch_pos_embed = patch_pos_embed.reshape( |
| 244 | 1, sqrt_num_positions, sqrt_num_positions, dim |
| 245 | ) |
| 246 | patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) |
| 247 | |
| 248 | patch_pos_embed = nn.functional.interpolate( |
| 249 | patch_pos_embed, |
| 250 | size=(new_height, new_width), |
| 251 | mode="bilinear", |
| 252 | align_corners=False, |
| 253 | ) |
| 254 | |
| 255 | patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) |
| 256 | return patch_pos_embed |
| 257 | |
| 258 | def fetch_position_embedding_lfu_cache(self, embeddings, h, w, max_cache: int = 20): |
| 259 | grid = (h, w) |
| 260 | if grid in self.cache_position_embedding: |
| 261 | self.cache_position_count[grid] += 1 |
| 262 | return self.cache_position_embedding[grid] |
| 263 | |
| 264 | if len(self.cache_position_embedding) >= max_cache: |
| 265 | min_hit_grid = min( |
| 266 | self.cache_position_count, |
| 267 | key=self.cache_position_count.__getitem__, |
| 268 | ) |
| 269 | self.cache_position_count.pop(min_hit_grid) |
| 270 | self.cache_position_embedding.pop(min_hit_grid) |
| 271 | |
| 272 | position_embedding = self.interpolate_pos_encoding(embeddings, h, w, True) |
| 273 | self.cache_position_count[grid] = 1 |
| 274 | self.cache_position_embedding[grid] = position_embedding |
| 275 | return position_embedding |
| 276 | |
| 277 | def forward( |
| 278 | self, |
| 279 | pixel_values: torch.FloatTensor, |
| 280 | position_ids: torch.Tensor | None = None, |
| 281 | image_grid_thw: list[tuple[int, int, int]] | None = None, |
| 282 | interpolate_pos_encoding=False, |
| 283 | ) -> torch.Tensor: |
| 284 | if pixel_values.dim() == 4: |
| 285 | pixel_values = pixel_values.unsqueeze(0) |
| 286 | if pixel_values.dim() == 5: |
| 287 | if position_ids is None: |
| 288 | raise ValueError( |
| 289 | "position_ids cannot be None when pixel_values.dim() is 5." |
| 290 | ) |
| 291 | ( |
| 292 | batch_size, |
| 293 | sequence_len, |
| 294 | channel, |
| 295 | height, |
| 296 | width, |
| 297 | ) = pixel_values.shape |
| 298 | target_dtype = self.patch_embedding.weight.dtype |
| 299 | pixel_values = rearrange(pixel_values, "b l c h w -> (b l) c h w") |
| 300 | patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) |
| 301 | embeddings = patch_embeds.flatten(-2).squeeze(-1) |
| 302 | |
| 303 | if interpolate_pos_encoding and image_grid_thw is not None: |
| 304 | start = 0 |
| 305 | tmp_embeddings = list() |
| 306 | for image_grid in image_grid_thw: |
| 307 | t, h, w = image_grid |
| 308 | end = start + t * h * w |
| 309 | image_embeddings = embeddings[start:end, :] |
| 310 | position_embedding = ( |
| 311 | self.interpolate_pos_encoding(image_embeddings, h, w, True) |
| 312 | .squeeze(0) |
| 313 | .repeat(t, 1) |
| 314 | ) |
| 315 | image_embeddings = image_embeddings + position_embedding |
| 316 | tmp_embeddings.append(image_embeddings) |
| 317 | start = end |
| 318 | embeddings = torch.concat(tmp_embeddings, dim=0).unsqueeze(0) |
| 319 | else: |
| 320 | embeddings = embeddings + self.packing_position_embedding(position_ids) |
| 321 | return embeddings |
| 322 | else: |
| 323 | raise ValueError( |
| 324 | "Unsupported pixel_values dimension:" |
| 325 | f" {pixel_values.dim()}. Expected 4 or 5." |
| 326 | ) |
| 327 | |
| 328 | |
| 329 | def apply_rotary_pos_emb_flashatt( |
| 330 | q: torch.Tensor, |
| 331 | k: torch.Tensor, |
| 332 | cos: torch.Tensor, |
| 333 | sin: torch.Tensor, |
| 334 | apply_rotary_emb: ApplyRotaryEmb, |
| 335 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 336 | cos = cos.chunk(2, dim=-1)[0].contiguous() |
| 337 | sin = sin.chunk(2, dim=-1)[0].contiguous() |
| 338 | |
| 339 | q_embed = apply_rotary_emb(q, cos, sin) |
| 340 | k_embed = apply_rotary_emb(k, cos, sin) |
| 341 | |
| 342 | return q_embed, k_embed |
| 343 | |
| 344 | |
| 345 | class KeyeSiglipAttention(nn.Module): |
| 346 | """Multi-headed attention from 'Attention Is All You |
| 347 | Need' paper.""" |
| 348 | |
| 349 | def __init__( |
| 350 | self, |
| 351 | config: PreTrainedConfig, |
| 352 | quant_config: QuantizationConfig | None = None, |
| 353 | prefix: str = "", |
| 354 | ): |
| 355 | super().__init__() |
| 356 | self.config = config |
| 357 | |
| 358 | hidden_size = config.hidden_size |
| 359 | self.hidden_size = config.hidden_size |
| 360 | use_data_parallel = is_vit_use_data_parallel() |
| 361 | tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size() |
| 362 | self.total_num_heads = config.num_attention_heads |
| 363 | assert self.total_num_heads % tp_size == 0 |
| 364 | self.num_heads = self.total_num_heads // tp_size |
| 365 | self.total_num_kv_heads = config.num_attention_heads |
| 366 | if self.total_num_kv_heads >= tp_size: |
| 367 | assert self.total_num_kv_heads % tp_size == 0 |
| 368 | else: |
| 369 | assert tp_size % self.total_num_kv_heads == 0 |
| 370 | self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) |
| 371 | self.head_dim = config.hidden_size // self.total_num_heads |
| 372 | self.q_size = self.num_heads * self.head_dim |
| 373 | self.kv_size = self.num_kv_heads * self.head_dim |
| 374 | self.scale = self.head_dim**-0.5 |
| 375 | |
| 376 | self.qkv_proj = QKVParallelLinear( |
| 377 | hidden_size, |
| 378 | self.head_dim, |
| 379 | self.total_num_heads, |
| 380 | self.total_num_kv_heads, |
| 381 | bias=True, |
| 382 | quant_config=quant_config, |
| 383 | prefix=f"{prefix}.qkv_proj", |
| 384 | ) |
| 385 | self.out_proj = RowParallelLinear( |
| 386 | input_size=hidden_size, |
| 387 | output_size=hidden_size, |
| 388 | quant_config=quant_config, |
| 389 | prefix=f"{prefix}.out_proj", |
| 390 | ) |
| 391 | |
| 392 | self.attn = MMEncoderAttention( |
| 393 | num_heads=self.num_heads, |
| 394 | head_size=self.head_dim, |
| 395 | scale=self.scale, |
| 396 | num_kv_heads=self.num_kv_heads, |
| 397 | prefix=f"{prefix}.attn", |
| 398 | ) |
| 399 | |
| 400 | self.apply_rotary_emb = ApplyRotaryEmb( |
| 401 | enforce_enable=True, |
| 402 | enable_fp32_compute=True, |
| 403 | ) |
| 404 | |
| 405 | def forward( |
| 406 | self, |
| 407 | hidden_states: torch.Tensor, |
| 408 | attention_mask: torch.Tensor | None = None, |
| 409 | output_attentions: bool | None = False, |
| 410 | cu_seqlens: torch.Tensor | None = None, |
| 411 | rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None, |
| 412 | ) -> torch.Tensor: |
| 413 | qkv, _ = self.qkv_proj(hidden_states) |
| 414 | q, k, v = qkv.split( |
| 415 | [self.q_size, self.kv_size, self.kv_size], |
| 416 | dim=-1, |
| 417 | ) |
| 418 | |
| 419 | if cu_seqlens is None: |
| 420 | raise ValueError("cu_seqlens cannot be None.") |
| 421 | max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() |
| 422 | |
| 423 | if rope_emb is None: |
| 424 | q = q.view(*q.shape[:-1], self.num_heads, self.head_dim) |
| 425 | k = k.view( |
| 426 | *k.shape[:-1], |
| 427 | self.num_kv_heads, |
| 428 | self.head_dim, |
| 429 | ) |
| 430 | v = v.view( |
| 431 | *v.shape[:-1], |
| 432 | self.num_kv_heads, |
| 433 | self.head_dim, |
| 434 | ) |
| 435 | else: |
| 436 | cos, sin = rope_emb |
| 437 | q = q.view(*q.shape[:-1], self.num_heads, self.head_dim) |
| 438 | k = k.view( |
| 439 | *k.shape[:-1], |
| 440 | self.num_kv_heads, |
| 441 | self.head_dim, |
| 442 | ) |
| 443 | q, k = apply_rotary_pos_emb_flashatt(q, k, cos, sin, self.apply_rotary_emb) |
| 444 | v = v.view( |
| 445 | *v.shape[:-1], |
| 446 | self.num_kv_heads, |
| 447 | self.head_dim, |
| 448 | ) |
| 449 | |
| 450 | context_layer = self.attn( |
| 451 | query=q, |
| 452 | key=k, |
| 453 | value=v, |
| 454 | cu_seqlens=cu_seqlens, |
| 455 | max_seqlen=max_seqlen, |
| 456 | ) |
| 457 | context_layer = rearrange(context_layer, "b s h d -> b s (h d)") |
| 458 | |
| 459 | output, _ = self.out_proj(context_layer) |
| 460 | return output |
| 461 | |
| 462 | |
| 463 | class SigLIPRotaryEmbedding(nn.Module): |
| 464 | def __init__(self, dim: int, theta: float = 10000.0) -> None: |
| 465 | super().__init__() |
| 466 | self.dim = dim |
| 467 | self.theta = theta |
| 468 | self.rope_init() |
| 469 | |
| 470 | def rope_init(self): |
| 471 | inv_freq = 1.0 / ( |
| 472 | self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim) |
| 473 | ) |
| 474 | self.register_buffer("inv_freq", inv_freq, persistent=False) |
| 475 | |
| 476 | def forward(self, seqlen: int) -> torch.Tensor: |
| 477 | seq = torch.arange( |
| 478 | seqlen, |
| 479 | device=self.inv_freq.device, |
| 480 | dtype=self.inv_freq.dtype, |
| 481 | ) |
| 482 | freqs = torch.outer(seq, self.inv_freq) |
| 483 | return freqs |
| 484 | |
| 485 | |
| 486 | class KeyeSiglipEncoderLayer(nn.Module): |
| 487 | def __init__( |
| 488 | self, |
| 489 | config: PreTrainedConfig, |
| 490 | quant_config: QuantizationConfig | None = None, |
| 491 | prefix: str = "", |
| 492 | ): |
| 493 | super().__init__() |
| 494 | self.embed_dim = config.hidden_size |
| 495 | self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) |
| 496 | self.self_attn = KeyeSiglipAttention( |
| 497 | config, |
| 498 | quant_config=quant_config, |
| 499 | prefix=f"{prefix}.self_attn", |
| 500 | ) |
| 501 | self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) |
| 502 | self.mlp = SiglipMLP( |
| 503 | config, |
| 504 | quant_config=quant_config, |
| 505 | prefix=f"{prefix}.mlp", |
| 506 | ) |
| 507 | |
| 508 | def forward( |
| 509 | self, |
| 510 | hidden_states: torch.Tensor, |
| 511 | attention_mask: torch.Tensor, |
| 512 | output_attentions: bool | None = False, |
| 513 | cu_seqlens: torch.Tensor | None = None, |
| 514 | rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None, |
| 515 | ) -> tuple[torch.FloatTensor]: |
| 516 | residual = hidden_states |
| 517 | |
| 518 | hidden_states = self.layer_norm1(hidden_states) |
| 519 | hidden_states = self.self_attn( |
| 520 | hidden_states=hidden_states, |
| 521 | attention_mask=attention_mask, |
| 522 | output_attentions=output_attentions, |
| 523 | cu_seqlens=cu_seqlens, |
| 524 | rope_emb=rope_emb, |
| 525 | ) |
| 526 | |
| 527 | hidden_states = residual + hidden_states |
| 528 | |
| 529 | residual = hidden_states |
| 530 | hidden_states = self.layer_norm2(hidden_states) |
| 531 | hidden_states = self.mlp(hidden_states) |
| 532 | |
| 533 | hidden_states = residual + hidden_states |
| 534 | |
| 535 | return hidden_states |
| 536 | |
| 537 | |
| 538 | class KeyeSiglipEncoder(nn.Module): |
| 539 | def __init__( |
| 540 | self, |
| 541 | config: PreTrainedConfig, |
| 542 | quant_config: QuantizationConfig | None = None, |
| 543 | prefix: str = "", |
| 544 | ): |
| 545 | super().__init__() |
| 546 | self.config = config |
| 547 | embed_dim = config.hidden_size |
| 548 | num_heads = config.num_attention_heads |
| 549 | head_dim = embed_dim // num_heads |
| 550 | self.layers = nn.ModuleList( |
| 551 | [ |
| 552 | KeyeSiglipEncoderLayer( |
| 553 | config, |
| 554 | quant_config=quant_config, |
| 555 | prefix=f"{prefix}.layers.{layer_idx}", |
| 556 | ) |
| 557 | for layer_idx in range(config.num_hidden_layers) |
| 558 | ] |
| 559 | ) |
| 560 | self.rotary_pos_emb = SigLIPRotaryEmbedding(head_dim // 2) |
| 561 | |
| 562 | @staticmethod |
| 563 | def flatten_list(image_grid_thw): |
| 564 | tmp_image_grid_thw = list() |
| 565 | for image_grid in image_grid_thw: |
| 566 | if isinstance(image_grid, list): |
| 567 | tmp_image_grid_thw.extend(image_grid) |
| 568 | else: |
| 569 | tmp_image_grid_thw.append(image_grid) |
| 570 | return tmp_image_grid_thw |
| 571 | |
| 572 | def forward( |
| 573 | self, |
| 574 | inputs_embeds, |
| 575 | attention_mask: torch.Tensor | None = None, |
| 576 | output_attentions: bool | None = None, |
| 577 | output_hidden_states: bool | None = None, |
| 578 | cu_seqlens: torch.Tensor | None = None, |
| 579 | image_grid_thw: list[tuple[int, int, int]] | None = None, |
| 580 | height_position_ids: torch.Tensor | None = None, |
| 581 | width_position_ids: torch.Tensor | None = None, |
| 582 | use_rope: bool | None = False, |
| 583 | window_size: int = -1, |
| 584 | vision_or_text: str = "vision", |
| 585 | ) -> BaseModelOutput: |
| 586 | device = inputs_embeds.device |
| 587 | hidden_states = inputs_embeds |
| 588 | if use_rope is True: |
| 589 | flatten_image_grid_thw = self.flatten_list(image_grid_thw) |
| 590 | |
| 591 | if width_position_ids is None or height_position_ids is None: |
| 592 | split_hids = list() |
| 593 | split_wids = list() |
| 594 | for t, h, w in flatten_image_grid_thw: |
| 595 | image_pids = torch.arange(t * h * w, device=device) % (h * w) |
| 596 | sample_hids = image_pids // w |
| 597 | sample_wids = image_pids % w |
| 598 | split_hids.append(sample_hids) |
| 599 | split_wids.append(sample_wids) |
| 600 | width_position_ids = torch.concat(split_wids, dim=0) |
| 601 | height_position_ids = torch.concat(split_hids, dim=0) |
| 602 | |
| 603 | pids = torch.stack( |
| 604 | [height_position_ids, width_position_ids], |
| 605 | dim=-1, |
| 606 | ) |
| 607 | # The ids are built from the grids above, so `h`/`w` bound them |
| 608 | # and the table size is known on the host. |
| 609 | max_grid_size = max(max(h, w) for _, h, w in flatten_image_grid_thw) |
| 610 | rope_emb_max_grid = self.rotary_pos_emb(max_grid_size) |
| 611 | rope_emb = rope_emb_max_grid[pids].flatten(1) |
| 612 | rope_emb = rope_emb.repeat(1, 2) |
| 613 | rope_emb = (rope_emb.cos(), rope_emb.sin()) |
| 614 | else: |
| 615 | rope_emb = None |
| 616 | |
| 617 | attn_cu_seqlens = cu_seqlens |
| 618 | hidden_states = inputs_embeds |
| 619 | assert attention_mask is None |
| 620 | |
| 621 | for encoder_layer in self.layers: |
| 622 | hidden_states = encoder_layer( |
| 623 | hidden_states, |
| 624 | attention_mask, |
| 625 | output_attentions=output_attentions, |
| 626 | cu_seqlens=attn_cu_seqlens, |
| 627 | rope_emb=rope_emb, |
| 628 | ) |
| 629 | return hidden_states |
| 630 | |
| 631 | |
| 632 | class KeyeSiglipVisionTransformer(nn.Module): |
| 633 | def __init__( |
| 634 | self, |
| 635 | config: PreTrainedConfig, |
| 636 | quant_config: QuantizationConfig | None = None, |
| 637 | prefix: str = "", |
| 638 | ): |
| 639 | super().__init__() |
| 640 | self.config = config |
| 641 | embed_dim = config.hidden_size |
| 642 | |
| 643 | self.embeddings = KeyeVisionEmbeddings(config) |
| 644 | self.encoder = KeyeSiglipEncoder( |
| 645 | config, |
| 646 | quant_config=quant_config, |
| 647 | prefix=f"{prefix}.encoder", |
| 648 | ) |
| 649 | self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) |
| 650 | |
| 651 | def forward( |
| 652 | self, |
| 653 | pixel_values, |
| 654 | output_attentions: bool | None = None, |
| 655 | output_hidden_states: bool | None = None, |
| 656 | interpolate_pos_encoding: bool | None = False, |
| 657 | attention_mask: torch.Tensor | None = None, |
| 658 | sample_indices: torch.Tensor | None = None, |
| 659 | image_indices: torch.Tensor | None = None, |
| 660 | position_ids: torch.Tensor | None = None, |
| 661 | height_position_ids: torch.Tensor | None = None, |
| 662 | width_position_ids: torch.Tensor | None = None, |
| 663 | cu_seqlens: torch.Tensor | None = None, |
| 664 | padding_mask: torch.Tensor | None = None, |
| 665 | vision_return_embed_list: bool | None = False, |
| 666 | image_grid_thw: list[tuple[int, int, int]] | None = None, |
| 667 | return_pooler_output: bool | None = True, |
| 668 | use_rope: bool | None = False, |
| 669 | window_size: int = -1, |
| 670 | ) -> BaseModelOutputWithPooling: |
| 671 | hidden_states = self.embeddings( |
| 672 | pixel_values, |
| 673 | interpolate_pos_encoding=interpolate_pos_encoding, |
| 674 | position_ids=position_ids, |
| 675 | image_grid_thw=image_grid_thw, |
| 676 | ) |
| 677 | |
| 678 | last_hidden_state = self.encoder( |
| 679 | inputs_embeds=hidden_states, |
| 680 | output_attentions=output_attentions, |
| 681 | output_hidden_states=output_hidden_states, |
| 682 | attention_mask=attention_mask, |
| 683 | cu_seqlens=cu_seqlens, |
| 684 | image_grid_thw=image_grid_thw, |
| 685 | use_rope=use_rope, |
| 686 | height_position_ids=height_position_ids, |
| 687 | width_position_ids=width_position_ids, |
| 688 | window_size=window_size, |
| 689 | vision_or_text="vision", |
| 690 | ) |
| 691 | |
| 692 | last_hidden_state = self.post_layernorm(last_hidden_state) |
| 693 | |
| 694 | if cu_seqlens is None: |
| 695 | raise ValueError( |
| 696 | "cu_seqlens cannot be None for " |
| 697 | "SiglipVisionTransformer output processing." |
| 698 | ) |
| 699 | # `cu_seqlens` is the running sum of the per-image `t * h * w`, so the |
| 700 | # split sizes are known on the host and slicing needs no device read. |
| 701 | split_sizes = [ |
| 702 | int(np.prod(thw)) for thw in self.encoder.flatten_list(image_grid_thw) |
| 703 | ] |
| 704 | return [ |
| 705 | tensor.squeeze(0) for tensor in last_hidden_state.split(split_sizes, dim=1) |
| 706 | ] |
| 707 | |
| 708 | |
| 709 | class KeyeSiglipVisionModel(nn.Module): |
| 710 | config_class = PreTrainedConfig |
| 711 | main_input_name = "pixel_values" |
| 712 | |
| 713 | hf_to_vllm_mapper = WeightsMapper( |
| 714 | orig_to_new_stacked={ |
| 715 | ".q_proj": (".qkv_proj", "q"), |
| 716 | ".k_proj": (".qkv_proj", "k"), |
| 717 | ".v_proj": (".qkv_proj", "v"), |
| 718 | }, |
| 719 | orig_to_new_prefix={"vision_model.head.": None}, |
| 720 | ) |
| 721 | |
| 722 | def __init__( |
| 723 | self, |
| 724 | config: PreTrainedConfig, |
| 725 | quant_config: QuantizationConfig | None = None, |
| 726 | prefix: str = "", |
| 727 | ): |
| 728 | super().__init__() |
| 729 | |
| 730 | self.vision_model = KeyeSiglipVisionTransformer( |
| 731 | config, |
| 732 | quant_config=quant_config, |
| 733 | prefix=maybe_prefix(prefix, "vision_model"), |
| 734 | ) |
| 735 | self.quant_config = quant_config |
| 736 | |
| 737 | @property |
| 738 | def dtype(self) -> torch.dtype: |
| 739 | return self.vision_model.embeddings.patch_embedding.weight.dtype |
| 740 | |
| 741 | @property |
| 742 | def device(self) -> torch.device: |
| 743 | return self.vision_model.embeddings.patch_embedding.weight.device |
| 744 | |
| 745 | def get_input_embeddings(self) -> nn.Module: |
| 746 | return self.vision_model.embeddings.patch_embedding |
| 747 | |
| 748 | def forward( |
| 749 | self, |
| 750 | pixel_values, |
| 751 | sample_indices: torch.Tensor | None = None, |
| 752 | output_attentions: bool | None = None, |
| 753 | output_hidden_states: bool | None = None, |
| 754 | interpolate_pos_encoding: bool = False, |
| 755 | position_ids: torch.Tensor | None = None, |
| 756 | vision_return_embed_list: bool | None = False, |
| 757 | image_grid_thw: list[tuple[int, int, int]] | None = None, |
| 758 | cu_seqlens: torch.Tensor | None = None, |
| 759 | return_pooler_output: bool | None = True, |
| 760 | use_rope: bool | None = False, |
| 761 | window_size: int = -1, |
| 762 | ) -> BaseModelOutputWithPooling: |
| 763 | return self.vision_model( |
| 764 | pixel_values=pixel_values, |
| 765 | output_attentions=output_attentions, |
| 766 | output_hidden_states=output_hidden_states, |
| 767 | interpolate_pos_encoding=interpolate_pos_encoding, |
| 768 | position_ids=position_ids, |
| 769 | vision_return_embed_list=vision_return_embed_list, |
| 770 | image_grid_thw=image_grid_thw, |
| 771 | sample_indices=sample_indices, |
| 772 | cu_seqlens=cu_seqlens, |
| 773 | return_pooler_output=return_pooler_output, |
| 774 | use_rope=use_rope, |
| 775 | window_size=window_size, |
| 776 | ) |
| 777 | |
| 778 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 779 | loader = AutoWeightsLoader(self) |
| 780 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 781 | |
| 782 | |
| 783 | class Projector(nn.Module): |
| 784 | def __init__( |
| 785 | self, |
| 786 | text_config: PreTrainedConfig, |
| 787 | vision_config: PreTrainedConfig, |
| 788 | quant_config: QuantizationConfig | None = None, |
| 789 | prefix: str = "", |
| 790 | ): |
| 791 | super().__init__() |
| 792 | self.text_config = text_config |
| 793 | self.vision_config = vision_config |
| 794 | self.merge_kernel_size = (2, 2) |
| 795 | |
| 796 | self.hidden_size = ( |
| 797 | self.vision_config.hidden_size |
| 798 | * self.merge_kernel_size[0] |
| 799 | * self.merge_kernel_size[1] |
| 800 | ) |
| 801 | |
| 802 | self.pre_norm = torch.nn.LayerNorm(self.vision_config.hidden_size, eps=1e-05) |
| 803 | self.act = GELUActivation() |
| 804 | |
| 805 | self.linear_1 = ColumnParallelLinear( |
| 806 | self.hidden_size, |
| 807 | self.hidden_size, |
| 808 | bias=True, |
| 809 | quant_config=quant_config, |
| 810 | prefix=f"{prefix}.linear_1", |
| 811 | ) |
| 812 | self.linear_2 = RowParallelLinear( |
| 813 | self.hidden_size, |
| 814 | self.text_config.hidden_size, |
| 815 | bias=True, |
| 816 | quant_config=quant_config, |
| 817 | prefix=f"{prefix}.linear_2", |
| 818 | ) |
| 819 | |
| 820 | def forward( |
| 821 | self, |
| 822 | image_features: torch.Tensor | list[torch.Tensor], |
| 823 | image_grid_thw: list[tuple[int, int, int]], |
| 824 | ) -> torch.Tensor | list[torch.Tensor]: |
| 825 | m1, m2 = self.merge_kernel_size |
| 826 | if isinstance(image_features, (list, tuple)): |
| 827 | processed_features = list() |
| 828 | for image_feature, image_grid in zip(image_features, image_grid_thw): |
| 829 | image_feature = self.pre_norm(image_feature) |
| 830 | t, h, w = image_grid |
| 831 | |
| 832 | image_feature = rearrange( |
| 833 | image_feature, |
| 834 | "(t h p1 w p2) d -> (t h w) (p1 p2 d)", |
| 835 | t=t, |
| 836 | h=h // m1, |
| 837 | p1=m1, |
| 838 | w=w // m2, |
| 839 | p2=m2, |
| 840 | ) |
| 841 | hidden_states, _ = self.linear_1(image_feature) |
| 842 | hidden_states = self.act(hidden_states) |
| 843 | hidden_states, _ = self.linear_2(hidden_states) |
| 844 | processed_features.append(hidden_states) |
| 845 | |
| 846 | return processed_features |
| 847 | |
| 848 | dims = image_features.shape[:-1] |
| 849 | dim = image_features.shape[-1] |
| 850 | image_features = image_features.view(np.prod(dims), dim) |
| 851 | hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size) |
| 852 | hidden_states = self.linear_1(hidden_states) |
| 853 | hidden_states = self.act(hidden_states) |
| 854 | hidden_states = self.linear_2(hidden_states) |
| 855 | |
| 856 | return hidden_states.view(*dims, -1) |
| 857 | |
| 858 | |
| 859 | def _keye_field_config( |
| 860 | hf_inputs: Mapping[str, torch.Tensor], |
| 861 | ): |
| 862 | image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) |
| 863 | image_grid_sizes = image_grid_thw.prod(-1) |
| 864 | |
| 865 | video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3))) |
| 866 | video_grid_sizes = video_grid_thw.prod(-1) |
| 867 | |
| 868 | return dict( |
| 869 | pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), |
| 870 | image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), |
| 871 | image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 872 | pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( |
| 873 | "video", video_grid_sizes |
| 874 | ), |
| 875 | video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes), |
| 876 | video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True), |
| 877 | ) |
| 878 | |
| 879 | |
| 880 | class KeyeMultiModalDataParser(MultiModalDataParser): |
| 881 | # The patch grid is what sizes the placeholder range. |
| 882 | embedding_fields = { |
| 883 | "image": {"image_embeds": "values", "image_grid_thw": "metadata"}, |
| 884 | "video": {"video_embeds": "values", "video_grid_thw": "metadata"}, |
| 885 | } |
| 886 | |
| 887 | def _parse_image_data( |
| 888 | self, |
| 889 | data: dict[str, torch.Tensor] | ModalityData[ImageItem], |
| 890 | ) -> ModalityDataItems[Any, Any] | None: |
| 891 | if isinstance(data, dict): |
| 892 | required, optional = self.embedding_field_sets("image") |
| 893 | return DictEmbeddingItems( |
| 894 | data, |
| 895 | modality="image", |
| 896 | required_fields=required, |
| 897 | optional_fields=optional, |
| 898 | fields_factory=_keye_field_config, |
| 899 | ) |
| 900 | |
| 901 | return super()._parse_image_data(data) |
| 902 | |
| 903 | def _parse_video_data( |
| 904 | self, |
| 905 | data: dict[str, torch.Tensor] | ModalityData[VideoItem], |
| 906 | ) -> ModalityDataItems[Any, Any] | None: |
| 907 | if isinstance(data, dict): |
| 908 | required, optional = self.embedding_field_sets("video") |
| 909 | return DictEmbeddingItems( |
| 910 | data, |
| 911 | modality="video", |
| 912 | required_fields=required, |
| 913 | optional_fields=optional, |
| 914 | fields_factory=_keye_field_config, |
| 915 | ) |
| 916 | |
| 917 | return super()._parse_video_data(data) |
| 918 | |
| 919 | |
| 920 | class KeyeProcessingInfo(BaseProcessingInfo): |
| 921 | def get_max_image_size(self) -> int: |
| 922 | return 9999999 # _MAX_IMAGE_SIZE |
| 923 | |
| 924 | def get_max_frame_per_video(self) -> int: |
| 925 | return 16 # _MAX_FRAMES_PER_VIDEO |
| 926 | |
| 927 | def get_image_processor(self, **kwargs: object): |
| 928 | return self.get_hf_processor(**kwargs).image_processor |
| 929 | |
| 930 | def get_data_parser(self): |
| 931 | return KeyeMultiModalDataParser( |
| 932 | expected_hidden_size=self._get_expected_hidden_size(), |
| 933 | allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, |
| 934 | ) |
| 935 | |
| 936 | def get_supported_mm_limits( |
| 937 | self, |
| 938 | ) -> Mapping[str, int | None]: |
| 939 | return {"image": None, "video": None} |
| 940 | |
| 941 | def get_mm_max_tokens_per_item( |
| 942 | self, |
| 943 | seq_len: int, |
| 944 | mm_counts: Mapping[str, int], |
| 945 | ) -> Mapping[str, int]: |
| 946 | return { |
| 947 | "image": self.get_max_image_tokens(), |
| 948 | "video": self.get_max_video_tokens(seq_len), |
| 949 | } |
| 950 | |
| 951 | def _get_vision_info( |
| 952 | self, |
| 953 | *, |
| 954 | image_width: int, |
| 955 | image_height: int, |
| 956 | num_frames: int = 1, |
| 957 | do_resize: bool = True, |
| 958 | image_processor: BaseImageProcessor, |
| 959 | mm_kwargs: Mapping[str, object], |
| 960 | ) -> tuple[ImageSize, int]: |
| 961 | hf_config = self.get_hf_config() |
| 962 | vision_config = hf_config.vision_config |
| 963 | patch_size = vision_config.patch_size |
| 964 | merge_size = vision_config.spatial_merge_size |
| 965 | temporal_patch_size = 1 |
| 966 | |
| 967 | mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs) |
| 968 | size = image_processor.size |
| 969 | if override_size := mm_kwargs.get("size"): |
| 970 | size = size | override_size |
| 971 | if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None: |
| 972 | size = size | {"min_pixels": override_min_pixels} |
| 973 | if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None: |
| 974 | size = size | {"max_pixels": override_max_pixels} |
| 975 | |
| 976 | if do_resize: |
| 977 | resized_height, resized_width = smart_resize( |
| 978 | height=image_height, |
| 979 | width=image_width, |
| 980 | factor=patch_size * merge_size, |
| 981 | min_pixels=size["min_pixels"], |
| 982 | max_pixels=size["max_pixels"], |
| 983 | ) |
| 984 | preprocessed_size = ImageSize(width=resized_width, height=resized_height) |
| 985 | else: |
| 986 | preprocessed_size = ImageSize(width=image_width, height=image_height) |
| 987 | |
| 988 | padded_num_frames = num_frames + (-num_frames % temporal_patch_size) |
| 989 | |
| 990 | grid_t = max(padded_num_frames // temporal_patch_size, 1) |
| 991 | grid_h = preprocessed_size.height // patch_size |
| 992 | grid_w = preprocessed_size.width // patch_size |
| 993 | |
| 994 | num_patches = grid_t * grid_h * grid_w |
| 995 | num_vision_tokens = num_patches // (merge_size**2) |
| 996 | |
| 997 | return preprocessed_size, num_vision_tokens |
| 998 | |
| 999 | def get_num_image_tokens( |
| 1000 | self, |
| 1001 | *, |
| 1002 | image_width: int, |
| 1003 | image_height: int, |
| 1004 | image_processor: BaseImageProcessor, |
| 1005 | mm_kwargs: Mapping[str, object], |
| 1006 | ) -> int: |
| 1007 | _, num_image_tokens = self._get_vision_info( |
| 1008 | image_width=image_width, |
| 1009 | image_height=image_height, |
| 1010 | image_processor=image_processor, |
| 1011 | mm_kwargs=mm_kwargs, |
| 1012 | ) |
| 1013 | return num_image_tokens |
| 1014 | |
| 1015 | def get_num_video_tokens( |
| 1016 | self, |
| 1017 | *, |
| 1018 | image_width: int, |
| 1019 | image_height: int, |
| 1020 | num_frames: int, |
| 1021 | image_processor: BaseImageProcessor, |
| 1022 | mm_kwargs: Mapping[str, object], |
| 1023 | ) -> int: |
| 1024 | _, num_video_tokens = self._get_vision_info( |
| 1025 | image_width=image_width, |
| 1026 | image_height=image_height, |
| 1027 | num_frames=num_frames, |
| 1028 | image_processor=image_processor, |
| 1029 | mm_kwargs=mm_kwargs, |
| 1030 | ) |
| 1031 | return num_video_tokens |
| 1032 | |
| 1033 | def get_image_size_with_most_features(self) -> ImageSize: |
| 1034 | image_processor = self.get_image_processor() |
| 1035 | |
| 1036 | max_image_size, _ = self._get_vision_info( |
| 1037 | image_width=self.get_max_image_size(), |
| 1038 | image_height=self.get_max_image_size(), |
| 1039 | image_processor=image_processor, |
| 1040 | mm_kwargs={}, |
| 1041 | ) |
| 1042 | return max_image_size |
| 1043 | |
| 1044 | def get_max_image_tokens(self) -> int: |
| 1045 | image_processor = self.get_image_processor() |
| 1046 | target_width, target_height = self.get_image_size_with_most_features() |
| 1047 | |
| 1048 | return self.get_num_image_tokens( |
| 1049 | image_width=target_width, |
| 1050 | image_height=target_height, |
| 1051 | image_processor=image_processor, |
| 1052 | mm_kwargs={}, |
| 1053 | ) |
| 1054 | |
| 1055 | def _get_max_video_frames(self, max_tokens: int) -> int: |
| 1056 | image_processor = self.get_image_processor() |
| 1057 | target_width, target_height = self.get_image_size_with_most_features() |
| 1058 | |
| 1059 | num_frames = 0 |
| 1060 | |
| 1061 | while True: |
| 1062 | next_num_frames = num_frames + 1 |
| 1063 | next_max_tokens = self.get_num_video_tokens( |
| 1064 | image_width=target_width, |
| 1065 | image_height=target_height, |
| 1066 | num_frames=next_num_frames, |
| 1067 | image_processor=image_processor, |
| 1068 | mm_kwargs={}, |
| 1069 | ) |
| 1070 | |
| 1071 | if next_max_tokens > max_tokens: |
| 1072 | break |
| 1073 | |
| 1074 | num_frames = next_num_frames |
| 1075 | |
| 1076 | return num_frames |
| 1077 | |
| 1078 | def get_num_frames_with_most_features(self, seq_len: int) -> int: |
| 1079 | mm_config = self.ctx.get_mm_config() |
| 1080 | max_images = mm_config.get_limit_per_prompt("image") |
| 1081 | max_videos = mm_config.get_limit_per_prompt("video") |
| 1082 | |
| 1083 | max_image_tokens = self.get_max_image_tokens() * max_images |
| 1084 | max_total_frames = self._get_max_video_frames(seq_len - max_image_tokens) |
| 1085 | max_frames_per_video = min( |
| 1086 | max_total_frames // max(max_videos, 1), |
| 1087 | self.get_max_frame_per_video(), |
| 1088 | ) |
| 1089 | |
| 1090 | return max(max_frames_per_video, 1) |
| 1091 | |
| 1092 | def get_max_video_tokens(self, seq_len: int) -> int: |
| 1093 | image_processor = self.get_image_processor() |
| 1094 | target_width, target_height = self.get_image_size_with_most_features() |
| 1095 | |
| 1096 | return self.get_num_video_tokens( |
| 1097 | image_width=target_width, |
| 1098 | image_height=target_height, |
| 1099 | num_frames=self.get_num_frames_with_most_features(seq_len), |
| 1100 | image_processor=image_processor, |
| 1101 | mm_kwargs={}, |
| 1102 | ) |
| 1103 | |
| 1104 | |
| 1105 | _I = TypeVar("_I", bound=KeyeProcessingInfo) |
| 1106 | |
| 1107 | |
| 1108 | class KeyeBaseDummyInputsBuilder(BaseDummyInputsBuilder[_I]): |
| 1109 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 1110 | num_images = mm_counts.get("image", 0) |
| 1111 | num_videos = mm_counts.get("video", 0) |
| 1112 | |
| 1113 | hf_processor = self.info.get_hf_processor() |
| 1114 | image_token: str = hf_processor.image_token |
| 1115 | video_token: str = hf_processor.video_token |
| 1116 | |
| 1117 | return image_token * num_images + video_token * num_videos |
| 1118 | |
| 1119 | def get_dummy_mm_data( |
| 1120 | self, |
| 1121 | seq_len: int, |
| 1122 | mm_counts: Mapping[str, int], |
| 1123 | mm_options: MultiModalDummyOptions, |
| 1124 | ) -> MultiModalDataDict: |
| 1125 | target_width, target_height = self.info.get_image_size_with_most_features() |
| 1126 | target_num_frames = self.info.get_num_frames_with_most_features(seq_len) |
| 1127 | |
| 1128 | mm_data = { |
| 1129 | "image": self._get_dummy_images( |
| 1130 | width=target_width, |
| 1131 | height=target_height, |
| 1132 | num_images=mm_counts.get("image", 0), |
| 1133 | overrides=mm_options.get("image"), |
| 1134 | ), |
| 1135 | "video": self._get_dummy_videos( |
| 1136 | width=target_width, |
| 1137 | height=target_height, |
| 1138 | num_frames=target_num_frames, |
| 1139 | num_videos=mm_counts.get("video", 0), |
| 1140 | overrides=mm_options.get("video"), |
| 1141 | ), |
| 1142 | } |
| 1143 | |
| 1144 | return mm_data |
| 1145 | |
| 1146 | |
| 1147 | class KeyeDummyInputsBuilder(KeyeBaseDummyInputsBuilder[KeyeProcessingInfo]): |
| 1148 | pass |
| 1149 | |
| 1150 | |
| 1151 | class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]): |
| 1152 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 1153 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 1154 | |
| 1155 | def _get_prompt_updates( |
| 1156 | self, |
| 1157 | mm_items: MultiModalDataItems, |
| 1158 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 1159 | out_mm_kwargs: MultiModalKwargsItems, |
| 1160 | ) -> Sequence[PromptUpdate]: |
| 1161 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 1162 | image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) |
| 1163 | tokenizer = self.info.get_tokenizer() |
| 1164 | vocab = tokenizer.get_vocab() |
| 1165 | |
| 1166 | placeholder = { |
| 1167 | "image": vocab[hf_processor.image_token], |
| 1168 | "video": vocab[hf_processor.video_token], |
| 1169 | } |
| 1170 | |
| 1171 | merge_length = image_processor.merge_size**2 |
| 1172 | |
| 1173 | def get_replacement_keye(item_idx: int, modality: str): |
| 1174 | out_item = out_mm_kwargs[modality][item_idx] |
| 1175 | grid_thw = out_item[f"{modality}_grid_thw"].data |
| 1176 | assert isinstance(grid_thw, torch.Tensor) |
| 1177 | |
| 1178 | num_tokens = int(grid_thw.prod()) // merge_length |
| 1179 | return [placeholder[modality]] * num_tokens |
| 1180 | |
| 1181 | return [ |
| 1182 | PromptReplacement( |
| 1183 | modality=modality, |
| 1184 | target=[placeholder[modality]], |
| 1185 | replacement=partial(get_replacement_keye, modality=modality), |
| 1186 | ) |
| 1187 | for modality in ("image", "video") |
| 1188 | ] |
| 1189 | |
| 1190 | def _get_mm_fields_config( |
| 1191 | self, |
| 1192 | hf_inputs: BatchFeature, |
| 1193 | hf_processor_mm_kwargs: Mapping[str, object], |
| 1194 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 1195 | return _keye_field_config(hf_inputs) |
| 1196 | |
| 1197 | |
| 1198 | class BaseKeyeModule(nn.Module, SupportsMultiModal): |
| 1199 | packed_modules_mapping = { |
| 1200 | "qkv_proj": [ |
| 1201 | "q_proj", |
| 1202 | "k_proj", |
| 1203 | "v_proj", |
| 1204 | ], |
| 1205 | "gate_up_proj": [ |
| 1206 | "gate_proj", |
| 1207 | "up_proj", |
| 1208 | ], |
| 1209 | } |
| 1210 | |
| 1211 | hf_to_vllm_mapper = WeightsMapper( |
| 1212 | orig_to_new_prefix={ |
| 1213 | "lm_head.": "language_model.lm_head.", |
| 1214 | "model.": "language_model.model.", |
| 1215 | } |
| 1216 | ) |
| 1217 | |
| 1218 | supports_tower_connector_lora = True |
| 1219 | |
| 1220 | @classmethod |
| 1221 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 1222 | if modality.startswith("image"): |
| 1223 | return "<|vision_start|><|image_pad|><|vision_end|>" |
| 1224 | if modality.startswith("video"): |
| 1225 | return "<|vision_start|><|video_pad|><|vision_end|>" |
| 1226 | |
| 1227 | raise ValueError("Only image or video modality is supported") |
| 1228 | |
| 1229 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1230 | super().__init__() |
| 1231 | config: PreTrainedConfig = vllm_config.model_config.hf_config |
| 1232 | quant_config = vllm_config.quant_config |
| 1233 | |
| 1234 | self.config = config |
| 1235 | |
| 1236 | with self._mark_tower_model(vllm_config, {"image", "video"}): |
| 1237 | self.visual = KeyeSiglipVisionModel( |
| 1238 | config.vision_config, |
| 1239 | quant_config=quant_config, |
| 1240 | prefix=maybe_prefix(prefix, "visual"), |
| 1241 | ) |
| 1242 | self.mlp_AR = self._build_projector( |
| 1243 | config, |
| 1244 | config.vision_config, |
| 1245 | quant_config=quant_config, |
| 1246 | prefix=maybe_prefix(prefix, "mlp_AR"), |
| 1247 | ) |
| 1248 | |
| 1249 | with self._mark_language_model(vllm_config): |
| 1250 | self.language_model = init_vllm_registered_model( |
| 1251 | vllm_config=vllm_config, |
| 1252 | prefix=maybe_prefix(prefix, "language_model"), |
| 1253 | architectures=["Qwen3ForCausalLM"], |
| 1254 | ) |
| 1255 | |
| 1256 | self.make_empty_intermediate_tensors = ( |
| 1257 | self.language_model.make_empty_intermediate_tensors |
| 1258 | ) |
| 1259 | |
| 1260 | @abstractmethod |
| 1261 | def _build_projector( |
| 1262 | self, |
| 1263 | text_config: PreTrainedConfig, |
| 1264 | vision_config: PreTrainedConfig, |
| 1265 | quant_config: QuantizationConfig | None = None, |
| 1266 | prefix: str = "", |
| 1267 | ) -> nn.Module: |
| 1268 | raise NotImplementedError("Need projector") |
| 1269 | |
| 1270 | def _process_image_input(self, image_input: Any) -> tuple[torch.Tensor, ...]: |
| 1271 | siglip_position_ids = list() |
| 1272 | image_grid_hws = list() |
| 1273 | sample_indices = list() |
| 1274 | cu_seqlens = [0] |
| 1275 | |
| 1276 | image_grid_thw = image_input["image_grid_thw"] |
| 1277 | assert image_grid_thw.ndim == 2 |
| 1278 | |
| 1279 | image_grid_thw_list = image_grid_thw.tolist() |
| 1280 | for idx, thw in enumerate(image_grid_thw_list): |
| 1281 | thw_tuple = tuple(thw) |
| 1282 | numel = np.prod(thw_tuple) |
| 1283 | image_grid_hws.append(thw_tuple) |
| 1284 | image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:]) |
| 1285 | siglip_position_ids.append(image_position_ids) |
| 1286 | sample_indices.append(torch.full((numel,), idx, dtype=torch.int64)) |
| 1287 | cu_seqlens.append(cu_seqlens[-1] + numel) |
| 1288 | |
| 1289 | if image_input["type"] == "image_embeds": |
| 1290 | raise ValueError( |
| 1291 | "Image embeddings are not supported for this processing path." |
| 1292 | ) |
| 1293 | else: |
| 1294 | pixel_values = image_input["pixel_values"].type(self.visual.dtype) |
| 1295 | # These are all built on the host; concat straight into pinned |
| 1296 | # buffers so the H2D copies stay non-blocking. |
| 1297 | siglip_position_ids = torch.concat( |
| 1298 | siglip_position_ids, |
| 1299 | dim=0, |
| 1300 | out=torch.empty( |
| 1301 | sum(t.numel() for t in siglip_position_ids), |
| 1302 | dtype=torch.int64, |
| 1303 | pin_memory=PIN_MEMORY, |
| 1304 | ), |
| 1305 | ).to(pixel_values.device, non_blocking=True) |
| 1306 | cu_seqlens = async_tensor_h2d( |
| 1307 | cu_seqlens, dtype=torch.int32, device=pixel_values.device |
| 1308 | ) |
| 1309 | sample_indices = torch.concat( |
| 1310 | sample_indices, |
| 1311 | dim=0, |
| 1312 | out=torch.empty( |
| 1313 | sum(t.numel() for t in sample_indices), |
| 1314 | dtype=torch.int64, |
| 1315 | pin_memory=PIN_MEMORY, |
| 1316 | ), |
| 1317 | ).to(pixel_values.device, non_blocking=True) |
| 1318 | |
| 1319 | image_embeds = self.visual( |
| 1320 | pixel_values=pixel_values, |
| 1321 | image_grid_thw=image_grid_hws, |
| 1322 | position_ids=siglip_position_ids, |
| 1323 | vision_return_embed_list=False, |
| 1324 | interpolate_pos_encoding=True, |
| 1325 | sample_indices=sample_indices, |
| 1326 | cu_seqlens=cu_seqlens, |
| 1327 | use_rope=True, |
| 1328 | window_size=-1, |
| 1329 | ) |
| 1330 | image_embeds = tuple(self.mlp_AR(image_embeds, image_grid_thw)) |
| 1331 | return image_embeds |
| 1332 | |
| 1333 | def _process_video_embeds( |
| 1334 | self, |
| 1335 | video_type: Literal["video_embeds", "pixel_values_videos"], |
| 1336 | video_grid_thw: torch.Tensor, |
| 1337 | pixel_values_videos: torch.Tensor | None = None, |
| 1338 | ) -> torch.Tensor | list[torch.Tensor]: |
| 1339 | siglip_position_ids = list() |
| 1340 | video_grid_hws = list() |
| 1341 | sample_indices = list() |
| 1342 | cu_seqlens = [0] |
| 1343 | |
| 1344 | assert video_grid_thw.ndim == 2 |
| 1345 | for idx, sub_thw in enumerate(video_grid_thw): |
| 1346 | thw_tuple = tuple(sub_thw.detach().cpu().numpy().tolist()) |
| 1347 | numel = np.prod(thw_tuple) |
| 1348 | |
| 1349 | video_grid_hws.append(thw_tuple) |
| 1350 | video_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:]) |
| 1351 | siglip_position_ids.append(video_position_ids) |
| 1352 | sample_indices.append(torch.full((numel,), idx, dtype=torch.int64)) |
| 1353 | cu_seqlens.append(cu_seqlens[-1] + numel) |
| 1354 | |
| 1355 | if video_type == "video_embeds": |
| 1356 | raise ValueError( |
| 1357 | "Video embeddings are not supported for this processing path." |
| 1358 | ) |
| 1359 | else: |
| 1360 | assert pixel_values_videos is not None |
| 1361 | pixel_values_videos = pixel_values_videos.type(self.visual.dtype) |
| 1362 | # These are host-built; concat straight into pinned buffers so |
| 1363 | # the H2D copies stay non-blocking. |
| 1364 | siglip_position_ids = torch.concat( |
| 1365 | siglip_position_ids, |
| 1366 | dim=0, |
| 1367 | out=torch.empty( |
| 1368 | sum(t.numel() for t in siglip_position_ids), |
| 1369 | dtype=torch.int64, |
| 1370 | pin_memory=PIN_MEMORY, |
| 1371 | ), |
| 1372 | ).to(pixel_values_videos.device, non_blocking=True) |
| 1373 | cu_seqlens = async_tensor_h2d( |
| 1374 | cu_seqlens, dtype=torch.int32, device=pixel_values_videos.device |
| 1375 | ) |
| 1376 | sample_indices = torch.concat( |
| 1377 | sample_indices, |
| 1378 | dim=0, |
| 1379 | out=torch.empty( |
| 1380 | sum(t.numel() for t in sample_indices), |
| 1381 | dtype=torch.int64, |
| 1382 | pin_memory=PIN_MEMORY, |
| 1383 | ), |
| 1384 | ).to(pixel_values_videos.device, non_blocking=True) |
| 1385 | |
| 1386 | video_embeds = self.visual( |
| 1387 | pixel_values=pixel_values_videos, |
| 1388 | image_grid_thw=video_grid_hws, |
| 1389 | position_ids=siglip_position_ids, |
| 1390 | vision_return_embed_list=True, |
| 1391 | interpolate_pos_encoding=True, |
| 1392 | sample_indices=sample_indices, |
| 1393 | cu_seqlens=cu_seqlens, |
| 1394 | use_rope=True, |
| 1395 | window_size=-1, |
| 1396 | ) |
| 1397 | video_embeds = self.mlp_AR(video_embeds, video_grid_thw) |
| 1398 | return video_embeds |
| 1399 | |
| 1400 | def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: |
| 1401 | modalities = {} |
| 1402 | |
| 1403 | for input_key in kwargs: |
| 1404 | if ( |
| 1405 | input_key in ("pixel_values", "image_embeds") |
| 1406 | and "images" not in modalities |
| 1407 | ): |
| 1408 | modalities["images"] = self._parse_and_validate_image_input(**kwargs) |
| 1409 | if ( |
| 1410 | input_key in ("pixel_values_videos", "video_embeds") |
| 1411 | and "videos" not in modalities |
| 1412 | ): |
| 1413 | modalities["videos"] = self._parse_and_validate_video_input(**kwargs) |
| 1414 | |
| 1415 | return modalities |
| 1416 | |
| 1417 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: |
| 1418 | modalities = self._parse_and_validate_multimodal_inputs(**kwargs) |
| 1419 | if not modalities: |
| 1420 | return None |
| 1421 | |
| 1422 | multimodal_embeddings: tuple[torch.Tensor, ...] = () |
| 1423 | |
| 1424 | for modality in modalities: |
| 1425 | if modality == "images": |
| 1426 | image_input = modalities["images"] |
| 1427 | image_embeddings = self._process_image_input(image_input) |
| 1428 | multimodal_embeddings += tuple(image_embeddings) |
| 1429 | if modality == "videos": |
| 1430 | video_input = modalities["videos"] |
| 1431 | video_embeddings = self._process_video_input(video_input) |
| 1432 | multimodal_embeddings += tuple(video_embeddings) |
| 1433 | return multimodal_embeddings |
| 1434 | |
| 1435 | def forward( |
| 1436 | self, |
| 1437 | input_ids: torch.Tensor | None, |
| 1438 | positions: torch.Tensor, |
| 1439 | intermediate_tensors: IntermediateTensors | None = None, |
| 1440 | inputs_embeds: torch.Tensor | None = None, |
| 1441 | **kwargs: object, |
| 1442 | ) -> torch.Tensor | IntermediateTensors: |
| 1443 | """Run forward pass for Keye-VL. |
| 1444 | |
| 1445 | Args: |
| 1446 | input_ids: Flattened (concatenated) input_ids corresponding to a |
| 1447 | batch. |
| 1448 | positions: Flattened (concatenated) position ids corresponding to a |
| 1449 | batch. |
| 1450 | **NOTE**: If mrope is enabled (default setting for Qwen2-VL |
| 1451 | opensource models), the shape will be `(3, seq_len)`, |
| 1452 | otherwise it will be `(seq_len,)`. |
| 1453 | intermediate_tensors: Intermediate tensors from prior forward pass. |
| 1454 | inputs_embeds: Optional tensor of input embeddings. |
| 1455 | **kwargs: Multimodal inputs for this batch, forwarded to the |
| 1456 | multimodal embedding path. |
| 1457 | |
| 1458 | """ |
| 1459 | if intermediate_tensors is not None: |
| 1460 | inputs_embeds = None |
| 1461 | |
| 1462 | hidden_states = self.language_model.model( |
| 1463 | input_ids=input_ids, |
| 1464 | positions=positions, |
| 1465 | intermediate_tensors=intermediate_tensors, |
| 1466 | inputs_embeds=inputs_embeds, |
| 1467 | ) |
| 1468 | |
| 1469 | return hidden_states |
| 1470 | |
| 1471 | def compute_logits( |
| 1472 | self, |
| 1473 | hidden_states: torch.Tensor, |
| 1474 | ) -> torch.Tensor | None: |
| 1475 | return self.language_model.compute_logits(hidden_states) |
| 1476 | |
| 1477 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1478 | loader = AutoWeightsLoader(self) |
| 1479 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 1480 | |
| 1481 | def get_mm_mapping(self) -> MultiModelKeys: |
| 1482 | """Get the module prefix in multimodal models.""" |
| 1483 | return MultiModelKeys.from_string_field( |
| 1484 | language_model="language_model", |
| 1485 | connector="mlp_AR.", |
| 1486 | tower_model="visual.", |
| 1487 | ) |
| 1488 | |
| 1489 | def get_mm_lora_token_counts( |
| 1490 | self, |
| 1491 | *, |
| 1492 | modality: str, |
| 1493 | mm_kwargs: MultiModalKwargsItem | None, |
| 1494 | num_mm_embeds: int, |
| 1495 | ) -> tuple[int, int | None]: |
| 1496 | del modality, mm_kwargs |
| 1497 | merge_size = self.config.vision_config.spatial_merge_size |
| 1498 | return num_mm_embeds * merge_size**2, num_mm_embeds |
| 1499 | |
| 1500 | |
| 1501 | @MULTIMODAL_REGISTRY.register_processor( |
| 1502 | KeyeMultiModalProcessor, |
| 1503 | info=KeyeProcessingInfo, |
| 1504 | dummy_inputs=KeyeDummyInputsBuilder, |
| 1505 | ) |
| 1506 | class KeyeForConditionalGeneration( |
| 1507 | BaseKeyeModule, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE |
| 1508 | ): |
| 1509 | def _build_projector( |
| 1510 | self, |
| 1511 | text_config: PreTrainedConfig, |
| 1512 | vision_config: PreTrainedConfig, |
| 1513 | quant_config: QuantizationConfig | None = None, |
| 1514 | prefix: str = "", |
| 1515 | ) -> nn.Module: |
| 1516 | return Projector(text_config, vision_config, quant_config, prefix) |
| 1517 | |
| 1518 | def _parse_and_validate_image_input( |
| 1519 | self, **kwargs: object |
| 1520 | ) -> KeyeImageInputs | None: |
| 1521 | pixel_values = kwargs.pop("pixel_values", None) |
| 1522 | image_embeds = kwargs.pop("image_embeds", None) |
| 1523 | image_grid_thw = kwargs.pop("image_grid_thw", None) |
| 1524 | |
| 1525 | if pixel_values is None and image_embeds is None: |
| 1526 | return None |
| 1527 | |
| 1528 | if pixel_values is not None: |
| 1529 | return KeyeImagePixelInputs( |
| 1530 | type="pixel_values", |
| 1531 | pixel_values=pixel_values, |
| 1532 | image_grid_thw=image_grid_thw, |
| 1533 | ) |
| 1534 | |
| 1535 | return KeyeImageEmbeddingInputs( |
| 1536 | type="image_embeds", |
| 1537 | image_embeds=image_embeds, |
| 1538 | image_grid_thw=image_grid_thw, |
| 1539 | ) |
| 1540 | |
| 1541 | def _parse_and_validate_video_input( |
| 1542 | self, **kwargs: object |
| 1543 | ) -> KeyeVideoInputs | None: |
| 1544 | pixel_values_videos = kwargs.pop("pixel_values_videos", None) |
| 1545 | video_embeds = kwargs.pop("video_embeds", None) |
| 1546 | video_grid_thw = kwargs.pop("video_grid_thw", None) |
| 1547 | |
| 1548 | if pixel_values_videos is None and video_embeds is None: |
| 1549 | return None |
| 1550 | |
| 1551 | if pixel_values_videos is not None: |
| 1552 | return KeyeVideoPixelInputs( |
| 1553 | type="pixel_values_videos", |
| 1554 | pixel_values_videos=pixel_values_videos, |
| 1555 | video_grid_thw=video_grid_thw, |
| 1556 | ) |
| 1557 | |
| 1558 | return KeyeVideoEmbeddingInputs( |
| 1559 | type="video_embeds", |
| 1560 | video_embeds=video_embeds, |
| 1561 | video_grid_thw=video_grid_thw, |
| 1562 | ) |
| 1563 | |
| 1564 | def _process_video_input( |
| 1565 | self, video_input: KeyeVideoInputs |
| 1566 | ) -> tuple[torch.Tensor, ...]: |
| 1567 | video_type = video_input["type"] |
| 1568 | video_grid_thw = video_input["video_grid_thw"] |
| 1569 | pixel_values_videos = video_input.get("pixel_values_videos", None) |
| 1570 | |
| 1571 | return tuple( |
| 1572 | self._process_video_embeds(video_type, video_grid_thw, pixel_values_videos) |
| 1573 | ) |
| 1574 | |
| 1575 | @staticmethod |
| 1576 | def _split_video_grid_thw( |
| 1577 | grid_thw: torch.Tensor | list[list[int]] | list[int], |
| 1578 | ) -> list[list[int]]: |
| 1579 | """Split video grid_thw along the t dimension into per-frame rows. |
| 1580 | |
| 1581 | This preserves Keye's current M-RoPE behavior, where a video is emitted |
| 1582 | as consecutive frame-level multimodal blocks rather than a single block |
| 1583 | spanning the whole video. |
| 1584 | """ |
| 1585 | if isinstance(grid_thw, list): |
| 1586 | if len(grid_thw) == 0: |
| 1587 | return [] |
| 1588 | if isinstance(grid_thw[0], int): |
| 1589 | grid_thw = torch.tensor([grid_thw], dtype=torch.long) |
| 1590 | else: |
| 1591 | grid_thw = torch.tensor(grid_thw, dtype=torch.long) |
| 1592 | elif grid_thw.ndim == 1: |
| 1593 | grid_thw = grid_thw.unsqueeze(0) |
| 1594 | |
| 1595 | if grid_thw.numel() == 0: |
| 1596 | return [] |
| 1597 | |
| 1598 | t, hw = grid_thw[:, 0], grid_thw[:, 1:] |
| 1599 | ones = torch.ones_like(hw[:, :1]) |
| 1600 | out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0) |
| 1601 | return out.tolist() |
| 1602 | |
| 1603 | def iter_mm_grid_thw( |
| 1604 | self, mm_features: list[MultiModalFeatureSpec] |
| 1605 | ) -> Iterator[tuple[int, int, int, int]]: |
| 1606 | spatial_merge_size = self.config.vision_config.spatial_merge_size |
| 1607 | |
| 1608 | for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): |
| 1609 | if mm_feature.data is None: |
| 1610 | raise ValueError("M-RoPE calculation requires multimodal feature data") |
| 1611 | |
| 1612 | if mm_feature.modality == "image": |
| 1613 | grid_thw = mm_feature.data["image_grid_thw"].data |
| 1614 | if isinstance(grid_thw, torch.Tensor): |
| 1615 | if grid_thw.ndim == 2: |
| 1616 | assert grid_thw.shape[0] == 1 |
| 1617 | t, h, w = grid_thw[0].tolist() |
| 1618 | else: |
| 1619 | t, h, w = grid_thw.tolist() |
| 1620 | else: |
| 1621 | if isinstance(grid_thw[0], list): |
| 1622 | assert len(grid_thw) == 1 |
| 1623 | t, h, w = grid_thw[0] |
| 1624 | else: |
| 1625 | t, h, w = grid_thw |
| 1626 | |
| 1627 | yield ( |
| 1628 | mm_feature.mm_position.offset, |
| 1629 | t, |
| 1630 | h // spatial_merge_size, |
| 1631 | w // spatial_merge_size, |
| 1632 | ) |
| 1633 | elif mm_feature.modality == "video": |
| 1634 | current_offset = mm_feature.mm_position.offset |
| 1635 | for t, h, w in self._split_video_grid_thw( |
| 1636 | mm_feature.data["video_grid_thw"].data |
| 1637 | ): |
| 1638 | llm_grid_h = h // spatial_merge_size |
| 1639 | llm_grid_w = w // spatial_merge_size |
| 1640 | yield (current_offset, t, llm_grid_h, llm_grid_w) |
| 1641 | current_offset += t * llm_grid_h * llm_grid_w |
| 1642 | else: |
| 1643 | raise ValueError(f"Unsupported modality: {mm_feature.modality}") |
| 1644 | |
| 1645 | def get_mrope_input_positions( |
| 1646 | self, |
| 1647 | input_tokens: list[int], |
| 1648 | mm_features: list[MultiModalFeatureSpec], |
| 1649 | ) -> tuple[torch.Tensor, int]: |
| 1650 | llm_pos_ids_list: list = [] |
| 1651 | st = 0 |
| 1652 | |
| 1653 | for ( |
| 1654 | offset, |
| 1655 | llm_grid_t, |
| 1656 | llm_grid_h, |
| 1657 | llm_grid_w, |
| 1658 | ) in self.iter_mm_grid_thw(mm_features): |
| 1659 | text_len = offset - st |
| 1660 | |
| 1661 | st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 |
| 1662 | llm_pos_ids_list.append( |
| 1663 | torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx |
| 1664 | ) |
| 1665 | |
| 1666 | t_index = ( |
| 1667 | ( |
| 1668 | torch.arange(llm_grid_t) |
| 1669 | .view(-1, 1) |
| 1670 | .expand(-1, llm_grid_h * llm_grid_w) |
| 1671 | ) |
| 1672 | .long() |
| 1673 | .flatten() |
| 1674 | ) |
| 1675 | |
| 1676 | h_index = ( |
| 1677 | torch.arange(llm_grid_h) |
| 1678 | .view(1, -1, 1) |
| 1679 | .expand(llm_grid_t, -1, llm_grid_w) |
| 1680 | .flatten() |
| 1681 | ) |
| 1682 | w_index = ( |
| 1683 | torch.arange(llm_grid_w) |
| 1684 | .view(1, 1, -1) |
| 1685 | .expand(llm_grid_t, llm_grid_h, -1) |
| 1686 | .flatten() |
| 1687 | ) |
| 1688 | llm_pos_ids_list.append( |
| 1689 | torch.stack([t_index, h_index, w_index]) + text_len + st_idx |
| 1690 | ) |
| 1691 | st = offset + llm_grid_t * llm_grid_h * llm_grid_w |
| 1692 | |
| 1693 | if st < len(input_tokens): |
| 1694 | st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 |
| 1695 | text_len = len(input_tokens) - st |
| 1696 | llm_pos_ids_list.append( |
| 1697 | torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx |
| 1698 | ) |
| 1699 | |
| 1700 | llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) |
| 1701 | mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() |
| 1702 | |
| 1703 | return llm_positions, mrope_position_delta |