| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # -------------------------------------------------------- |
| 4 | # Adapted from |
| 5 | # https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/internvl.py |
| 6 | # under Apache-2.0 License |
| 7 | # LICENSE is in root directory. |
| 8 | # -------------------------------------------------------- |
| 9 | |
| 10 | import math |
| 11 | from abc import ABC, abstractmethod |
| 12 | from collections.abc import Sequence |
| 13 | from dataclasses import dataclass |
| 14 | from functools import cached_property |
| 15 | from typing import Any, TypeVar |
| 16 | |
| 17 | import einops |
| 18 | import numpy as np |
| 19 | import numpy.typing as npt |
| 20 | import regex as re |
| 21 | import torch |
| 22 | from PIL import Image |
| 23 | from transformers import BatchFeature, PreTrainedConfig, TensorType |
| 24 | |
| 25 | from vllm.model_executor.models.parakeet import ParakeetExtractor |
| 26 | from vllm.multimodal.inputs import AudioItem |
| 27 | from vllm.multimodal.processing.processor import PromptUpdateDetails, cached_encode |
| 28 | from vllm.multimodal.video_prune.evs import compute_retained_tokens_count |
| 29 | from vllm.platforms import current_platform |
| 30 | from vllm.tokenizers.hf import HfTokenizer |
| 31 | |
| 32 | from .internvl import calculate_internvl_targets, get_internvl_target_ratios |
| 33 | |
| 34 | _T = TypeVar("_T") |
| 35 | |
| 36 | |
| 37 | IMG_START = "<img>" |
| 38 | IMG_END = "</img>" |
| 39 | IMG_CONTEXT = "<image>" |
| 40 | AUDIO_START = "<so_start>" |
| 41 | AUDIO_END = "<so_end>" |
| 42 | AUDIO_CONTEXT = "<so_embedding>" |
| 43 | |
| 44 | # Profiling |
| 45 | # MAX_FRAMES = 16 |
| 46 | DEFAULT_NUM_TILES = 12 |
| 47 | |
| 48 | |
| 49 | def calculate_timestamps( |
| 50 | indices: list[int] | torch.Tensor, |
| 51 | frame_duration_ms: int, |
| 52 | ): |
| 53 | if not isinstance(indices, list): |
| 54 | indices = indices.tolist() |
| 55 | |
| 56 | timestamps = [int(i) * frame_duration_ms / 1000.0 for i in indices] |
| 57 | return timestamps |
| 58 | |
| 59 | |
| 60 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 61 | def _bicubic_resize_and_normalize( |
| 62 | tensor: torch.Tensor, |
| 63 | size: tuple[int, int] | None = None, |
| 64 | norm_mean: torch.Tensor | None = None, |
| 65 | norm_std: torch.Tensor | None = None, |
| 66 | dtype: torch.dtype = torch.float32, |
| 67 | ) -> torch.Tensor: |
| 68 | """Permute NHWC→NCHW, optional bicubic resize, rescale + normalize. |
| 69 | |
| 70 | Input must be a raw 4-D **NHWC** tensor. |
| 71 | |
| 72 | *size*: target ``(H, W)``; skips interpolation when ``None``. |
| 73 | *norm_mean* / *norm_std*: when both provided, fused |
| 74 | ``(x/255 - mean) / std`` + dtype cast; otherwise ``x/255`` + cast. |
| 75 | """ |
| 76 | tensor = tensor.permute(0, 3, 1, 2).to(dtype=torch.float32) |
| 77 | if size is not None: |
| 78 | tensor = torch.nn.functional.interpolate( |
| 79 | tensor, size=size, mode="bicubic", align_corners=False, antialias=True |
| 80 | ) |
| 81 | if norm_mean is not None and norm_std is not None: |
| 82 | return ((tensor / 255.0 - norm_mean) / norm_std).to(dtype=dtype).contiguous() |
| 83 | return (tensor / 255.0).to(dtype=dtype).contiguous() |
| 84 | |
| 85 | |
| 86 | def _pil_to_nhwc_tensor(image: Image.Image) -> torch.Tensor: |
| 87 | """Convert a PIL image to a 4-D NHWC tensor suitable for compiled ops.""" |
| 88 | array = np.asarray( |
| 89 | image.convert("RGB") if image.mode != "RGB" else image, dtype=np.uint8 |
| 90 | ) |
| 91 | return torch.from_numpy(np.expand_dims(array, axis=0)) |
| 92 | |
| 93 | |
| 94 | def dynamic_preprocess( |
| 95 | image: Image.Image, |
| 96 | *, |
| 97 | image_size: int = 512, |
| 98 | max_num_tiles: int = 12, |
| 99 | use_thumbnail: bool = True, |
| 100 | norm_mean: torch.Tensor | None = None, |
| 101 | norm_std: torch.Tensor | None = None, |
| 102 | dtype: torch.dtype = torch.float32, |
| 103 | ) -> torch.Tensor: |
| 104 | orig_width, orig_height = image.size |
| 105 | |
| 106 | target_ratios = get_internvl_target_ratios(1, max_num_tiles) |
| 107 | |
| 108 | blocks, target_width, target_height = calculate_internvl_targets( |
| 109 | orig_width=orig_width, |
| 110 | orig_height=orig_height, |
| 111 | target_ratios=target_ratios, |
| 112 | image_size=image_size, |
| 113 | use_thumbnail=False, |
| 114 | ) |
| 115 | |
| 116 | tensor = _pil_to_nhwc_tensor(image) |
| 117 | |
| 118 | resized_img = _bicubic_resize_and_normalize( |
| 119 | tensor, |
| 120 | size=(target_height, target_width), |
| 121 | norm_mean=norm_mean, |
| 122 | norm_std=norm_std, |
| 123 | dtype=dtype, |
| 124 | ) |
| 125 | B, C, H, W = resized_img.shape |
| 126 | hp, wp = H // image_size, W // image_size |
| 127 | patches = ( |
| 128 | resized_img.reshape(B, C, hp, image_size, wp, image_size) |
| 129 | .permute(0, 2, 4, 1, 3, 5) |
| 130 | .reshape(B * hp * wp, C, image_size, image_size) |
| 131 | ) |
| 132 | |
| 133 | if use_thumbnail and patches.shape[0] > 1: |
| 134 | thumb = _bicubic_resize_and_normalize( |
| 135 | tensor, |
| 136 | size=(image_size, image_size), |
| 137 | norm_mean=norm_mean, |
| 138 | norm_std=norm_std, |
| 139 | dtype=dtype, |
| 140 | ) |
| 141 | patches = torch.cat([patches, thumb], dim=0) |
| 142 | |
| 143 | return patches |
| 144 | |
| 145 | |
| 146 | def _compute_aspect_preserving_size( |
| 147 | orig_w: int, |
| 148 | orig_h: int, |
| 149 | target_num_patches: int, |
| 150 | patch_size: int, |
| 151 | downsample_ratio: float, |
| 152 | ) -> tuple[int, int]: |
| 153 | """Compute target pixel dimensions that preserve aspect ratio. |
| 154 | |
| 155 | Mirrors Megatron-LM image_processing.py video frame resizing: |
| 156 | target area in patch-grid space is *target_num_patches*, distributed |
| 157 | according to the source aspect ratio, then snapped to a multiple of |
| 158 | the required divisor (2 for pixel-shuffle). |
| 159 | """ |
| 160 | aspect_wh = orig_w / max(orig_h, 1) |
| 161 | ph = round(math.sqrt(target_num_patches / aspect_wh)) |
| 162 | pw = round(math.sqrt(target_num_patches * aspect_wh)) |
| 163 | ph = max(ph, 1) |
| 164 | pw = max(pw, 1) |
| 165 | |
| 166 | reduction_factor = int(round(1 / downsample_ratio)) |
| 167 | required_divisor = reduction_factor # 2 for pixel-shuffle |
| 168 | if required_divisor > 1: |
| 169 | rem_h = ph % required_divisor |
| 170 | rem_w = pw % required_divisor |
| 171 | ph_up = ph + (required_divisor - rem_h if rem_h else 0) |
| 172 | ph_down = ph - rem_h |
| 173 | pw_up = pw + (required_divisor - rem_w if rem_w else 0) |
| 174 | pw_down = pw - rem_w |
| 175 | if ph_up * pw_up <= target_num_patches: |
| 176 | ph, pw = ph_up, pw_up |
| 177 | else: |
| 178 | ph = max(required_divisor, ph_down) |
| 179 | pw = max(required_divisor, pw_down) |
| 180 | |
| 181 | return pw * patch_size, ph * patch_size # (width, height) in pixels |
| 182 | |
| 183 | |
| 184 | def get_video_target_size_and_feature_size( |
| 185 | orig_w: int, |
| 186 | orig_h: int, |
| 187 | target_patches: int, |
| 188 | maintain_aspect_ratio: bool, |
| 189 | patch_size: int, |
| 190 | downsample_ratio: float, |
| 191 | ) -> tuple[int, int, int]: |
| 192 | """Compute target (width, height) and feature_size for video resize and token count. |
| 193 | |
| 194 | Used by video_to_pixel_values (resize) and get_video_replacement_internvl |
| 195 | (seq length calc) so both use the same dimensions. |
| 196 | """ |
| 197 | if maintain_aspect_ratio: |
| 198 | target_w, target_h = _compute_aspect_preserving_size( |
| 199 | orig_w=orig_w, |
| 200 | orig_h=orig_h, |
| 201 | target_num_patches=target_patches, |
| 202 | patch_size=patch_size, |
| 203 | downsample_ratio=downsample_ratio, |
| 204 | ) |
| 205 | else: |
| 206 | reduction_factor = int(round(1 / downsample_ratio)) |
| 207 | side = int(math.sqrt(target_patches)) |
| 208 | side = max(reduction_factor, (side // reduction_factor) * reduction_factor) |
| 209 | target_w = side * patch_size |
| 210 | target_h = side * patch_size |
| 211 | |
| 212 | feature_size = int((target_h // patch_size) * downsample_ratio) * int( |
| 213 | (target_w // patch_size) * downsample_ratio |
| 214 | ) |
| 215 | return target_w, target_h, feature_size |
| 216 | |
| 217 | |
| 218 | def video_to_pixel_values( |
| 219 | video: npt.NDArray, |
| 220 | *, |
| 221 | input_size: int, |
| 222 | video_target_num_patches: int | None = None, |
| 223 | video_maintain_aspect_ratio: bool = False, |
| 224 | patch_size: int = 16, |
| 225 | downsample_ratio: float = 0.5, |
| 226 | norm_mean: torch.Tensor | None = None, |
| 227 | norm_std: torch.Tensor | None = None, |
| 228 | dtype: torch.dtype = torch.float32, |
| 229 | ) -> torch.Tensor: |
| 230 | """Convert video ndarray (T, H, W, C) to normalized pixel tensor (T, C, H, W).""" |
| 231 | orig_h, orig_w = video.shape[1], video.shape[2] |
| 232 | size: tuple[int, int] | None = None |
| 233 | |
| 234 | if video_target_num_patches is not None: |
| 235 | tw, th, _ = get_video_target_size_and_feature_size( |
| 236 | orig_w=orig_w, |
| 237 | orig_h=orig_h, |
| 238 | target_patches=video_target_num_patches, |
| 239 | maintain_aspect_ratio=video_maintain_aspect_ratio, |
| 240 | patch_size=patch_size, |
| 241 | downsample_ratio=downsample_ratio, |
| 242 | ) |
| 243 | if orig_h != th or orig_w != tw: |
| 244 | size = (th, tw) |
| 245 | elif orig_h != input_size or orig_w != input_size: |
| 246 | size = (input_size, input_size) |
| 247 | |
| 248 | tensor = torch.from_numpy(video) |
| 249 | return _bicubic_resize_and_normalize(tensor, size, norm_mean, norm_std, dtype) |
| 250 | |
| 251 | |
| 252 | class DynamicResolutionImageTiler: |
| 253 | CONV_MERGING = False |
| 254 | PIXEL_SHUFFLE = True |
| 255 | USE_THUMBNAIL = False |
| 256 | |
| 257 | def __init__( |
| 258 | self, |
| 259 | *, |
| 260 | max_model_len: int, |
| 261 | patch_size: int, |
| 262 | min_num_patches: int, |
| 263 | max_num_patches: int, |
| 264 | downsample_ratio: int, |
| 265 | norm_mean: Sequence[float], |
| 266 | norm_std: Sequence[float], |
| 267 | factor_max: float = 1.0, |
| 268 | use_thumbnail: bool = False, |
| 269 | ) -> None: |
| 270 | assert use_thumbnail is False, "use_thumbnail is not supported" |
| 271 | self._patch_size: int = patch_size |
| 272 | self._max_model_len = max_model_len |
| 273 | self._min_num_patches = min_num_patches |
| 274 | self._max_num_patches = max_num_patches if max_num_patches > 0 else float("inf") |
| 275 | self._factor_max = factor_max |
| 276 | self.norm_mean = torch.tensor(norm_mean).reshape(3, 1, 1) |
| 277 | self.norm_std = torch.tensor(norm_std).reshape(3, 1, 1) |
| 278 | assert downsample_ratio < 1 |
| 279 | reduction_factor = 1 / downsample_ratio |
| 280 | assert reduction_factor == 2.0 |
| 281 | self._downsample_ratio = int(reduction_factor) ** ( |
| 282 | self.PIXEL_SHUFFLE + self.CONV_MERGING |
| 283 | ) |
| 284 | assert self._downsample_ratio == 2 |
| 285 | |
| 286 | def _get_num_embeddings(self, width: int, height: int) -> int: |
| 287 | num_patches = (width // self._patch_size) * (height // self._patch_size) |
| 288 | num_tokens = num_patches // (self._downsample_ratio**2) |
| 289 | return num_tokens |
| 290 | |
| 291 | def width_and_height_for_max_num_tokens_available( |
| 292 | self, |
| 293 | target_num_tokens_post_shuffle: int, |
| 294 | ) -> tuple[int, int]: |
| 295 | """TODO: optimize this so it squeezes closer to target number of tokens. |
| 296 | Calculate image dimensions that produce approximately `target` tokens after |
| 297 | pixel_shuffle. |
| 298 | |
| 299 | With pixel_shuffle enabled, each 2x2 patch grid becomes 1 token, so we |
| 300 | need 4*B patches to get B tokens. |
| 301 | |
| 302 | Examples: |
| 303 | >>> PATCH_SIZE = 16 |
| 304 | >>> DOWNSAMPLE_RATIO = 0.5 |
| 305 | >>> tiler = DynamicResolutionImageTiler( |
| 306 | ... max_model_len=16384, |
| 307 | ... patch_size=PATCH_SIZE, |
| 308 | ... downsample_ratio=DOWNSAMPLE_RATIO, |
| 309 | ... min_num_patches=4, |
| 310 | ... max_num_patches=0, |
| 311 | ... ) |
| 312 | >>> width, height = tiler.width_and_height_for_max_num_tokens_available( |
| 313 | ... target_num_tokens_post_shuffle=8192, |
| 314 | ... ) |
| 315 | >>> assert width, height == (2880, 2880) |
| 316 | >>> assert (width // PATCH_SIZE) * ( |
| 317 | ... height // PATCH_SIZE |
| 318 | ... ) // 2**2 == 8100 # tokens post-shuffle |
| 319 | >>> assert tiler._get_num_embeddings(width=width, height=height) == 8100 |
| 320 | |
| 321 | """ |
| 322 | side_pixels = ( |
| 323 | math.isqrt(target_num_tokens_post_shuffle) |
| 324 | * self._downsample_ratio |
| 325 | * self._patch_size |
| 326 | ) |
| 327 | assert isinstance(side_pixels, int) and side_pixels % self._patch_size == 0 |
| 328 | return side_pixels, side_pixels |
| 329 | |
| 330 | def max_num_tokens_available(self, text_prompt_length: int) -> int: |
| 331 | return self._max_model_len - text_prompt_length - 4 |
| 332 | |
| 333 | def _images_to_pixel_values_lst( |
| 334 | self, |
| 335 | text_prompt_length: int, |
| 336 | images: list[Image.Image], |
| 337 | dtype: torch.dtype = torch.float32, |
| 338 | ) -> tuple[list[torch.Tensor], list[int]]: |
| 339 | num_tokens_available = self.max_num_tokens_available(text_prompt_length) |
| 340 | params_per_image = self.compute_params(images, num_tokens_available) |
| 341 | |
| 342 | feature_sizes = [] |
| 343 | images = [] |
| 344 | for param in params_per_image: |
| 345 | for t in self.apply_params(param, dtype=dtype): |
| 346 | assert t.ndim == 3, f"{t.ndim=}: expected 3 dim tensor" |
| 347 | images.append(t) |
| 348 | feature_sizes.append(param.num_embeddings) |
| 349 | return images, feature_sizes |
| 350 | |
| 351 | @dataclass |
| 352 | class DynamicResolutionParams: |
| 353 | media: Image.Image |
| 354 | num_tiles: int |
| 355 | num_embeddings: int |
| 356 | patch_size: tuple[int, int] |
| 357 | |
| 358 | def apply_params( |
| 359 | self, |
| 360 | params: DynamicResolutionParams, |
| 361 | dtype: torch.dtype = torch.float32, |
| 362 | ) -> list[torch.Tensor]: |
| 363 | target_size = ( |
| 364 | params.patch_size[1] * self._patch_size, |
| 365 | params.patch_size[0] * self._patch_size, |
| 366 | ) |
| 367 | tensor = _pil_to_nhwc_tensor(params.media) |
| 368 | resized_img = _bicubic_resize_and_normalize( |
| 369 | tensor, |
| 370 | size=target_size, |
| 371 | norm_mean=self.norm_mean, |
| 372 | norm_std=self.norm_std, |
| 373 | dtype=dtype, |
| 374 | ) |
| 375 | return list(resized_img) |
| 376 | |
| 377 | def process_media( |
| 378 | self, |
| 379 | media: Image.Image, |
| 380 | num_tokens_available: int, |
| 381 | ) -> tuple[DynamicResolutionParams, int]: |
| 382 | """Process a single media item and return its parameters. |
| 383 | |
| 384 | Args: |
| 385 | media: The media item to process |
| 386 | num_tokens_available: Number of tokens available for this media |
| 387 | Returns: |
| 388 | DynamicResolutionParams for the media |
| 389 | |
| 390 | """ |
| 391 | current_num_tokens_available = num_tokens_available |
| 392 | assert isinstance(media, Image.Image), ( |
| 393 | "Dynamic resolution is only supported for image media" |
| 394 | ) |
| 395 | orig_width, orig_height = media.width, media.height |
| 396 | closest_patch_height = round(orig_height / self._patch_size + 0.5) |
| 397 | closest_patch_width = round(orig_width / self._patch_size + 0.5) |
| 398 | patches = closest_patch_height * closest_patch_width |
| 399 | |
| 400 | factor = min( |
| 401 | math.sqrt(current_num_tokens_available / patches), self._factor_max |
| 402 | ) |
| 403 | target_patch_height = math.floor(factor * closest_patch_height) |
| 404 | target_patch_width = math.floor(factor * closest_patch_width) |
| 405 | |
| 406 | # Consider self._min_num_patches if > current_num_tokens_available. |
| 407 | if ( |
| 408 | current_num_tokens_available > self._min_num_patches |
| 409 | and target_patch_height * target_patch_width < self._min_num_patches |
| 410 | ): |
| 411 | up_factor = math.sqrt( |
| 412 | self._min_num_patches / (target_patch_height * target_patch_width) |
| 413 | ) |
| 414 | target_patch_height = math.ceil(up_factor * target_patch_height) |
| 415 | target_patch_width = math.ceil(up_factor * target_patch_width) |
| 416 | |
| 417 | # Round patch grid to be divisible by 2 (pixel-shuffle OR conv-merging) |
| 418 | # or by 4 when BOTH are enabled (two successive 2x reductions) |
| 419 | if self.PIXEL_SHUFFLE or self.CONV_MERGING: |
| 420 | required_divisor = 4 if (self.PIXEL_SHUFFLE and self.CONV_MERGING) else 2 |
| 421 | |
| 422 | rem_h = target_patch_height % required_divisor |
| 423 | if rem_h != 0: |
| 424 | inc_h = required_divisor - rem_h |
| 425 | if ( |
| 426 | target_patch_height + inc_h |
| 427 | ) * target_patch_width <= current_num_tokens_available: |
| 428 | target_patch_height += inc_h |
| 429 | else: |
| 430 | target_patch_height = max( |
| 431 | required_divisor, target_patch_height - rem_h |
| 432 | ) |
| 433 | |
| 434 | rem_w = target_patch_width % required_divisor |
| 435 | if rem_w != 0: |
| 436 | inc_w = required_divisor - rem_w |
| 437 | if ( |
| 438 | target_patch_height * (target_patch_width + inc_w) |
| 439 | <= current_num_tokens_available |
| 440 | ): |
| 441 | target_patch_width += inc_w |
| 442 | else: |
| 443 | target_patch_width = max( |
| 444 | required_divisor, target_patch_width - rem_w |
| 445 | ) |
| 446 | |
| 447 | # Calculate embeddings for the main dynamic resolution image |
| 448 | num_embeddings = self._get_num_embeddings( |
| 449 | target_patch_width * self._patch_size, |
| 450 | target_patch_height * self._patch_size, |
| 451 | ) |
| 452 | |
| 453 | token_count = target_patch_width * target_patch_height |
| 454 | |
| 455 | # Add thumbnail embeddings if enabled and image area is below threshold |
| 456 | num_tiles = 1 # Base dynamic resolution image |
| 457 | |
| 458 | return self.DynamicResolutionParams( |
| 459 | media=media, |
| 460 | num_tiles=num_tiles, |
| 461 | num_embeddings=num_embeddings, |
| 462 | patch_size=(target_patch_width, target_patch_height), |
| 463 | ), token_count |
| 464 | |
| 465 | def compute_params( |
| 466 | self, |
| 467 | media_list: list[Image.Image], |
| 468 | num_tokens_available: int, |
| 469 | ) -> list[DynamicResolutionParams]: |
| 470 | """Compute parameters for all media with iterative token budgeting. |
| 471 | |
| 472 | Args: |
| 473 | media_list: List of media items to process |
| 474 | num_tokens_available: Total number of tokens available across all media |
| 475 | Returns: |
| 476 | List of ImageTilingParams for each media item |
| 477 | |
| 478 | """ |
| 479 | num_tokens_available = ( |
| 480 | num_tokens_available |
| 481 | * (4 if self.PIXEL_SHUFFLE else 1) |
| 482 | * (4 if self.CONV_MERGING else 1) |
| 483 | ) |
| 484 | # When the number of available token is too small, |
| 485 | # allow self._min_num_patches per media and let the sample be truncated. |
| 486 | num_tokens_available = max( |
| 487 | num_tokens_available, self._min_num_patches * len(media_list) |
| 488 | ) |
| 489 | |
| 490 | # Clip the number of tokens available per media to >min and <max patches. |
| 491 | num_tokens_available_per_media = [ |
| 492 | int( |
| 493 | max( |
| 494 | min(num_tokens_available, self._max_num_patches), |
| 495 | self._min_num_patches, |
| 496 | ) |
| 497 | ) |
| 498 | for _ in range(len(media_list)) |
| 499 | ] |
| 500 | |
| 501 | # prevent infinite loop in any case |
| 502 | for _ in range(10): |
| 503 | # Step 1: Process each media with current token budget |
| 504 | params = [] |
| 505 | token_counts = [] |
| 506 | |
| 507 | for media, tokens_for_media in zip( |
| 508 | media_list, num_tokens_available_per_media |
| 509 | ): |
| 510 | param, token_count = self.process_media(media, tokens_for_media) |
| 511 | params.append(param) |
| 512 | token_counts.append(token_count) |
| 513 | |
| 514 | # Step 2: Check if total tokens is within budget |
| 515 | total_tokens = sum(token_counts) |
| 516 | |
| 517 | if total_tokens <= num_tokens_available: |
| 518 | # We're within budget, return the params |
| 519 | return params |
| 520 | |
| 521 | # Step 3: We're over budget, need to scale down |
| 522 | # Calculate scaling factor to get under budget |
| 523 | scaling_factor = num_tokens_available / total_tokens |
| 524 | |
| 525 | # Recalculate token budgets for each media based on scaling |
| 526 | # Each media gets a proportional share of the total budget |
| 527 | scaled_down_num_tokens_available_per_media = [ |
| 528 | max(self._min_num_patches, int(token_count * scaling_factor)) |
| 529 | for token_count in token_counts |
| 530 | ] |
| 531 | scaled_down = any( |
| 532 | [ |
| 533 | scaled_down_num_tokens_available_per_media[i] |
| 534 | < num_tokens_available_per_media[i] |
| 535 | for i in range(len(num_tokens_available_per_media)) |
| 536 | ] |
| 537 | ) |
| 538 | # If there wasn't scaling down, we're stuck with min_num_patches per media, |
| 539 | # else try with the scaled down num_tokens_available_per_media. |
| 540 | if not scaled_down: |
| 541 | num_tokens_available_per_media = [self._min_num_patches] * len( |
| 542 | media_list |
| 543 | ) |
| 544 | else: |
| 545 | num_tokens_available_per_media = ( |
| 546 | scaled_down_num_tokens_available_per_media |
| 547 | ) |
| 548 | ctx = f"{params=} {total_tokens=} {num_tokens_available=}" |
| 549 | raise ValueError( |
| 550 | f"Should be unreachable - `return params` above must be reached: {ctx}" |
| 551 | ) |
| 552 | |
| 553 | @staticmethod |
| 554 | def stack( |
| 555 | images: list[torch.Tensor] | torch.Tensor, patch_size: int |
| 556 | ) -> torch.Tensor: |
| 557 | assert len(images) > 0, "No images to stack" |
| 558 | |
| 559 | def rearrange_img(x): |
| 560 | py = x.shape[-2] // patch_size |
| 561 | px = x.shape[-1] // patch_size |
| 562 | x = einops.rearrange( |
| 563 | x, |
| 564 | "c (py yy) (px xx) -> (py px) (c yy xx)", |
| 565 | py=py, |
| 566 | yy=patch_size, |
| 567 | px=px, |
| 568 | xx=patch_size, |
| 569 | ) |
| 570 | return x |
| 571 | |
| 572 | imgs = [rearrange_img(img) for img in images] |
| 573 | pixel_values_flat = torch.cat(imgs, dim=0).unsqueeze(0) |
| 574 | return pixel_values_flat |
| 575 | |
| 576 | |
| 577 | class BaseNanoNemotronVLProcessor(ABC): |
| 578 | """This model doesn't define its own HF processor, |
| 579 | so we implement our own one here. |
| 580 | |
| 581 | The code to insert image tokens is based on: |
| 582 | https://huggingface.co/OpenGVLab/InternVL2-1B/blob/main/modeling_internvl_chat.py#L252 |
| 583 | """ |
| 584 | |
| 585 | def __init__( |
| 586 | self, |
| 587 | config: PreTrainedConfig, |
| 588 | tokenizer: HfTokenizer, |
| 589 | *args, |
| 590 | max_model_len: int, |
| 591 | max_num_tiles: int | None = None, |
| 592 | **kwargs, |
| 593 | ) -> None: |
| 594 | super().__init__() |
| 595 | |
| 596 | self.config = config |
| 597 | self.tokenizer = tokenizer |
| 598 | |
| 599 | self.max_num_tiles = max_num_tiles or DEFAULT_NUM_TILES |
| 600 | image_size: int = config.force_image_size |
| 601 | patch_size: int = config.patch_size |
| 602 | downsample_ratio: int = config.downsample_ratio |
| 603 | |
| 604 | self.num_image_token = int( |
| 605 | (image_size // patch_size) ** 2 * (downsample_ratio**2) |
| 606 | ) |
| 607 | self.image_size = image_size |
| 608 | self.use_thumbnail: bool = config.use_thumbnail |
| 609 | self.norm_mean = torch.Tensor(config.norm_mean).reshape(1, 3, 1, 1) |
| 610 | self.norm_std = torch.Tensor(config.norm_std).reshape(1, 3, 1, 1) |
| 611 | |
| 612 | self.dynamic_tiler: DynamicResolutionImageTiler | None = None |
| 613 | if self.use_dynamic_resolution(config): |
| 614 | self.dynamic_tiler = DynamicResolutionImageTiler( |
| 615 | max_model_len=max_model_len, |
| 616 | patch_size=patch_size, |
| 617 | downsample_ratio=downsample_ratio, |
| 618 | min_num_patches=config.vision_config.args["min_num_patches"], |
| 619 | max_num_patches=config.vision_config.args["max_num_patches"], |
| 620 | norm_mean=config.norm_mean, |
| 621 | norm_std=config.norm_std, |
| 622 | ) |
| 623 | self.dtype: torch.dtype = getattr(config, "dtype", torch.float32) |
| 624 | |
| 625 | @staticmethod |
| 626 | def use_dynamic_resolution(config: PreTrainedConfig) -> bool: |
| 627 | return "min_num_patches" in config.vision_config.args |
| 628 | |
| 629 | @property |
| 630 | @abstractmethod |
| 631 | def image_token_id(self) -> int: |
| 632 | raise NotImplementedError |
| 633 | |
| 634 | @abstractmethod |
| 635 | def get_image_repl( |
| 636 | self, |
| 637 | feature_size: int, |
| 638 | num_patches: int | None, |
| 639 | ) -> PromptUpdateDetails: |
| 640 | raise NotImplementedError |
| 641 | |
| 642 | def get_num_image_tokens( |
| 643 | self, |
| 644 | *, |
| 645 | image_width: int, |
| 646 | image_height: int, |
| 647 | max_num_tiles: int, |
| 648 | ) -> int: |
| 649 | target_ratios = get_internvl_target_ratios(1, max_num_tiles) |
| 650 | |
| 651 | num_patches, _, _ = calculate_internvl_targets( |
| 652 | orig_width=image_width, |
| 653 | orig_height=image_height, |
| 654 | target_ratios=target_ratios, |
| 655 | image_size=self.image_size, |
| 656 | use_thumbnail=self.use_thumbnail, |
| 657 | ) |
| 658 | |
| 659 | return num_patches * self.num_image_token |
| 660 | |
| 661 | def _images_to_pixel_values_lst( |
| 662 | self, |
| 663 | images: list[Image.Image], |
| 664 | max_num_tiles: int, |
| 665 | ) -> list[torch.Tensor]: |
| 666 | return [ |
| 667 | dynamic_preprocess( |
| 668 | image, |
| 669 | image_size=self.image_size, |
| 670 | max_num_tiles=max_num_tiles, |
| 671 | use_thumbnail=self.use_thumbnail, |
| 672 | norm_mean=self.norm_mean, |
| 673 | norm_std=self.norm_std, |
| 674 | dtype=self.dtype, |
| 675 | ) |
| 676 | for image in images |
| 677 | ] |
| 678 | |
| 679 | def _preprocess_image( |
| 680 | self, |
| 681 | text: list[str], |
| 682 | images: list[Image.Image], |
| 683 | max_num_tiles: int, |
| 684 | ) -> tuple[list[str], dict[str, Any]]: |
| 685 | if len(images) == 0: |
| 686 | return text, {} |
| 687 | |
| 688 | image_inputs: dict[str, Any] |
| 689 | if tiler := self.dynamic_tiler: |
| 690 | sans_images = text[0].replace("<image>", "") |
| 691 | text_prompt_length = len( |
| 692 | self.tokenizer(sans_images, add_special_tokens=False).input_ids |
| 693 | ) |
| 694 | pixel_values_lst, num_tokens_per_image = tiler._images_to_pixel_values_lst( |
| 695 | text_prompt_length=text_prompt_length, |
| 696 | images=images, |
| 697 | dtype=self.dtype, |
| 698 | ) |
| 699 | imgs_sizes = [(pv.shape[-2], pv.shape[-1]) for pv in pixel_values_lst] |
| 700 | image_num_patches = torch.tensor([1] * len(num_tokens_per_image)) |
| 701 | image_inputs = { |
| 702 | "pixel_values_flat": pixel_values_lst, |
| 703 | "imgs_sizes": imgs_sizes, |
| 704 | "num_tokens_per_image": num_tokens_per_image, |
| 705 | } |
| 706 | else: |
| 707 | pixel_values_lst = self._images_to_pixel_values_lst(images, max_num_tiles) |
| 708 | image_num_patches = torch.tensor([len(item) for item in pixel_values_lst]) |
| 709 | pixel_values_flat = ( |
| 710 | torch.cat(pixel_values_lst) |
| 711 | if len(pixel_values_lst) > 1 |
| 712 | else pixel_values_lst[0] |
| 713 | ) |
| 714 | image_inputs = { |
| 715 | "pixel_values_flat": pixel_values_flat, |
| 716 | "image_num_patches": image_num_patches, |
| 717 | } |
| 718 | num_tokens_per_image = [ |
| 719 | self.num_image_token * len(item) for item in pixel_values_lst |
| 720 | ] |
| 721 | |
| 722 | assert len(text) == 1, ( |
| 723 | "hf_processor is called on the output of get_dummy_text, " |
| 724 | "which should be a single string" |
| 725 | ) |
| 726 | parts = [x for x in re.split(r"(<image>)", text[0]) if x] |
| 727 | assert parts.count("<image>") == len(num_tokens_per_image), ( |
| 728 | f"Expected {len(num_tokens_per_image)} <image> tokens in text " |
| 729 | f"but found {parts.count('<image>')}" |
| 730 | ) |
| 731 | |
| 732 | for i, (feature_size, num_patches) in enumerate( |
| 733 | zip(num_tokens_per_image, image_num_patches, strict=True) |
| 734 | ): |
| 735 | image_repl = self.get_image_repl(feature_size, num_patches) |
| 736 | # image_repl.full is a list of token IDs |
| 737 | # Convert token IDs back to text for the HF processor flow |
| 738 | image_repl_text = self.tokenizer.decode( |
| 739 | image_repl.full, skip_special_tokens=False |
| 740 | ) |
| 741 | parts[i] = parts[i].replace("<image>", image_repl_text) |
| 742 | text = ["".join(parts)] |
| 743 | |
| 744 | return text, image_inputs |
| 745 | |
| 746 | def _make_batch_input(self, input_item: _T | list[_T] | None = None) -> list[_T]: |
| 747 | if input_item is None: |
| 748 | input_item = [] |
| 749 | if not isinstance(input_item, list): |
| 750 | input_item = [input_item] |
| 751 | return input_item |
| 752 | |
| 753 | @abstractmethod |
| 754 | def __call__( |
| 755 | self, |
| 756 | text: str | list[str] | None = None, |
| 757 | images: Image.Image | list[Image.Image] | None = None, |
| 758 | *, |
| 759 | return_tensors: str | TensorType | None = None, |
| 760 | max_num_tiles: int | None = None, |
| 761 | **kwargs, |
| 762 | ) -> BatchFeature: |
| 763 | raise NotImplementedError |
| 764 | |
| 765 | |
| 766 | class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): |
| 767 | """HF Processor with extended video processing logic. |
| 768 | Code for video processing is adapted from video example: |
| 769 | https://huggingface.co/OpenGVLab/InternVL3-1B#inference-with-transformers |
| 770 | """ |
| 771 | |
| 772 | def __init__( |
| 773 | self, |
| 774 | config: PreTrainedConfig, |
| 775 | tokenizer: HfTokenizer, |
| 776 | *, |
| 777 | max_model_len: int, |
| 778 | max_num_tiles: int | None = None, |
| 779 | video_token: str | None = None, |
| 780 | video_pruning_rate: float | None = None, |
| 781 | use_audio_in_video: bool = False, |
| 782 | ) -> None: |
| 783 | super().__init__( |
| 784 | config=config, |
| 785 | tokenizer=tokenizer, |
| 786 | max_model_len=max_model_len, |
| 787 | max_num_tiles=max_num_tiles, |
| 788 | ) |
| 789 | # add extra video token for video processing |
| 790 | self.video_token = video_token |
| 791 | self.video_pruning_rate = video_pruning_rate |
| 792 | self.use_audio_in_video = use_audio_in_video |
| 793 | |
| 794 | # Video params live exclusively in vision_config |
| 795 | vision_config = getattr(config, "vision_config", config) |
| 796 | self.video_temporal_patch_size: int = getattr( |
| 797 | vision_config, "video_temporal_patch_size", 1 |
| 798 | ) |
| 799 | self.video_maintain_aspect_ratio: bool = getattr( |
| 800 | vision_config, "video_maintain_aspect_ratio", False |
| 801 | ) |
| 802 | |
| 803 | # Resolve video frame target size: exactly one of video_target_num_patches |
| 804 | # or video_target_img_size may be set (mirrors Megatron's |
| 805 | # DynamicResolutionImageTilingStrategy validation). |
| 806 | target_num_patches = getattr(vision_config, "video_target_num_patches", None) |
| 807 | target_img_size = getattr(vision_config, "video_target_img_size", None) |
| 808 | if target_num_patches is not None and target_img_size is not None: |
| 809 | raise ValueError( |
| 810 | "Exactly one of video_target_num_patches or " |
| 811 | "video_target_img_size must be set, got both" |
| 812 | ) |
| 813 | if target_num_patches is not None: |
| 814 | self.video_target_num_patches: int | None = target_num_patches |
| 815 | elif target_img_size is not None: |
| 816 | base_patches = math.ceil(target_img_size / config.patch_size) |
| 817 | self.video_target_num_patches = base_patches * base_patches |
| 818 | else: |
| 819 | self.video_target_num_patches = None |
| 820 | |
| 821 | self.audio_extractor: ParakeetExtractor | None = None |
| 822 | raw_sound_config = getattr(config, "sound_config", None) |
| 823 | if raw_sound_config is not None: |
| 824 | self.audio_extractor = ParakeetExtractor(raw_sound_config) |
| 825 | |
| 826 | # Pre-tokenize special tokens for video processing |
| 827 | # to avoid repeated tokenization |
| 828 | self._img_start_token_ids = tokenizer.encode( |
| 829 | IMG_START, add_special_tokens=False |
| 830 | ) |
| 831 | self._img_end_token_ids = tokenizer.encode(IMG_END, add_special_tokens=False) |
| 832 | self._img_context_token_ids = tokenizer.encode( |
| 833 | IMG_CONTEXT, add_special_tokens=False |
| 834 | ) |
| 835 | |
| 836 | @cached_property |
| 837 | def num_video_token(self) -> int: |
| 838 | """Token count per video frame, accounting for video_target_num_patches. |
| 839 | |
| 840 | When video_target_num_patches is set the per-frame feature count |
| 841 | differs from the image-based num_image_token. We use a square |
| 842 | dummy (1:1) to compute the feature_size because the dummy video is |
| 843 | square and the user confirmed that is acceptable. |
| 844 | """ |
| 845 | if self.video_target_num_patches is not None: |
| 846 | _, _, feature_size = get_video_target_size_and_feature_size( |
| 847 | orig_w=self.image_size, |
| 848 | orig_h=self.image_size, |
| 849 | target_patches=self.video_target_num_patches, |
| 850 | maintain_aspect_ratio=self.video_maintain_aspect_ratio, |
| 851 | patch_size=self.config.patch_size, |
| 852 | downsample_ratio=self.config.downsample_ratio, |
| 853 | ) |
| 854 | return feature_size |
| 855 | return self.num_image_token |
| 856 | |
| 857 | @property |
| 858 | def supports_video(self) -> bool: |
| 859 | return True |
| 860 | |
| 861 | @property |
| 862 | def video_token_id(self) -> int: |
| 863 | assert self.video_token is not None |
| 864 | return self.tokenizer.get_vocab()[self.video_token] |
| 865 | |
| 866 | @property |
| 867 | def image_token_id(self) -> int: |
| 868 | return self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT) |
| 869 | |
| 870 | def _videos_to_pixel_values_lst( |
| 871 | self, |
| 872 | videos: list[npt.NDArray], |
| 873 | *, |
| 874 | dtype: torch.dtype = torch.float32, |
| 875 | ) -> list[torch.Tensor]: |
| 876 | return [ |
| 877 | video_to_pixel_values( |
| 878 | video, |
| 879 | input_size=self.image_size, |
| 880 | video_target_num_patches=self.video_target_num_patches, |
| 881 | video_maintain_aspect_ratio=self.video_maintain_aspect_ratio, |
| 882 | patch_size=self.config.patch_size, |
| 883 | downsample_ratio=self.config.downsample_ratio, |
| 884 | norm_mean=self.norm_mean, |
| 885 | norm_std=self.norm_std, |
| 886 | dtype=dtype, |
| 887 | ) |
| 888 | for video in videos |
| 889 | ] |
| 890 | |
| 891 | def _preprocess_video( |
| 892 | self, |
| 893 | text: list[str], |
| 894 | videos: list[tuple[npt.NDArray, dict[str, Any]]], |
| 895 | ) -> tuple[list[str], dict[str, Any]]: |
| 896 | if len(videos) == 0 or not self.supports_video: |
| 897 | return text, {} |
| 898 | |
| 899 | videos_lst = [v[0] for v in videos] |
| 900 | video_metadata_lst = [v[1] for v in videos] |
| 901 | |
| 902 | pixel_values_lst_video = self._videos_to_pixel_values_lst( |
| 903 | videos_lst, |
| 904 | dtype=self.dtype, |
| 905 | ) |
| 906 | |
| 907 | # We use frame duration in milliseconds (as integer) to ensure |
| 908 | # we have consistent timestamps calculation. At preprocessing |
| 909 | # fps parameter is given in fp32, while at inference it is bf16 |
| 910 | # which leads to inaccurate timestamp calculation and causes |
| 911 | # timestamp values to differ.In rare cases this causes |
| 912 | # mismatching number of output tokens for tokenized frame prefixes |
| 913 | frame_duration_ms_lst = [ |
| 914 | int(1000.0 / metadata["fps"]) for metadata in video_metadata_lst |
| 915 | ] |
| 916 | frames_indices_lst = [ |
| 917 | metadata["frames_indices"] for metadata in video_metadata_lst |
| 918 | ] |
| 919 | video_num_patches = torch.tensor([len(item) for item in pixel_values_lst_video]) |
| 920 | |
| 921 | # Normalization already fused into resize above. |
| 922 | # Skip the torch.cat copy when there is exactly one video |
| 923 | if len(pixel_values_lst_video) == 1: |
| 924 | pixel_values_flat = pixel_values_lst_video[0] |
| 925 | else: |
| 926 | pixel_values_flat = torch.cat(pixel_values_lst_video) |
| 927 | video_inputs = { |
| 928 | "pixel_values_flat_video": pixel_values_flat, |
| 929 | "video_num_patches": video_num_patches, |
| 930 | "frames_indices": frames_indices_lst, |
| 931 | "frame_duration_ms": torch.tensor(frame_duration_ms_lst), |
| 932 | } |
| 933 | |
| 934 | patch_size: int = self.config.patch_size |
| 935 | downsample_ratio = self.config.downsample_ratio |
| 936 | |
| 937 | T = self.video_temporal_patch_size |
| 938 | |
| 939 | for pixel_values, video_metadata, frames_indices, frame_duration_ms in zip( |
| 940 | pixel_values_lst_video, |
| 941 | video_metadata_lst, |
| 942 | frames_indices_lst, |
| 943 | frame_duration_ms_lst, |
| 944 | ): |
| 945 | num_frames = pixel_values.shape[0] |
| 946 | frame_h, frame_w = pixel_values.shape[-2], pixel_values.shape[-1] |
| 947 | tokens_in_single_frame = int( |
| 948 | (frame_h * frame_w // patch_size**2) * (downsample_ratio**2) |
| 949 | ) |
| 950 | num_tubelets = math.ceil(num_frames / T) if T > 1 else num_frames |
| 951 | |
| 952 | if self.video_pruning_rate is not None and self.video_pruning_rate > 0.0: |
| 953 | # Start of EVS-specific code |
| 954 | num_tokens = compute_retained_tokens_count( |
| 955 | tokens_per_frame=tokens_in_single_frame, |
| 956 | num_frames=num_tubelets, |
| 957 | q=self.video_pruning_rate, |
| 958 | ) |
| 959 | |
| 960 | # Here we just need placeholders that won't actually be replaced - |
| 961 | # we just need to make sure the total number of tokens is correct |
| 962 | # assign all tokens to the first frame |
| 963 | tokens_per_frame = [num_tokens] + [0] * (num_tubelets - 1) |
| 964 | |
| 965 | # End of EVS-specific code |
| 966 | else: |
| 967 | tokens_per_frame = [tokens_in_single_frame] * num_tubelets |
| 968 | |
| 969 | video_repl = self.get_video_repl( |
| 970 | tokens_per_frame=tokens_per_frame, |
| 971 | frames_indices=frames_indices, |
| 972 | frame_duration_ms=frame_duration_ms, |
| 973 | tokenizer=self.tokenizer, |
| 974 | img_start_token_ids=self._img_start_token_ids, |
| 975 | img_end_token_ids=self._img_end_token_ids, |
| 976 | img_context_token_ids=self._img_context_token_ids, |
| 977 | video_temporal_patch_size=T, |
| 978 | ) |
| 979 | |
| 980 | # video_repl.full is a list of token IDs |
| 981 | # Convert token IDs back to text for the HF processor flow |
| 982 | video_repl_text = self.tokenizer.decode( |
| 983 | video_repl.full, skip_special_tokens=False |
| 984 | ) |
| 985 | text = [t.replace("<video>", video_repl_text, 1) for t in text] |
| 986 | |
| 987 | return text, video_inputs |
| 988 | |
| 989 | def _preprocess_audio( |
| 990 | self, |
| 991 | text: list[str], |
| 992 | audios: list[npt.NDArray], |
| 993 | ) -> tuple[list[str], dict[str, Any]]: |
| 994 | if len(audios) == 0: |
| 995 | return text, {"audio_num_clips": []} |
| 996 | |
| 997 | assert self.audio_extractor is not None |
| 998 | extractor = self.audio_extractor |
| 999 | |
| 1000 | parts = [x for x in re.split(f"({re.escape(AUDIO_CONTEXT)})", text[0]) if x] |
| 1001 | token_count = parts.count(AUDIO_CONTEXT) |
| 1002 | if token_count != len(audios): |
| 1003 | raise ValueError( |
| 1004 | "Number of audio tokens in text does not match the number " |
| 1005 | f"of audios (tokens={token_count}, audios={len(audios)})." |
| 1006 | ) |
| 1007 | audio_index = 0 |
| 1008 | for idx, part in enumerate(parts): |
| 1009 | if part == AUDIO_CONTEXT: |
| 1010 | audio_repl = self.get_audio_repl(audios[audio_index]) |
| 1011 | # audio_repl.full is a list of token IDs |
| 1012 | # Convert token IDs back to text for the HF processor flow |
| 1013 | parts[idx] = self.tokenizer.decode( |
| 1014 | audio_repl.full, skip_special_tokens=False |
| 1015 | ) |
| 1016 | audio_index += 1 |
| 1017 | text = ["".join(parts)] |
| 1018 | audio_inputs = extractor(audios) |
| 1019 | return text, audio_inputs |
| 1020 | |
| 1021 | def __call__( |
| 1022 | self, |
| 1023 | text: str | list[str] | None = None, |
| 1024 | images: Image.Image | list[Image.Image] | None = None, |
| 1025 | videos: tuple[npt.NDArray, dict[str, Any]] |
| 1026 | | list[tuple[npt.NDArray, dict[str, Any]]] |
| 1027 | | None = None, |
| 1028 | audios: AudioItem | list[AudioItem] | None = None, |
| 1029 | *, |
| 1030 | return_tensors: str | TensorType | None = None, |
| 1031 | max_num_tiles: int | None = None, |
| 1032 | **kwargs, |
| 1033 | ) -> BatchFeature: |
| 1034 | # Use default if not provided |
| 1035 | if max_num_tiles is None: |
| 1036 | max_num_tiles = self.max_num_tiles |
| 1037 | |
| 1038 | text = self._make_batch_input(text) |
| 1039 | images = self._make_batch_input(images) |
| 1040 | videos = self._make_batch_input(videos) |
| 1041 | audios = self._make_batch_input(audios) |
| 1042 | |
| 1043 | text, image_inputs = self._preprocess_image( |
| 1044 | text=text, |
| 1045 | images=images, |
| 1046 | max_num_tiles=max_num_tiles, |
| 1047 | ) |
| 1048 | |
| 1049 | text, video_inputs = self._preprocess_video( |
| 1050 | text=text, |
| 1051 | videos=videos, |
| 1052 | ) |
| 1053 | |
| 1054 | text, audio_inputs = self._preprocess_audio( |
| 1055 | text=text, |
| 1056 | audios=audios, |
| 1057 | ) |
| 1058 | |
| 1059 | text_inputs = self.tokenizer(text, add_special_tokens=False) |
| 1060 | |
| 1061 | combined_inputs = {**text_inputs, **video_inputs, **audio_inputs} |
| 1062 | frames_indices = combined_inputs.get("frames_indices") |
| 1063 | ragged_frames_indices = ( |
| 1064 | isinstance(frames_indices, list) |
| 1065 | and len({len(frame_indices) for frame_indices in frames_indices}) > 1 |
| 1066 | ) |
| 1067 | if ragged_frames_indices: |
| 1068 | combined_inputs.pop("frames_indices") |
| 1069 | |
| 1070 | if self.dynamic_tiler is None: |
| 1071 | batch = BatchFeature( |
| 1072 | {**combined_inputs, **image_inputs}, |
| 1073 | tensor_type=return_tensors, |
| 1074 | ) |
| 1075 | else: |
| 1076 | batch = BatchFeature(combined_inputs, tensor_type=return_tensors) |
| 1077 | # allow images to be exempt from the BatchFeature validation: |
| 1078 | # We will .stack() them in _parse_and_validate_image_input |
| 1079 | batch.update(image_inputs) |
| 1080 | if ragged_frames_indices: |
| 1081 | assert isinstance(frames_indices, list) |
| 1082 | batch["frames_indices"] = [ |
| 1083 | torch.as_tensor(frame_indices, dtype=torch.int64) |
| 1084 | for frame_indices in frames_indices |
| 1085 | ] |
| 1086 | return batch |
| 1087 | |
| 1088 | def get_image_repl( |
| 1089 | self, |
| 1090 | feature_size: int, |
| 1091 | num_patches: int | None, |
| 1092 | ) -> PromptUpdateDetails: |
| 1093 | repl_features = IMG_CONTEXT * feature_size |
| 1094 | repl_full = IMG_START + repl_features + IMG_END |
| 1095 | |
| 1096 | full_ids = cached_encode(self.tokenizer, repl_full, add_special_tokens=False) |
| 1097 | |
| 1098 | return PromptUpdateDetails.select_token_ids( |
| 1099 | full_ids, self._img_context_token_ids |
| 1100 | ) |
| 1101 | |
| 1102 | def get_audio_repl( |
| 1103 | self, |
| 1104 | audio: npt.NDArray, |
| 1105 | ) -> PromptUpdateDetails: |
| 1106 | assert self.audio_extractor is not None |
| 1107 | num_tokens = self.audio_extractor.audio_token_count(len(audio)) |
| 1108 | repl_full = f"{AUDIO_START}{AUDIO_CONTEXT * num_tokens}{AUDIO_END}" |
| 1109 | |
| 1110 | full_ids = cached_encode(self.tokenizer, repl_full, add_special_tokens=False) |
| 1111 | ctx_token_ids = cached_encode( |
| 1112 | self.tokenizer, AUDIO_CONTEXT, add_special_tokens=False |
| 1113 | ) |
| 1114 | |
| 1115 | return PromptUpdateDetails.select_token_ids(full_ids, ctx_token_ids) |
| 1116 | |
| 1117 | @classmethod |
| 1118 | def get_video_repl( |
| 1119 | cls, |
| 1120 | *, |
| 1121 | tokens_per_frame: list[int], |
| 1122 | frames_indices: list[int], |
| 1123 | frame_duration_ms: int, |
| 1124 | tokenizer: HfTokenizer, |
| 1125 | img_start_token_ids: list[int], |
| 1126 | img_end_token_ids: list[int], |
| 1127 | img_context_token_ids: list[int], |
| 1128 | video_temporal_patch_size: int = 1, |
| 1129 | ) -> PromptUpdateDetails: |
| 1130 | """Build prompt replacement for a video. |
| 1131 | The replacement returned is not actually used to replace the placeholder |
| 1132 | tokens - it's just used to make sure we allocate the correct number |
| 1133 | of tokens. |
| 1134 | Actual replacement is done in embed_multimodal of |
| 1135 | NemotronH_Nano_VL_V2 |
| 1136 | (specifically in _process_video_input -> _create_final_video_embeddings). |
| 1137 | There, we create the final embeddings with text embeddings for indicator tokens |
| 1138 | and video embeddings for video tokens. |
| 1139 | This is a single function that handles all cases - non EVS, EVS dummy, EVS real. |
| 1140 | The differentiation is done via tokens_per_frame parameter. |
| 1141 | - non EVS case - constant value same value across all frames |
| 1142 | - EVS dummy - Doesn't matter how tokens are distributed between frames - just |
| 1143 | make sure the total number of tokens is correct. |
| 1144 | - EVS real (called from get_real_video_repl_for_evs) - different value per frame |
| 1145 | Args: |
| 1146 | tokens_per_frame (list[int]): number of tokens per frame |
| 1147 | (one per tubelet when T > 1) |
| 1148 | frames_indices (list[int]): orig. frame indices |
| 1149 | (one per frame, before tubelet subsampling) |
| 1150 | frame_duration_ms (int): duration of each frame in milliseconds |
| 1151 | tokenizer (TokenizerLike): tokenizer to use for tokenizing frame separators |
| 1152 | img_start_token_ids (list[int]): pre-tokenized IMG_START tokens |
| 1153 | img_end_token_ids (list[int]): pre-tokenized IMG_END tokens |
| 1154 | img_context_token_ids (list[int]): pre-tokenized IMG_CONTEXT tokens |
| 1155 | video_temporal_patch_size (int): temporal patch size for videos |
| 1156 | """ |
| 1157 | # TODO: Add support of frame_duration_ms to be None |
| 1158 | # At preprocessing step we should allow absent / metadata without |
| 1159 | # frames_indices field. |
| 1160 | timestamps_enabled = frame_duration_ms is not None |
| 1161 | T = video_temporal_patch_size |
| 1162 | num_frames = len(frames_indices) |
| 1163 | |
| 1164 | if T > 1 and timestamps_enabled: |
| 1165 | all_timestamps = calculate_timestamps(frames_indices, frame_duration_ms) |
| 1166 | |
| 1167 | frame_separators = [] |
| 1168 | for group_idx, i in enumerate(range(0, num_frames, T)): |
| 1169 | group_frames = [] |
| 1170 | for j in range(T): # Every frame in the group |
| 1171 | frame_idx = i + j |
| 1172 | if frame_idx < num_frames: |
| 1173 | # Valid idx (haven't padded to mult. of T yet) |
| 1174 | ts = all_timestamps[frame_idx] |
| 1175 | frame_str = "Frame" if j == 0 else "frame" |
| 1176 | group_frames.append( |
| 1177 | f"{frame_str} {frame_idx + 1} sampled at {ts:.2f} seconds" |
| 1178 | ) |
| 1179 | if group_frames: |
| 1180 | # Join by `and` if there are >1 frame, otherwise no `and` |
| 1181 | # Prepend \n to match training format (except first group) |
| 1182 | sep = " and ".join(group_frames) + ": " |
| 1183 | if group_idx > 0: |
| 1184 | sep = "\n" + sep |
| 1185 | frame_separators.append(sep) |
| 1186 | elif timestamps_enabled: |
| 1187 | timestamps = calculate_timestamps(frames_indices, frame_duration_ms) |
| 1188 | |
| 1189 | assert len(timestamps) == len(tokens_per_frame), ( |
| 1190 | "timestamps and tokens_per_frame must have the same length" |
| 1191 | ) |
| 1192 | frame_separators = [ |
| 1193 | ("\n" if i > 0 else "") |
| 1194 | + f"Frame {i + 1} sampled at {timestamp:.2f} seconds: " |
| 1195 | for i, timestamp in enumerate(timestamps) |
| 1196 | ] |
| 1197 | else: |
| 1198 | frame_separators = [ |
| 1199 | ("\n" if i > 0 else "") + f"Frame {i + 1}: " |
| 1200 | for i, _ in enumerate(tokens_per_frame) |
| 1201 | ] |
| 1202 | |
| 1203 | # Batch-tokenize all frame separators at once — the HuggingFace |
| 1204 | # tokenizers Rust backend parallelizes batch encoding across threads. |
| 1205 | batch_encoded = tokenizer( |
| 1206 | frame_separators, |
| 1207 | add_special_tokens=False, |
| 1208 | return_attention_mask=False, |
| 1209 | ) |
| 1210 | frame_separators_tokenized: list[list[int]] = batch_encoded["input_ids"] |
| 1211 | |
| 1212 | # Tokenize each component independently to avoid tokenizer merging tokens |
| 1213 | # across boundaries. This ensures consistent tokenization regardless of |
| 1214 | # num_tokens_per_frame values. |
| 1215 | all_token_ids = [] |
| 1216 | for i, num_tokens in enumerate(tokens_per_frame): |
| 1217 | all_token_ids.extend(frame_separators_tokenized[i]) |
| 1218 | all_token_ids.extend(img_start_token_ids) |
| 1219 | all_token_ids.extend(img_context_token_ids * num_tokens) |
| 1220 | all_token_ids.extend(img_end_token_ids) |
| 1221 | |
| 1222 | return PromptUpdateDetails.from_seq(all_token_ids) |