| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. |
| 4 | # |
| 5 | # NVIDIA CORPORATION and its licensors retain all intellectual property |
| 6 | # and proprietary rights in and to this software, related documentation |
| 7 | # and any modifications thereto. Any use, reproduction, disclosure or |
| 8 | # distribution of this software and related documentation without an express |
| 9 | # license agreement from NVIDIA CORPORATION is strictly prohibited. |
| 10 | |
| 11 | import math |
| 12 | from collections.abc import Iterable |
| 13 | from dataclasses import dataclass |
| 14 | from itertools import accumulate, repeat |
| 15 | from typing import TypeAlias |
| 16 | |
| 17 | import torch |
| 18 | import torch.nn as nn |
| 19 | import torch.nn.functional as F |
| 20 | from einops import rearrange |
| 21 | from transformers import PreTrainedConfig |
| 22 | |
| 23 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 24 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 25 | from vllm.model_executor.models.intern_vit import ( |
| 26 | InternParallelAttention, |
| 27 | InternVisionEncoder, |
| 28 | InternVisionEncoderLayer, |
| 29 | ) |
| 30 | |
| 31 | input_dim_t: TypeAlias = int | tuple[int, int] |
| 32 | norm_t: TypeAlias = tuple[float, float, float] | torch.Tensor |
| 33 | |
| 34 | |
| 35 | def _ntuple(n): |
| 36 | def parse(x): |
| 37 | if isinstance(x, Iterable) and not isinstance(x, str): |
| 38 | return tuple(x) |
| 39 | return tuple(repeat(x, n)) |
| 40 | |
| 41 | return parse |
| 42 | |
| 43 | |
| 44 | to_1tuple = _ntuple(1) |
| 45 | to_2tuple = _ntuple(2) |
| 46 | to_3tuple = _ntuple(3) |
| 47 | to_4tuple = _ntuple(4) |
| 48 | to_ntuple = _ntuple |
| 49 | |
| 50 | |
| 51 | def calc_seq_len(size: tuple[int, int], patch_size: int) -> int: |
| 52 | h, w = size |
| 53 | return (h // patch_size) * (w // patch_size) |
| 54 | |
| 55 | |
| 56 | def calc_seq_lens(sizes: list[tuple[int, int]], patch_size: int) -> list[int]: |
| 57 | return [calc_seq_len(size, patch_size) for size in sizes] |
| 58 | |
| 59 | |
| 60 | class ClsToken(nn.Module): |
| 61 | def __init__( |
| 62 | self, |
| 63 | ndim: int, |
| 64 | num_tokens: int = 1, |
| 65 | enabled: bool = True, |
| 66 | register_multiple: int | None = None, |
| 67 | num_registers: int | None = None, |
| 68 | ): |
| 69 | super().__init__() |
| 70 | |
| 71 | self.ndim = ndim |
| 72 | self.enabled = enabled |
| 73 | self.num_registers = 0 |
| 74 | self.num_tokens = num_tokens |
| 75 | if enabled: |
| 76 | if num_registers: |
| 77 | self.num_registers = num_registers |
| 78 | elif register_multiple: |
| 79 | self.num_registers = register_multiple - ( |
| 80 | num_tokens % register_multiple |
| 81 | ) |
| 82 | |
| 83 | scale = ndim**-0.5 |
| 84 | self.token = nn.Parameter( |
| 85 | torch.randn(num_tokens + self.num_registers, ndim) * scale |
| 86 | ) |
| 87 | |
| 88 | else: |
| 89 | self.token = None |
| 90 | |
| 91 | self.num_patches = self.num_tokens + self.num_registers |
| 92 | |
| 93 | def forward(self, x: torch.Tensor): |
| 94 | if self.token is None: |
| 95 | return x |
| 96 | |
| 97 | token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1) |
| 98 | x = torch.cat( |
| 99 | [ |
| 100 | token, |
| 101 | x, |
| 102 | ], |
| 103 | dim=1, |
| 104 | ) |
| 105 | |
| 106 | return x |
| 107 | |
| 108 | |
| 109 | class ViTPatchGenerator(nn.Module): |
| 110 | def __init__( |
| 111 | self, |
| 112 | # config: PreTrainedConfig, |
| 113 | patch_size: int, |
| 114 | embed_dim: int, |
| 115 | input_dims: input_dim_t, |
| 116 | abs_pos: bool = True, |
| 117 | normalize_patches: bool = False, |
| 118 | cls_token: bool = False, |
| 119 | max_input_dims: input_dim_t | None = None, |
| 120 | pos_dropout: float = 0.0, |
| 121 | return_pos_enc: bool = False, |
| 122 | num_cls_tokens: int = 1, |
| 123 | register_multiple: int | None = None, |
| 124 | num_registers: int | None = None, |
| 125 | patch_bias: bool = False, |
| 126 | temporal_patch_size: int = 1, |
| 127 | separate_video_embedder: bool = True, |
| 128 | device=None, |
| 129 | dtype=None, |
| 130 | ): |
| 131 | super().__init__() |
| 132 | if isinstance(input_dims, int): |
| 133 | input_dims = (input_dims, input_dims) |
| 134 | |
| 135 | if max_input_dims is None: |
| 136 | max_input_dims = input_dims |
| 137 | if isinstance(max_input_dims, int): |
| 138 | max_input_dims = (max_input_dims, max_input_dims) |
| 139 | |
| 140 | max_input_dims = ( |
| 141 | int(math.ceil(max_input_dims[0] / patch_size) * patch_size), |
| 142 | int(math.ceil(max_input_dims[1] / patch_size) * patch_size), |
| 143 | ) |
| 144 | |
| 145 | self.cpe_mode = max_input_dims != input_dims |
| 146 | self.pos_dropout = pos_dropout |
| 147 | self.return_pos_enc = return_pos_enc |
| 148 | |
| 149 | factory = dict(device=device, dtype=dtype) |
| 150 | |
| 151 | self.patch_size = patch_size |
| 152 | self.abs_pos = abs_pos |
| 153 | self.embed_dim = embed_dim |
| 154 | self.temporal_patch_size = temporal_patch_size |
| 155 | |
| 156 | self.num_rows = max_input_dims[0] // patch_size |
| 157 | self.num_cols = max_input_dims[1] // patch_size |
| 158 | self.input_dims = ( |
| 159 | input_dims[0] // patch_size, |
| 160 | input_dims[1] // patch_size, |
| 161 | ) |
| 162 | self.num_patches = self.num_rows * self.num_cols |
| 163 | self.max_input_dims = max_input_dims |
| 164 | |
| 165 | self.im_to_patches = Im2Patches(patch_size) |
| 166 | self.embedder = ViTPatchLinear( |
| 167 | patch_size, embed_dim, bias=patch_bias, **factory |
| 168 | ) |
| 169 | |
| 170 | if temporal_patch_size > 1: |
| 171 | if not separate_video_embedder: |
| 172 | raise NotImplementedError( |
| 173 | "Only separate_video_embedder=True is supported for" |
| 174 | " temporal compression (temporal_patch_size > 1)" |
| 175 | ) |
| 176 | self.video_embedder = ViTPatchLinear( |
| 177 | patch_size, |
| 178 | embed_dim, |
| 179 | bias=patch_bias, |
| 180 | temporal_patch_size=temporal_patch_size, |
| 181 | **factory, |
| 182 | ) |
| 183 | |
| 184 | if abs_pos: |
| 185 | scale = embed_dim**-0.5 |
| 186 | self.pos_embed = nn.Parameter( |
| 187 | torch.randn(1, self.num_patches, embed_dim, **factory) * scale |
| 188 | ) |
| 189 | |
| 190 | self.cls_token = ClsToken( |
| 191 | embed_dim, |
| 192 | num_tokens=num_cls_tokens, |
| 193 | enabled=cls_token, |
| 194 | register_multiple=register_multiple, |
| 195 | num_registers=num_registers, |
| 196 | ) |
| 197 | |
| 198 | self.patch_normalizer = ( |
| 199 | nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity() |
| 200 | ) |
| 201 | |
| 202 | def forward( |
| 203 | self, x: torch.Tensor, imgs_sizes: list[tuple[int, int]] | None = None |
| 204 | ) -> torch.Tensor: |
| 205 | if imgs_sizes is not None: |
| 206 | patches = self.embedder(x) |
| 207 | patches, pos_enc = self.apply_pos_enc_dynamic( |
| 208 | patches, imgs_sizes=imgs_sizes |
| 209 | ) |
| 210 | patches = self.cls_token_dynamic(patches, imgs_sizes=imgs_sizes) |
| 211 | else: |
| 212 | patches = self.embed_patches(x) |
| 213 | patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:]) |
| 214 | patches = self.cls_token(patches) |
| 215 | patches = self.patch_normalizer(patches) |
| 216 | if self.return_pos_enc: |
| 217 | return patches, pos_enc |
| 218 | return patches |
| 219 | |
| 220 | def forward_video(self, x: torch.Tensor) -> torch.Tensor: |
| 221 | """Process video frames with temporal compression. |
| 222 | |
| 223 | Groups T consecutive frames into tubelets before embedding. |
| 224 | |
| 225 | Args: |
| 226 | x: [num_frames, 3, H, W] tensor of video frames |
| 227 | |
| 228 | Returns: |
| 229 | Embedded patches with temporal compression applied. |
| 230 | |
| 231 | """ |
| 232 | assert self.temporal_patch_size > 1 |
| 233 | T = self.temporal_patch_size |
| 234 | input_size = x.shape[2:] |
| 235 | |
| 236 | patches = self.im_to_patches(x) # [N, num_patches, 3*P*P] |
| 237 | num_frames, num_spatial, feat_dim = patches.shape |
| 238 | |
| 239 | # Pad to a multiple of T by repeating the last frame so that |
| 240 | # all tubelets have exactly T frames. |
| 241 | num_pad_frames = (-num_frames) % T |
| 242 | if num_pad_frames > 0: |
| 243 | last_frame_dup = patches[-1:].expand(num_pad_frames, -1, -1) |
| 244 | patches = torch.cat([patches, last_frame_dup], dim=0) |
| 245 | |
| 246 | # Group T frames per tubelet: for each spatial position, concatenate |
| 247 | # features across T consecutive frames; order follows Megatron training |
| 248 | num_frames_padded = patches.shape[0] |
| 249 | num_tublets = num_frames_padded // T |
| 250 | patches = rearrange( |
| 251 | patches, |
| 252 | "(tubelets frames) spatial feat -> tubelets spatial (frames feat)", |
| 253 | tubelets=num_tublets, |
| 254 | frames=T, |
| 255 | spatial=num_spatial, |
| 256 | feat=feat_dim, |
| 257 | ) |
| 258 | |
| 259 | patches = self.video_embedder(patches) |
| 260 | |
| 261 | patches, pos_enc = self.apply_pos_enc(patches, input_size=input_size) |
| 262 | |
| 263 | patches = self.cls_token(patches) |
| 264 | |
| 265 | patches = self.patch_normalizer(patches) |
| 266 | if self.return_pos_enc: |
| 267 | return patches, pos_enc |
| 268 | return patches |
| 269 | |
| 270 | def apply_pos_enc_dynamic( |
| 271 | self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]] |
| 272 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 273 | if not self.abs_pos: |
| 274 | return patches, None |
| 275 | |
| 276 | current_length = 0 |
| 277 | pos_enc_list = [] |
| 278 | |
| 279 | for size in imgs_sizes: |
| 280 | seq_length = calc_seq_len(size, self.patch_size) |
| 281 | |
| 282 | img_patches = patches[:, current_length : current_length + seq_length, :] |
| 283 | pos_enc = self.get_pos_enc(patches.shape[0], input_size=size) |
| 284 | img_patches_with_pos = img_patches + pos_enc |
| 285 | |
| 286 | patches = torch.cat( |
| 287 | [ |
| 288 | patches[:, :current_length, :], |
| 289 | img_patches_with_pos, |
| 290 | patches[:, current_length + seq_length :, :], |
| 291 | ], |
| 292 | dim=1, |
| 293 | ) |
| 294 | pos_enc_list.append(pos_enc) |
| 295 | current_length += seq_length |
| 296 | |
| 297 | full_pos_enc = torch.cat(pos_enc_list, dim=1) if pos_enc_list else None |
| 298 | return patches, full_pos_enc |
| 299 | |
| 300 | def cls_token_dynamic( |
| 301 | self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]] |
| 302 | ) -> torch.Tensor: |
| 303 | if not self.cls_token.enabled: |
| 304 | return patches |
| 305 | |
| 306 | out = [] |
| 307 | current_length = 0 |
| 308 | |
| 309 | for seq_len in calc_seq_lens(imgs_sizes, self.patch_size): |
| 310 | class_token = self.cls_token.token.unsqueeze(0).expand( |
| 311 | patches.shape[0], -1, -1 |
| 312 | ) |
| 313 | out.append(class_token) |
| 314 | out.append(patches[:, current_length : current_length + seq_len, :]) |
| 315 | current_length += seq_len |
| 316 | |
| 317 | return torch.cat(out, dim=1) |
| 318 | |
| 319 | @property |
| 320 | def apply_cls_token(self): |
| 321 | return self.cls_token.enabled |
| 322 | |
| 323 | @property |
| 324 | def num_cls_tokens(self): |
| 325 | return self.cls_token.num_tokens |
| 326 | |
| 327 | @property |
| 328 | def num_cls_patches(self): |
| 329 | return self.cls_token.num_patches |
| 330 | |
| 331 | @property |
| 332 | def num_registers(self): |
| 333 | return self.cls_token.num_registers |
| 334 | |
| 335 | @property |
| 336 | def num_skip(self): |
| 337 | return self.num_cls_tokens + self.num_registers |
| 338 | |
| 339 | def embed_patches(self, x: torch.Tensor) -> torch.Tensor: |
| 340 | patches = self.im_to_patches(x) |
| 341 | patches = self.embedder(patches) |
| 342 | return patches |
| 343 | |
| 344 | def apply_pos_enc( |
| 345 | self, |
| 346 | patches: torch.Tensor, |
| 347 | patch_idxs: torch.Tensor | None = None, |
| 348 | input_size: tuple[int, int] | None = None, |
| 349 | ) -> torch.Tensor: |
| 350 | if not self.abs_pos: |
| 351 | return patches |
| 352 | |
| 353 | pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size) |
| 354 | |
| 355 | if self.training and self.pos_dropout > 0: |
| 356 | keeps = ( |
| 357 | torch.rand( |
| 358 | patches.shape[0], 1, 1, dtype=pos_enc.dtype, device=pos_enc.device |
| 359 | ) |
| 360 | > self.pos_dropout |
| 361 | ) |
| 362 | pos_enc_drop = torch.where(keeps, pos_enc, 0) |
| 363 | else: |
| 364 | pos_enc_drop = pos_enc |
| 365 | |
| 366 | return patches + pos_enc_drop, pos_enc |
| 367 | |
| 368 | def get_pos_enc( |
| 369 | self, |
| 370 | batch_size: int, |
| 371 | patch_idxs: torch.Tensor | None = None, |
| 372 | input_size: tuple[int, int] | None = None, |
| 373 | ) -> torch.Tensor: |
| 374 | if input_size is None: |
| 375 | input_dims = self.input_dims |
| 376 | else: |
| 377 | input_dims = ( |
| 378 | input_size[0] // self.patch_size, |
| 379 | input_size[1] // self.patch_size, |
| 380 | ) |
| 381 | |
| 382 | pos_embed = self._get_pos_embeddings(batch_size, input_dims) |
| 383 | |
| 384 | if patch_idxs is None: |
| 385 | return pos_embed |
| 386 | |
| 387 | exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1]) |
| 388 | |
| 389 | pos_embed = torch.gather( |
| 390 | pos_embed.expand(patch_idxs.shape[0], -1, -1), dim=1, index=exp_patch_idxs |
| 391 | ) |
| 392 | return pos_embed |
| 393 | |
| 394 | def _get_pos_embeddings(self, batch_size: int, input_dims: tuple[int, int]): |
| 395 | if (self.num_rows, self.num_cols) == input_dims: |
| 396 | return self.pos_embed |
| 397 | |
| 398 | pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute( |
| 399 | 0, 3, 1, 2 |
| 400 | ) |
| 401 | |
| 402 | def window_select(pos_embed): |
| 403 | if input_dims[0] < pos_embed.shape[-2]: |
| 404 | pos_embed = pos_embed[..., : input_dims[0], :] |
| 405 | if input_dims[1] < pos_embed.shape[-1]: |
| 406 | pos_embed = pos_embed[..., :, : input_dims[1]] |
| 407 | return pos_embed |
| 408 | |
| 409 | if self.cpe_mode: |
| 410 | max_dim = max(input_dims) |
| 411 | pos_embed = F.interpolate( |
| 412 | pos_embed.float(), |
| 413 | size=(max_dim, max_dim), |
| 414 | align_corners=False, |
| 415 | mode="bilinear", |
| 416 | ).to(pos_embed.dtype) |
| 417 | |
| 418 | pos_embed = window_select(pos_embed) |
| 419 | else: |
| 420 | pos_embed = window_select(pos_embed) |
| 421 | |
| 422 | if pos_embed.shape[-2:] != input_dims: |
| 423 | pos_embed = F.interpolate( |
| 424 | pos_embed.float(), size=input_dims, align_corners=False, mode="bilinear" |
| 425 | ).to(pos_embed.dtype) |
| 426 | |
| 427 | pos_embed = pos_embed.flatten(2).permute(0, 2, 1) |
| 428 | |
| 429 | return pos_embed |
| 430 | |
| 431 | |
| 432 | class Im2Patches(nn.Module): |
| 433 | def __init__(self, patch_size: int): |
| 434 | super().__init__() |
| 435 | self.patch_size = patch_size |
| 436 | |
| 437 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 438 | if self.patch_size == 1: |
| 439 | patches = x.flatten(2) |
| 440 | patches = patches.permute(0, 2, 1) |
| 441 | return patches |
| 442 | |
| 443 | py = x.shape[-2] // self.patch_size |
| 444 | px = x.shape[-1] // self.patch_size |
| 445 | patches = rearrange( |
| 446 | x, |
| 447 | "b c (py yy) (px xx) -> b (py px) (c yy xx)", |
| 448 | py=py, |
| 449 | yy=self.patch_size, |
| 450 | px=px, |
| 451 | xx=self.patch_size, |
| 452 | ) |
| 453 | return patches |
| 454 | |
| 455 | |
| 456 | class ViTPatchLinear(nn.Linear): |
| 457 | def __init__( |
| 458 | self, |
| 459 | patch_size: int, |
| 460 | embed_dim: int, |
| 461 | bias: bool = False, |
| 462 | temporal_patch_size: int = 1, |
| 463 | **factory, |
| 464 | ): |
| 465 | super().__init__( |
| 466 | 3 * temporal_patch_size * (patch_size**2), embed_dim, bias=bias, **factory |
| 467 | ) |
| 468 | self.patch_size = patch_size |
| 469 | self.temporal_patch_size = temporal_patch_size |
| 470 | |
| 471 | |
| 472 | @dataclass(frozen=True, kw_only=True) |
| 473 | class MaskMetadata: |
| 474 | cu_seqlens: torch.Tensor |
| 475 | max_seqlen: torch.Tensor |
| 476 | |
| 477 | |
| 478 | class RadioParallelAttention(InternParallelAttention): |
| 479 | def forward( |
| 480 | self, x: torch.Tensor, mask_meta: MaskMetadata | None = None |
| 481 | ) -> torch.Tensor: |
| 482 | qkv, _ = self.qkv(x) |
| 483 | q, k, v = qkv.chunk(3, dim=-1) |
| 484 | |
| 485 | if self.qk_normalization: |
| 486 | q, k = self._apply_qk_norm(q, k) |
| 487 | |
| 488 | cu_seqlens, max_seqlen = None, None |
| 489 | if mask_meta is not None: |
| 490 | cu_seqlens = mask_meta.cu_seqlens |
| 491 | max_seqlen = mask_meta.max_seqlen |
| 492 | out = self.attn(q, k, v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) |
| 493 | out, _ = self.proj(out) |
| 494 | return out |
| 495 | |
| 496 | |
| 497 | class RadioVisionEncoderLayer(InternVisionEncoderLayer): |
| 498 | def __init__(self, *args, **kwargs) -> None: |
| 499 | super().__init__(*args, attn_cls=RadioParallelAttention, **kwargs) |
| 500 | |
| 501 | def forward( |
| 502 | self, |
| 503 | hidden_states: torch.Tensor, |
| 504 | mask_meta: MaskMetadata | None = None, |
| 505 | ): |
| 506 | hidden_states = ( |
| 507 | hidden_states |
| 508 | + self.attn(self.norm1(hidden_states), mask_meta=mask_meta) * self.ls1 |
| 509 | ) |
| 510 | |
| 511 | hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) * self.ls2 |
| 512 | |
| 513 | return hidden_states |
| 514 | |
| 515 | |
| 516 | class RadioVisionEncoder(InternVisionEncoder): |
| 517 | def __init__(self, *args, **kwargs) -> None: |
| 518 | super().__init__(*args, layer_cls=RadioVisionEncoderLayer, **kwargs) |
| 519 | |
| 520 | def forward( |
| 521 | self, |
| 522 | inputs_embeds: torch.Tensor, |
| 523 | mask_meta: MaskMetadata | None = None, |
| 524 | ): |
| 525 | hidden_states = inputs_embeds |
| 526 | for encoder_layer in self.layers: |
| 527 | hidden_states = encoder_layer(hidden_states, mask_meta=mask_meta) |
| 528 | return hidden_states |
| 529 | |
| 530 | |
| 531 | class RadioInternVisionModel(nn.Module): |
| 532 | packed_modules_mapping = { |
| 533 | "qkv": ["qkv"], |
| 534 | } |
| 535 | |
| 536 | def __init__( |
| 537 | self, |
| 538 | config: PreTrainedConfig = None, |
| 539 | quant_config: QuantizationConfig | None = None, |
| 540 | *, |
| 541 | num_hidden_layers_override: int | None = None, |
| 542 | num_dummy_heads: int = 0, |
| 543 | prefix: str = "", |
| 544 | ) -> None: |
| 545 | super().__init__() |
| 546 | |
| 547 | self.config = config |
| 548 | self.img_size, self.grid_size, self.num_patches = self._init_img_size( |
| 549 | config.patch_size, config.image_size |
| 550 | ) |
| 551 | assert self.img_size is not None |
| 552 | max_img_size = int( |
| 553 | round(config.cpe_max_size / config.patch_size) * config.patch_size |
| 554 | ) |
| 555 | self.temporal_patch_size = config.video_temporal_patch_size |
| 556 | unique_teachers = set(t["name"] for t in config.teachers) |
| 557 | self.patch_generator = ViTPatchGenerator( |
| 558 | config.patch_size, |
| 559 | config.hidden_size, |
| 560 | input_dims=self.img_size, |
| 561 | max_input_dims=max_img_size, |
| 562 | cls_token=True, |
| 563 | num_cls_tokens=len(unique_teachers) if config.cls_token_per_teacher else 1, |
| 564 | register_multiple=config.register_multiple, |
| 565 | temporal_patch_size=self.temporal_patch_size, |
| 566 | separate_video_embedder=config.separate_video_embedder, |
| 567 | ) |
| 568 | |
| 569 | self.encoder = RadioVisionEncoder( |
| 570 | config=config, |
| 571 | quant_config=quant_config, |
| 572 | num_hidden_layers_override=num_hidden_layers_override, |
| 573 | num_dummy_heads=num_dummy_heads, |
| 574 | prefix=f"{prefix}.encoder", |
| 575 | ) |
| 576 | |
| 577 | def _init_img_size( |
| 578 | self, |
| 579 | patch_size: input_dim_t, |
| 580 | img_size: input_dim_t | None, |
| 581 | ) -> tuple[tuple[int, int] | None, tuple[int, int] | None, int | None]: |
| 582 | if img_size is None: |
| 583 | return None, None, None |
| 584 | if isinstance(img_size, int): |
| 585 | img_size = (img_size, img_size) |
| 586 | if isinstance(patch_size, int): |
| 587 | patch_size = (patch_size, patch_size) |
| 588 | grid_size = ( |
| 589 | img_size[0] // patch_size[0], |
| 590 | img_size[1] // patch_size[1], |
| 591 | ) |
| 592 | num_patches = grid_size[0] * grid_size[1] |
| 593 | return img_size, grid_size, num_patches |
| 594 | |
| 595 | def get_input_embeddings(self): |
| 596 | return self.embeddings |
| 597 | |
| 598 | def inter_image_mask_metadata( |
| 599 | self, imgs_sizes: list[tuple[int, int]], device: torch.device |
| 600 | ) -> MaskMetadata: |
| 601 | """Build mask metadata from image pixel sizes. Adds num_skip to each |
| 602 | sequence length (cls/register tokens) to match patch generator output.""" |
| 603 | patch_size = self.patch_generator.patch_size |
| 604 | num_skip = self.patch_generator.num_skip |
| 605 | |
| 606 | seq_lens = calc_seq_lens(imgs_sizes, patch_size) |
| 607 | adjusted = [s + num_skip for s in seq_lens] |
| 608 | return self._inter_image_mask_metadata_from_seq_lens(adjusted, device=device) |
| 609 | |
| 610 | def _inter_image_mask_metadata_from_seq_lens( |
| 611 | self, seq_lens: list[int], device: torch.device |
| 612 | ) -> MaskMetadata: |
| 613 | """Build mask metadata from actual sequence lengths (already including |
| 614 | cls/register tokens, i.e. patch_count + num_skip per item). |
| 615 | Use inter_image_mask_metadata() when you only have imgs_sizes.""" |
| 616 | assert len(seq_lens) > 0 |
| 617 | cu_seqlens = torch.tensor( |
| 618 | list(accumulate(seq_lens, initial=0)), dtype=torch.int32, device=device |
| 619 | ) |
| 620 | # Keep max_seqlen on CPU to avoid .item() sync |
| 621 | # See: https://github.com/vllm-project/vllm/blob/20b6b01/vllm/v1/attention/ops/vit_attn_wrappers.py#L48 |
| 622 | max_seqlen = torch.tensor(max(seq_lens), dtype=torch.int32) |
| 623 | return MaskMetadata(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) |
| 624 | |
| 625 | def forward( |
| 626 | self, |
| 627 | x: torch.Tensor, |
| 628 | imgs_sizes: list[tuple[int, int]] | None = None, |
| 629 | num_frames: int | None = None, |
| 630 | ) -> torch.FloatTensor: |
| 631 | T = self.temporal_patch_size |
| 632 | |
| 633 | # Build packed-sequence metadata for MMEncoderAttention when needed. |
| 634 | mask_meta = None |
| 635 | packed_batch_size = None # Original batch size before packing |
| 636 | |
| 637 | if num_frames is not None and T > 1: |
| 638 | # Conv3d video: all tubelets have the same sequence length. |
| 639 | # Pack [num_tubelets, seq_per_tubelet, hidden] → [1, total, hidden] |
| 640 | hidden_states = self.patch_generator.forward_video(x) |
| 641 | packed_batch_size, seq_per_tubelet, hidden_dim = hidden_states.shape |
| 642 | hidden_states = hidden_states.reshape(1, -1, hidden_dim) |
| 643 | mask_meta = self._inter_image_mask_metadata_from_seq_lens( |
| 644 | [seq_per_tubelet] * packed_batch_size, device=hidden_states.device |
| 645 | ) |
| 646 | else: |
| 647 | # Images for any model, or video for non-conv3d model |
| 648 | hidden_states = self.patch_generator(x, imgs_sizes=imgs_sizes) |
| 649 | if imgs_sizes is not None and len(imgs_sizes) > 1: |
| 650 | # Dynamic resolution w/ > 1 image, create attn mask |
| 651 | mask_meta = self.inter_image_mask_metadata( |
| 652 | imgs_sizes, device=hidden_states.device |
| 653 | ) |
| 654 | |
| 655 | encoder_outputs = self.encoder(inputs_embeds=hidden_states, mask_meta=mask_meta) |
| 656 | |
| 657 | # Unpack back to original batch shape if we packed for video |
| 658 | if packed_batch_size is not None: |
| 659 | encoder_outputs = encoder_outputs.reshape( |
| 660 | packed_batch_size, seq_per_tubelet, -1 |
| 661 | ) |
| 662 | |
| 663 | return encoder_outputs |
| 664 | |
| 665 | |
| 666 | class RadioModel(nn.Module): |
| 667 | packed_modules_mapping = { |
| 668 | "qkv": ["qkv"], |
| 669 | } |
| 670 | |
| 671 | def __init__( |
| 672 | self, |
| 673 | config: PreTrainedConfig, |
| 674 | quant_config: QuantizationConfig | None = None, |
| 675 | *, |
| 676 | num_hidden_layers_override: int | None = None, |
| 677 | num_dummy_heads: int = 0, |
| 678 | prefix: str = "", |
| 679 | ) -> None: |
| 680 | super().__init__() |
| 681 | |
| 682 | self.config = config |
| 683 | self.model = RadioInternVisionModel( |
| 684 | config=config, |
| 685 | quant_config=quant_config, |
| 686 | num_hidden_layers_override=num_hidden_layers_override, |
| 687 | num_dummy_heads=num_dummy_heads, |
| 688 | prefix=prefix, |
| 689 | ) |
| 690 | |
| 691 | summary_idxs = None |
| 692 | if config.teachers: |
| 693 | summary_idxs = torch.tensor( |
| 694 | [i for i, t in enumerate(config.teachers) if t.get("use_summary", True)] |
| 695 | ) |
| 696 | if summary_idxs.numel() > 0: |
| 697 | self.register_buffer("summary_idxs", summary_idxs) |
| 698 | self.summary_idxs = summary_idxs |
| 699 | |
| 700 | def forward( |
| 701 | self, |
| 702 | pixel_values: torch.Tensor | None = None, |
| 703 | pixel_embeds: torch.Tensor | None = None, |
| 704 | *, |
| 705 | imgs_sizes: list[tuple[int, int]] | None = None, |
| 706 | num_frames: int | None = None, |
| 707 | ) -> tuple[torch.FloatTensor, torch.FloatTensor]: |
| 708 | y = self.model( |
| 709 | pixel_values, |
| 710 | imgs_sizes=imgs_sizes, |
| 711 | num_frames=num_frames, |
| 712 | ) |
| 713 | return self._extract_final(y, imgs_sizes=imgs_sizes) |
| 714 | |
| 715 | def load_weights(self, weights) -> set[str]: |
| 716 | loaded_params: set[str] = set() |
| 717 | params_dict = dict(self.named_parameters()) |
| 718 | |
| 719 | if isinstance(weights, dict): |
| 720 | weights_list = list(weights.items()) |
| 721 | else: |
| 722 | weights_list = list(weights) |
| 723 | |
| 724 | for name, weight in weights_list: |
| 725 | if not name.startswith("radio_model."): |
| 726 | # Skip non-radio weights |
| 727 | continue |
| 728 | |
| 729 | sub = name[len("radio_model.") :] # drop "radio_model." prefix |
| 730 | |
| 731 | # Skip buffers not used in vLLM |
| 732 | if sub in {"summary_idxs"}: |
| 733 | continue |
| 734 | if sub.startswith("input_conditioner."): |
| 735 | # we normalize in the input processor, |
| 736 | # based on norm and std values from the config |
| 737 | continue |
| 738 | |
| 739 | vllm_key = None |
| 740 | if sub.startswith("model.patch_generator."): |
| 741 | vllm_key = f"model.patch_generator.{sub.split('.', 2)[-1]}" |
| 742 | elif sub.startswith("input_conditioner."): |
| 743 | vllm_key = f"input_conditioner.{sub.split('.', 1)[-1]}" |
| 744 | elif sub.startswith("model.blocks."): |
| 745 | # Encoder blocks: HF 'model.blocks.{i}.' -> |
| 746 | # vLLM 'model.encoder.layers.{i}.' |
| 747 | parts = sub.split(".") |
| 748 | if len(parts) >= 4: |
| 749 | layer_idx = parts[2] |
| 750 | suffix = ".".join(parts[3:]) |
| 751 | # Skip layer-scale entries that vLLM doesn't use |
| 752 | if suffix in {"ls1", "ls2"} or suffix.startswith(("ls1.", "ls2.")): |
| 753 | continue |
| 754 | vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}" |
| 755 | |
| 756 | if vllm_key and vllm_key in params_dict: |
| 757 | param = params_dict[vllm_key] |
| 758 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 759 | weight_loader(param, weight) |
| 760 | loaded_params.add(vllm_key) |
| 761 | |
| 762 | return loaded_params |
| 763 | |
| 764 | def _extract_final( |
| 765 | self, y: torch.Tensor, imgs_sizes: list[tuple[int, int]] | None = None |
| 766 | ) -> tuple[torch.FloatTensor, torch.FloatTensor]: |
| 767 | # Remove CLS + REGISTERS tokens |
| 768 | num_skip = self.model.patch_generator.num_skip |
| 769 | patch_size = self.model.patch_generator.patch_size |
| 770 | num_cls_tokens = self.model.patch_generator.num_cls_tokens |
| 771 | if imgs_sizes is None: |
| 772 | all_summary = y[:, :num_cls_tokens] |
| 773 | all_feat = y[:, num_skip:] |
| 774 | else: |
| 775 | all_patches = [] |
| 776 | summaries = [] |
| 777 | current_pos = 0 |
| 778 | for num_patches in calc_seq_lens(imgs_sizes, patch_size): |
| 779 | patches = y[ |
| 780 | :, current_pos + num_skip : current_pos + num_skip + num_patches, : |
| 781 | ] |
| 782 | all_patches.append(patches) |
| 783 | summary = y[:, current_pos : current_pos + num_cls_tokens, :] |
| 784 | summaries.append(summary) |
| 785 | current_pos += num_skip + num_patches |
| 786 | all_summary = torch.cat(summaries, dim=1) |
| 787 | all_feat = torch.cat(all_patches, dim=1) |
| 788 | |
| 789 | if self.summary_idxs is not None: |
| 790 | bb_summary = all_summary[:, self.summary_idxs] |
| 791 | else: |
| 792 | bb_summary = all_summary |
| 793 | return bb_summary.flatten(1), all_feat |