@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cd3ac5b797
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 00:15 UTC

support dynamic resolution image encoding for Nemotron Nano VL (#32121)

Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>

Netanel Haber committed Jan 19, 2026 at 20:15 UTC cd3ac5b79703b350c7b66053dcec56300b33deb0
3 files changed +755 -164
vllm/model_executor/models/intern_vit.py
+6 -2
@@ -282,12 +282,14 @@ class InternVisionEncoderLayer(nn.Module):
282 num_dummy_heads: int = 0,
283 prefix: str = "",
284 use_data_parallel: bool = False,
285 + attn_cls: type[InternParallelAttention] = InternParallelAttention,
286 ) -> None:
287 super().__init__()
288
289 self.embed_dim = config.hidden_size
290 self.intermediate_size = config.intermediate_size
291 self.norm_type = config.norm_type
292 + self.attn_cls = attn_cls
293
294 self.attn = self._init_attn(
295 config,
@@ -327,7 +329,7 @@ class InternVisionEncoderLayer(nn.Module):
329 use_data_parallel = (
330 use_data_parallel or (num_heads + num_dummy_heads) % tp_size != 0
331 )
330 - return InternParallelAttention(
332 + return self.attn_cls(
333 config,
334 quant_config=quant_config,
335 num_dummy_heads=num_dummy_heads,
@@ -356,10 +358,12 @@ class InternVisionEncoder(nn.Module):
358 num_dummy_heads: int = 0,
359 prefix: str = "",
360 use_data_parallel: bool = False,
361 + layer_cls: type[InternVisionEncoderLayer] = InternVisionEncoderLayer,
362 ):
363 super().__init__()
364
365 self.config = config
366 + self.layer_cls = layer_cls
367
368 if num_hidden_layers_override is None:
369 num_hidden_layers = config.num_hidden_layers
@@ -368,7 +372,7 @@ class InternVisionEncoder(nn.Module):
372
373 self.layers = nn.ModuleList(
374 [
371 - InternVisionEncoderLayer(
375 + self.layer_cls(
376 config,
377 quant_config,
378 num_dummy_heads=num_dummy_heads,
vllm/model_executor/models/nano_nemotron_vl.py
+544 -141
@@ -8,11 +8,15 @@
8 # --------------------------------------------------------
9
10 import copy
11 +import math
12 import warnings
13 from abc import ABC, abstractmethod
14 from collections.abc import Iterable, Mapping, Sequence
15 +from dataclasses import dataclass
16 +from functools import cached_property
17 from typing import Annotated, Any, Literal, TypeAlias, TypeVar
18
19 +import einops
20 import numpy.typing as npt
21 import regex as re
22 import torch
@@ -23,6 +27,7 @@ from transformers import BatchFeature, PretrainedConfig, TensorType
27
28 from vllm.config import VllmConfig
29 from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
30 +from vllm.logger import init_logger
31 from vllm.model_executor.layers.activation import ReLUSquaredActivation
32 from vllm.model_executor.layers.layernorm import RMSNorm
33 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -39,7 +44,7 @@ from vllm.model_executor.models.internvl import (
44 )
45 from vllm.model_executor.models.module_mapping import MultiModelKeys
46 from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
42 -from vllm.model_executor.models.radio import RadioModel
47 +from vllm.model_executor.models.radio import RadioModel, calc_seq_lens
48 from vllm.model_executor.models.utils import (
49 init_vllm_registered_model,
50 maybe_prefix,
@@ -78,6 +83,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
83
84 from .utils import _merge_multimodal_embeddings
85
86 +logger = init_logger(__name__)
87 # Configure PIL to handle large images without warnings
88 # This prevents DecompressionBombWarning for legitimate large images
89 Image.MAX_IMAGE_PIXELS = None # Disable the limit entirely
@@ -103,11 +109,25 @@ class NanoNemotronVLImagePixelInputs(TensorSchema):
109 - w: Width of each image patch
110 """
111
106 - type: Literal["pixel_values"]
112 + type: Literal["pixel_values"] = "pixel_values"
113 pixel_values_flat: Annotated[torch.Tensor, TensorShape("bnp", 3, "h", "w")]
114 num_patches: Annotated[torch.Tensor, TensorShape("bn")]
115
116
117 +class NanoNemotronVLImagePixelInputsDynamic(TensorSchema):
118 + """
119 + Dynamic-resolution image inputs.
120 +
121 + imgs_sizes: per-image (height, width) in pixels.
122 + num_tokens_per_image: per-image number of embedding tokens (post downsample).
123 + """
124 +
125 + type: Literal["pixel_values_dynamic"] = "pixel_values_dynamic"
126 + pixel_values_flat: Annotated[torch.Tensor, TensorShape("bn", "h", "w")]
127 + imgs_sizes: list[tuple[int, int]]
128 + num_tokens_per_image: list[int]
129 +
130 +
131 class NanoNemotronVLImageEmbeddingInputs(TensorSchema):
132 """
133 Dimensions:
@@ -121,7 +141,9 @@ class NanoNemotronVLImageEmbeddingInputs(TensorSchema):
141
142
143 NanoNemotronVLImageInputs: TypeAlias = (
124 - NanoNemotronVLImagePixelInputs | NanoNemotronVLImageEmbeddingInputs
144 + NanoNemotronVLImagePixelInputs
145 + | NanoNemotronVLImagePixelInputsDynamic
146 + | NanoNemotronVLImageEmbeddingInputs
147 )
148
149
@@ -267,6 +289,329 @@ def calculate_timestamps(
289 return timestamps
290
291
292 +class DynamicResolutionImageTiler:
293 + CONV_MERGING = False
294 + PIXEL_SHUFFLE = True
295 + USE_THUMBNAIL = False
296 +
297 + def __init__(
298 + self,
299 + *,
300 + max_model_len: int,
301 + patch_size: int,
302 + min_num_patches: int,
303 + max_num_patches: int,
304 + downsample_ratio: int,
305 + norm_mean: Sequence[float],
306 + norm_std: Sequence[float],
307 + factor_max: float = 1.0,
308 + use_thumbnail: bool = False,
309 + ) -> None:
310 + assert use_thumbnail is False, "use_thumbnail is not supported"
311 + self._patch_size: int = patch_size
312 + self._max_model_len = max_model_len
313 + self._min_num_patches = min_num_patches
314 + self._max_num_patches = max_num_patches if max_num_patches > 0 else float("inf")
315 + self._factor_max = factor_max
316 + self.norm_mean = torch.tensor(norm_mean).reshape(3, 1, 1)
317 + self.norm_std = torch.tensor(norm_std).reshape(3, 1, 1)
318 + self._transform = T.Compose(
319 + [
320 + T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
321 + T.ToTensor(),
322 + ]
323 + )
324 + assert downsample_ratio < 1
325 + reduction_factor = 1 / downsample_ratio
326 + assert reduction_factor == 2.0
327 + self._downsample_ratio = int(reduction_factor) ** (
328 + self.PIXEL_SHUFFLE + self.CONV_MERGING
329 + )
330 + assert self._downsample_ratio == 2
331 +
332 + def _get_num_embeddings(self, width: int, height: int) -> int:
333 + num_patches = (width // self._patch_size) * (height // self._patch_size)
334 + num_tokens = num_patches // (self._downsample_ratio**2)
335 + return num_tokens
336 +
337 + def width_and_height_for_max_num_tokens_available(
338 + self,
339 + target_num_tokens_post_shuffle: int,
340 + ) -> tuple[int, int]:
341 + """
342 + TODO: optimize this so it squeezes closer to target number of tokens.
343 + Calculate image dimensions that produce approximately `target` tokens after
344 + pixel_shuffle.
345 +
346 + With pixel_shuffle enabled, each 2x2 patch grid becomes 1 token, so we
347 + need 4*B patches to get B tokens.
348 +
349 + Examples:
350 + >>> PATCH_SIZE = 16
351 + >>> DOWNSAMPLE_RATIO = 0.5
352 + >>> tiler = DynamicResolutionImageTiler(
353 + ... max_model_len=16384,
354 + ... patch_size=PATCH_SIZE,
355 + ... downsample_ratio=DOWNSAMPLE_RATIO,
356 + ... min_num_patches=4,
357 + ... max_num_patches=0,
358 + ... )
359 + >>> width, height = tiler.width_and_height_for_max_num_tokens_available(
360 + ... target_num_tokens_post_shuffle=8192,
361 + ... )
362 + >>> assert width, height == (2880, 2880)
363 + >>> assert (width // PATCH_SIZE) * (
364 + ... height // PATCH_SIZE
365 + ... ) // 2**2 == 8100 # tokens post-shuffle
366 + >>> assert tiler._get_num_embeddings(width=width, height=height) == 8100
367 + """
368 + side_pixels = (
369 + math.isqrt(target_num_tokens_post_shuffle)
370 + * self._downsample_ratio
371 + * self._patch_size
372 + )
373 + assert isinstance(side_pixels, int) and side_pixels % self._patch_size == 0
374 + return side_pixels, side_pixels
375 +
376 + def max_num_tokens_available(self, text_prompt_length: int) -> int:
377 + return self._max_model_len - text_prompt_length - 4
378 +
379 + def _images_to_pixel_values_lst(
380 + self,
381 + text_prompt_length: int,
382 + images: list[Image.Image],
383 + ) -> tuple[list[torch.Tensor], list[int]]:
384 + num_tokens_available = self.max_num_tokens_available(text_prompt_length)
385 + params_per_image = self.compute_params(images, num_tokens_available)
386 +
387 + feature_sizes = []
388 + images = []
389 + for param in params_per_image:
390 + for t in self.apply_params(param):
391 + assert t.ndim == 3, f"{t.ndim=}: expected 3 dim tensor"
392 + images.append(t)
393 + feature_sizes.append(param.num_embeddings)
394 + return images, feature_sizes
395 +
396 + feature_size_cache: dict[Image.Image, int] = {}
397 +
398 + @classmethod
399 + def get_cached_feature_size(cls, image: Image.Image) -> int:
400 + feature_size = cls.feature_size_cache[id(image)]
401 + # hard assert that we only use the feature size once
402 + del cls.feature_size_cache[id(image)]
403 + return feature_size
404 +
405 + @dataclass
406 + class DynamicResolutionParams:
407 + media: Image.Image
408 + num_tiles: int
409 + num_embeddings: int
410 + patch_size: tuple[int, int]
411 +
412 + def apply_params(self, params: DynamicResolutionParams) -> list[torch.Tensor]:
413 + resized_img = params.media.resize(
414 + (
415 + params.patch_size[0] * self._patch_size,
416 + params.patch_size[1] * self._patch_size,
417 + )
418 + )
419 + processed_images = [resized_img]
420 +
421 + return [self._transform(img) for img in processed_images]
422 +
423 + def process_media(
424 + self,
425 + media: Image.Image,
426 + num_tokens_available: int,
427 + ) -> tuple[DynamicResolutionParams, int]:
428 + """Process a single media item and return its parameters.
429 +
430 + Args:
431 + media: The media item to process
432 + num_tokens_available: Number of tokens available for this media
433 + Returns:
434 + DynamicResolutionParams for the media
435 + """
436 + current_num_tokens_available = num_tokens_available
437 + assert isinstance(media, Image.Image), (
438 + "Dynamic resolution is only supported for image media"
439 + )
440 + orig_width, orig_height = media.width, media.height
441 + closest_patch_height = round(orig_height / self._patch_size + 0.5)
442 + closest_patch_width = round(orig_width / self._patch_size + 0.5)
443 + patches = closest_patch_height * closest_patch_width
444 +
445 + factor = min(
446 + math.sqrt(current_num_tokens_available / patches), self._factor_max
447 + )
448 + target_patch_height = math.floor(factor * closest_patch_height)
449 + target_patch_width = math.floor(factor * closest_patch_width)
450 +
451 + # Consider self._min_num_patches if > current_num_tokens_available.
452 + if (
453 + current_num_tokens_available > self._min_num_patches
454 + and target_patch_height * target_patch_width < self._min_num_patches
455 + ):
456 + up_factor = math.sqrt(
457 + self._min_num_patches / (target_patch_height * target_patch_width)
458 + )
459 + target_patch_height = math.ceil(up_factor * target_patch_height)
460 + target_patch_width = math.ceil(up_factor * target_patch_width)
461 +
462 + # Round patch grid to be divisible by 2 (pixel-shuffle OR conv-merging)
463 + # or by 4 when BOTH are enabled (two successive 2x reductions)
464 + if self.PIXEL_SHUFFLE or self.CONV_MERGING:
465 + required_divisor = 4 if (self.PIXEL_SHUFFLE and self.CONV_MERGING) else 2
466 +
467 + rem_h = target_patch_height % required_divisor
468 + if rem_h != 0:
469 + inc_h = required_divisor - rem_h
470 + if (
471 + target_patch_height + inc_h
472 + ) * target_patch_width <= current_num_tokens_available:
473 + target_patch_height += inc_h
474 + else:
475 + target_patch_height = max(
476 + required_divisor, target_patch_height - rem_h
477 + )
478 +
479 + rem_w = target_patch_width % required_divisor
480 + if rem_w != 0:
481 + inc_w = required_divisor - rem_w
482 + if (
483 + target_patch_height * (target_patch_width + inc_w)
484 + <= current_num_tokens_available
485 + ):
486 + target_patch_width += inc_w
487 + else:
488 + target_patch_width = max(
489 + required_divisor, target_patch_width - rem_w
490 + )
491 +
492 + # Calculate embeddings for the main dynamic resolution image
493 + num_embeddings = self._get_num_embeddings(
494 + target_patch_width * self._patch_size,
495 + target_patch_height * self._patch_size,
496 + )
497 +
498 + token_count = target_patch_width * target_patch_height
499 +
500 + # Add thumbnail embeddings if enabled and image area is below threshold
501 + num_tiles = 1 # Base dynamic resolution image
502 +
503 + return self.DynamicResolutionParams(
504 + media=media,
505 + num_tiles=num_tiles,
506 + num_embeddings=num_embeddings,
507 + patch_size=(target_patch_width, target_patch_height),
508 + ), token_count
509 +
510 + def compute_params(
511 + self,
512 + media_list: list[Image.Image],
513 + num_tokens_available: int | None = None,
514 + ) -> list[DynamicResolutionParams]:
515 + """Compute parameters for all media with iterative token budgeting.
516 +
517 + Args:
518 + media_list: List of media items to process
519 + num_tokens_available: Total number of tokens available across all media
520 + Returns:
521 + List of ImageTilingParams for each media item
522 + """
523 + num_tokens_available = (
524 + num_tokens_available
525 + * (4 if self.PIXEL_SHUFFLE else 1)
526 + * (4 if self.CONV_MERGING else 1)
527 + )
528 + # When the number of available token is too small,
529 + # allow self._min_num_patches per media and let the sample be truncated.
530 + num_tokens_available = max(
531 + num_tokens_available, self._min_num_patches * len(media_list)
532 + )
533 +
534 + # Clip the number of tokens available per media to >min and <max patches.
535 + num_tokens_available_per_media = [
536 + max(min(num_tokens_available, self._max_num_patches), self._min_num_patches)
537 + for _ in range(len(media_list))
538 + ]
539 +
540 + # prevent infinite loop in any case
541 + for _ in range(10):
542 + # Step 1: Process each media with current token budget
543 + params = []
544 + token_counts = []
545 +
546 + for media, tokens_for_media in zip(
547 + media_list, num_tokens_available_per_media
548 + ):
549 + param, token_count = self.process_media(media, tokens_for_media)
550 + params.append(param)
551 + token_counts.append(token_count)
552 + self.feature_size_cache[id(param.media)] = param.num_embeddings
553 +
554 + # Step 2: Check if total tokens is within budget
555 + total_tokens = sum(token_counts)
556 +
557 + if total_tokens <= num_tokens_available:
558 + # We're within budget, return the params
559 + return params
560 +
561 + # Step 3: We're over budget, need to scale down
562 + # Calculate scaling factor to get under budget
563 + scaling_factor = num_tokens_available / total_tokens
564 +
565 + # Recalculate token budgets for each media based on scaling
566 + # Each media gets a proportional share of the total budget
567 + scaled_down_num_tokens_available_per_media = [
568 + max(self._min_num_patches, int(token_count * scaling_factor))
569 + for token_count in token_counts
570 + ]
571 + scaled_down = any(
572 + [
573 + scaled_down_num_tokens_available_per_media[i]
574 + < num_tokens_available_per_media[i]
575 + for i in range(len(num_tokens_available_per_media))
576 + ]
577 + )
578 + # If there wasn't scaling down, we're stuck with min_num_patches per media,
579 + # else try with the scaled down num_tokens_available_per_media.
580 + if not scaled_down:
581 + num_tokens_available_per_media = [self._min_num_patches] * len(
582 + media_list
583 + )
584 + else:
585 + num_tokens_available_per_media = (
586 + scaled_down_num_tokens_available_per_media
587 + )
588 + ctx = f"{params=} {total_tokens=} {num_tokens_available=}"
589 + raise ValueError(
590 + f"Should be unreachable - `return params` above must be reached: {ctx}"
591 + )
592 +
593 + @staticmethod
594 + def stack(images: list[torch.Tensor], patch_size: int) -> torch.Tensor:
595 + assert len(images) > 0, "No images to stack"
596 +
597 + def rearrange_img(x):
598 + py = x.shape[-2] // patch_size
599 + px = x.shape[-1] // patch_size
600 + x = einops.rearrange(
601 + x,
602 + "c (py yy) (px xx) -> (py px) (c yy xx)",
603 + py=py,
604 + yy=patch_size,
605 + px=px,
606 + xx=patch_size,
607 + )
608 + return x
609 +
610 + imgs = [rearrange_img(img) for img in images]
611 + pixel_values_flat = torch.cat(imgs, dim=0).unsqueeze(0)
612 + return pixel_values_flat
613 +
614 +
615 class BaseNanoNemotronVLProcessor(ABC):
616 """
617 This model doesn't define its own HF processor,
@@ -281,6 +626,7 @@ class BaseNanoNemotronVLProcessor(ABC):
626 config: PretrainedConfig,
627 tokenizer: TokenizerLike,
628 *args,
629 + max_model_len: int,
630 max_num_tiles: int | None = None,
631 **kwargs,
632 ) -> None:
@@ -292,15 +638,32 @@ class BaseNanoNemotronVLProcessor(ABC):
638 self.max_num_tiles = max_num_tiles or DEFAULT_NUM_TILES
639 image_size: int = config.force_image_size
640 patch_size: int = config.patch_size
641 + downsample_ratio: int = config.downsample_ratio
642
643 self.num_image_token = int(
297 - (image_size // patch_size) ** 2 * (config.downsample_ratio**2)
644 + (image_size // patch_size) ** 2 * (downsample_ratio**2)
645 )
646 self.image_size = image_size
647 self.use_thumbnail: bool = config.use_thumbnail
648 self.norm_mean = torch.Tensor(config.norm_mean).reshape(1, 3, 1, 1)
649 self.norm_std = torch.Tensor(config.norm_std).reshape(1, 3, 1, 1)
650
651 + self.dynamic_tiler: DynamicResolutionImageTiler | None = None
652 + if self.use_dynamic_resolution(config):
653 + self.dynamic_tiler = DynamicResolutionImageTiler(
654 + max_model_len=max_model_len,
655 + patch_size=patch_size,
656 + downsample_ratio=downsample_ratio,
657 + min_num_patches=config.vision_config.args["min_num_patches"],
658 + max_num_patches=config.vision_config.args["max_num_patches"],
659 + norm_mean=config.norm_mean,
660 + norm_std=config.norm_std,
661 + )
662 +
663 + @staticmethod
664 + def use_dynamic_resolution(config: PretrainedConfig) -> bool:
665 + return "min_num_patches" in config.vision_config.args
666 +
667 @property
668 @abstractmethod
669 def image_token_id(self) -> int:
@@ -354,36 +717,61 @@ class BaseNanoNemotronVLProcessor(ABC):
717 text: list[str],
718 images: list[Image.Image],
719 max_num_tiles: int,
357 - ) -> tuple[list[str], dict[str, torch.Tensor]]:
720 + ) -> tuple[list[str], dict[str, Any]]:
721 if len(images) == 0:
722 image_inputs = {}
723 + return text, image_inputs
724 +
725 + if tiler := self.dynamic_tiler:
726 + sans_images = text[0].replace("<image>", "")
727 + text_prompt_length = len(
728 + self.tokenizer(sans_images, add_special_tokens=False).input_ids
729 + )
730 + pixel_values_lst, num_tokens_per_image = tiler._images_to_pixel_values_lst(
731 + text_prompt_length=text_prompt_length,
732 + images=images,
733 + )
734 + imgs_sizes = [(pv.shape[-2], pv.shape[-1]) for pv in pixel_values_lst]
735 + normalized = [
736 + input_conditioner(img, tiler.norm_mean, tiler.norm_std)
737 + for img in pixel_values_lst
738 + ]
739 + image_num_patches = torch.tensor([1] * len(num_tokens_per_image))
740 + image_inputs = {
741 + "pixel_values_flat": normalized,
742 + "imgs_sizes": imgs_sizes,
743 + "num_tokens_per_image": num_tokens_per_image,
744 + }
745 else:
746 pixel_values_lst = self._images_to_pixel_values_lst(images, max_num_tiles)
747 + image_num_patches = torch.tensor([len(item) for item in pixel_values_lst])
748 + pixel_values_flat = input_conditioner(
749 + torch.cat(pixel_values_lst), self.norm_mean, self.norm_std
750 + )
751 image_inputs = {
363 - "pixel_values_flat": input_conditioner(
364 - torch.cat(pixel_values_lst), self.norm_mean, self.norm_std
365 - ),
366 - "image_num_patches": torch.tensor(
367 - [len(item) for item in pixel_values_lst]
368 - ),
752 + "pixel_values_flat": pixel_values_flat,
753 + "image_num_patches": image_num_patches,
754 }
755 + num_tokens_per_image = [
756 + self.num_image_token * len(item) for item in pixel_values_lst
757 + ]
758
371 - assert len(text) == 1, (
372 - "hf_processor is called on the output of get_dummy_text, "
373 - "which should be a single string"
374 - )
375 - parts = [x for x in re.split(r"(<image>)", text[0]) if x]
376 - assert parts.count("<image>") == len(pixel_values_lst), (
377 - "the number of <image> tokens in the text should be the "
378 - "same as the number of images"
379 - )
759 + assert len(text) == 1, (
760 + "hf_processor is called on the output of get_dummy_text, "
761 + "which should be a single string"
762 + )
763 + parts = [x for x in re.split(r"(<image>)", text[0]) if x]
764 + assert parts.count("<image>") == len(pixel_values_lst), (
765 + "the number of <image> tokens in the text should be the "
766 + "same as the number of images"
767 + )
768
381 - for i, pixel_values in enumerate(pixel_values_lst):
382 - num_patches = pixel_values.shape[0]
383 - feature_size = num_patches * self.num_image_token
384 - image_repl = self.get_image_repl(feature_size, num_patches)
385 - parts[i] = parts[i].replace("<image>", image_repl.full)
386 - text = ["".join(parts)]
769 + for i, (feature_size, num_patches) in enumerate(
770 + zip(num_tokens_per_image, image_num_patches, strict=True)
771 + ):
772 + image_repl = self.get_image_repl(feature_size, num_patches)
773 + parts[i] = parts[i].replace("<image>", image_repl.full)
774 + text = ["".join(parts)]
775 return text, image_inputs
776
777 def _make_batch_input(self, input_item: Any | list[Any] | None = None):
@@ -393,6 +781,7 @@ class BaseNanoNemotronVLProcessor(ABC):
781 input_item = [input_item]
782 return input_item
783
784 + @abstractmethod
785 def __call__(
786 self,
787 text: str | list[str] | None = None,
@@ -400,23 +789,7 @@ class BaseNanoNemotronVLProcessor(ABC):
789 return_tensors: str | TensorType | None = None,
790 max_num_tiles: int | None = None,
791 ) -> BatchFeature:
403 - # Use default if not provided
404 - if max_num_tiles is None:
405 - max_num_tiles = self.max_num_tiles
406 -
407 - text, images = [self._make_batch_input(x) for x in (text, images)]
408 -
409 - text, image_inputs = self._preprocess_image(
410 - text=text,
411 - images=images,
412 - max_num_tiles=max_num_tiles,
413 - )
414 -
415 - text_inputs = self.tokenizer(text, add_special_tokens=False)
416 -
417 - combined_outputs = {**text_inputs, **image_inputs}
418 -
419 - return BatchFeature(combined_outputs, tensor_type=return_tensors)
792 + raise NotImplementedError
793
794
795 class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
@@ -431,20 +804,16 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
804 config: PretrainedConfig,
805 tokenizer: TokenizerLike,
806 *,
807 + max_model_len: int,
808 max_num_tiles: int | None = None,
435 - min_dynamic_patch: int | None = None,
436 - max_dynamic_patch: int | None = None,
437 - dynamic_image_size: bool | None = None,
809 video_token: str | None = None,
810 video_pruning_rate: float | None = None,
811 ) -> None:
812 super().__init__(
813 config=config,
814 tokenizer=tokenizer,
815 + max_model_len=max_model_len,
816 max_num_tiles=max_num_tiles,
445 - min_dynamic_patch=min_dynamic_patch,
446 - max_dynamic_patch=max_dynamic_patch,
447 - dynamic_image_size=dynamic_image_size,
817 )
818 # add extra video token for video processing
819 self.video_token = video_token
@@ -478,7 +847,6 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
847 self,
848 videos: list[npt.NDArray],
849 max_num_tiles: int,
481 - dynamic_image_size: bool | None = None,
850 ) -> list[torch.Tensor]:
851 return [
852 video_to_pixel_values(
@@ -495,7 +863,6 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
863 text: list[str],
864 videos: list[tuple[npt.NDArray, dict[str, Any]]],
865 max_num_tiles: int,
498 - dynamic_image_size: bool | None = None,
866 ):
867 if len(videos) == 0 or not self.supports_video:
868 video_inputs = {}
@@ -505,7 +872,6 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
872 pixel_values_lst_video = self._videos_to_pixel_values_lst(
873 videos_lst,
874 max_num_tiles=max_num_tiles,
508 - dynamic_image_size=dynamic_image_size,
875 )
876
877 # We use frame duration in milliseconds (as integer) to ensure
@@ -592,7 +958,6 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
958 videos: list[tuple[npt.NDArray, dict[str, Any]]] | None = None,
959 return_tensors: str | TensorType | None = None,
960 max_num_tiles: int | None = None,
595 - dynamic_image_size: bool | None = None,
961 ) -> BatchFeature:
962 # Use default if not provided
963 if max_num_tiles is None:
@@ -612,14 +977,23 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
977 text=text,
978 videos=videos,
979 max_num_tiles=1,
615 - dynamic_image_size=dynamic_image_size,
980 )
981
982 text_inputs = self.tokenizer(text, add_special_tokens=False)
983
620 - combined_outputs = {**text_inputs, **image_inputs, **video_inputs}
621 -
622 - return BatchFeature(combined_outputs, tensor_type=return_tensors)
984 + if self.dynamic_tiler is None:
985 + batch = BatchFeature(
986 + {**text_inputs, **video_inputs, **image_inputs},
987 + tensor_type=return_tensors,
988 + )
989 + else:
990 + batch = BatchFeature(
991 + {**text_inputs, **video_inputs}, tensor_type=return_tensors
992 + )
993 + # allow images to be exempt from the BatchFeature validation:
994 + # We will .stack() them in _parse_and_validate_image_input
995 + batch.update(image_inputs)
996 + return batch
997
998 def get_image_repl(
999 self,
@@ -722,23 +1096,6 @@ class BaseNanoNemotronVLProcessingInfo(BaseProcessingInfo):
1096 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
1097 return {"image": None}
1098
725 - def get_num_image_tokens(
726 - self,
727 - *,
728 - image_width: int,
729 - image_height: int,
730 - max_num_tiles: int,
731 - processor: BaseNanoNemotronVLProcessor | None,
732 - ) -> int:
733 - if processor is None:
734 - processor = self.get_hf_processor()
735 -
736 - return processor.get_num_image_tokens(
737 - image_width=image_width,
738 - image_height=image_height,
739 - max_num_tiles=max_num_tiles,
740 - )
741 -
1099 def get_image_size_with_most_features(self, max_num_tiles: int) -> ImageSize:
1100 processor = self.get_hf_processor()
1101
@@ -749,11 +1106,8 @@ class BaseNanoNemotronVLProcessingInfo(BaseProcessingInfo):
1106 for wr, hr in target_ratios:
1107 width, height = base_size * wr, base_size * hr
1108
752 - feat_size = self.get_num_image_tokens(
753 - image_width=width,
754 - image_height=height,
755 - max_num_tiles=max_num_tiles,
756 - processor=processor,
1109 + feat_size = processor.get_num_image_tokens(
1110 + image_width=width, image_height=height, max_num_tiles=max_num_tiles
1111 )
1112 if feat_size > largest_feature_size:
1113 largest_feature_size = feat_size
@@ -772,11 +1126,10 @@ class BaseNanoNemotronVLProcessingInfo(BaseProcessingInfo):
1126 max_num_tiles
1127 )
1128
775 - return self.get_num_image_tokens(
1129 + return processor.get_num_image_tokens(
1130 image_width=target_width,
1131 image_height=target_height,
1132 max_num_tiles=max_num_tiles,
779 - processor=processor,
1133 )
1134
1135
@@ -822,6 +1175,7 @@ class NanoNemotronVLProcessingInfo(BaseNanoNemotronVLProcessingInfo):
1175 tokenizer=self.get_tokenizer(),
1176 video_token=self.get_video_token(),
1177 video_pruning_rate=self.get_video_pruning_rate(),
1178 + max_model_len=self.ctx.model_config.max_model_len,
1179 **kwargs,
1180 )
1181
@@ -829,19 +1183,29 @@ class NanoNemotronVLProcessingInfo(BaseNanoNemotronVLProcessingInfo):
1183 class NanoNemotronBaseVLMultiModalProcessor(BaseMultiModalProcessor[_I]):
1184 """Basic image-only MultiModalProcessor for InternVL-style models."""
1185
1186 + @cached_property
1187 + def is_dynamic_tiler(self) -> bool:
1188 + return self.info.get_hf_processor().dynamic_tiler is not None
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]:
837 - image_num_patches = hf_inputs.get("image_num_patches", torch.empty(0))
1195 + if self.is_dynamic_tiler:
1196 + pixel_values_flat = MultiModalFieldConfig.batched("image")
1197 + else:
1198 + image_num_patches = hf_inputs.get("image_num_patches", torch.empty(0))
1199 + pixel_values_flat = MultiModalFieldConfig.flat_from_sizes(
1200 + "image", image_num_patches
1201 + )
1202
1203 return dict(
840 - pixel_values_flat=MultiModalFieldConfig.flat_from_sizes(
841 - "image", image_num_patches
842 - ),
1204 + pixel_values_flat=pixel_values_flat,
1205 image_num_patches=MultiModalFieldConfig.batched("image"),
1206 image_embeds=MultiModalFieldConfig.batched("image"),
1207 + num_tokens_per_image=MultiModalFieldConfig.batched("image"),
1208 + imgs_sizes=MultiModalFieldConfig.batched("image"),
1209 )
1210
1211 def _get_prompt_updates(
@@ -870,17 +1234,19 @@ class NanoNemotronBaseVLMultiModalProcessor(BaseMultiModalProcessor[_I]):
1234
1235 if isinstance(images, ImageEmbeddingItems):
1236 feature_size = images.get_feature_size(item_idx)
1237 + elif tiler := hf_processor.dynamic_tiler:
1238 + image = images.get(item_idx)
1239 + feature_size = tiler.get_cached_feature_size(image)
1240 else:
1241 image_size = images.get_image_size(item_idx)
1242 # Extract max_num_tiles from kwargs, default to 12
1243 max_num_tiles = hf_processor_mm_kwargs.get(
1244 "max_num_tiles", hf_processor.max_num_tiles
1245 )
879 - feature_size = self.info.get_num_image_tokens(
1246 + feature_size = hf_processor.get_num_image_tokens(
1247 image_width=image_size.width,
1248 image_height=image_size.height,
1249 max_num_tiles=max_num_tiles,
883 - processor=hf_processor,
1250 )
1251
1252 num_patches = None
@@ -1017,12 +1383,18 @@ class NanoNemotronVLDummyInputsBuilder(BaseDummyInputsBuilder[_I]):
1383 mm_counts: Mapping[str, int],
1384 mm_options: Mapping[str, BaseDummyOptions] | None = None,
1385 ) -> MultiModalDataDict:
1020 - # Use default max_num_tiles for dummy data generation
1021 - max_num_tiles = 12
1022 - target_width, target_height = self.info.get_image_size_with_most_features(
1023 - max_num_tiles
1024 - )
1386 num_images = mm_counts.get("image", 0)
1387 + processor = self.info.get_hf_processor()
1388 + if tiler := processor.dynamic_tiler:
1389 + budget = tiler.max_num_tokens_available(text_prompt_length=num_images)
1390 + target_width, target_height = (
1391 + tiler.width_and_height_for_max_num_tokens_available(budget)
1392 + )
1393 + else:
1394 + max_num_tiles = 12
1395 + target_width, target_height = self.info.get_image_size_with_most_features(
1396 + max_num_tiles
1397 + )
1398
1399 image_overrides = mm_options.get("image") if mm_options else None
1400
@@ -1181,6 +1553,11 @@ class NemotronH_Nano_VL_V2(
1553 self._img_context_token_ids = tokenizer.encode(
1554 IMG_CONTEXT, add_special_tokens=False
1555 )
1556 + self.dynamic_resolution = BaseNanoNemotronVLProcessor.use_dynamic_resolution(
1557 + config
1558 + )
1559 + if self.dynamic_resolution:
1560 + logger.info("Dynamic resolution is enabled for NanoNemotronVLProcessor")
1561
1562 def pixel_shuffle(self, x, scale_factor=0.5):
1563 n, w, h, c = x.size()
@@ -1211,7 +1588,51 @@ class NemotronH_Nano_VL_V2(
1588 x = x.permute(0, 2, 1, 3).contiguous()
1589 return x
1590
1214 - def extract_feature(self, pixel_values):
1591 + def pixel_shuffle_dynamic_res(
1592 + self, x: torch.Tensor, *, imgs_sizes: list[tuple[int, int]]
1593 + ) -> torch.Tensor:
1594 + scale_factor = self.downsample_ratio
1595 + patch_dim = self.patch_size
1596 + seq_lens = calc_seq_lens(imgs_sizes, patch_dim)
1597 + splits = torch.split(x, seq_lens, dim=-2)
1598 + out = []
1599 + for i, sv in enumerate(splits):
1600 + h = imgs_sizes[i][0] // patch_dim
1601 + w = imgs_sizes[i][1] // patch_dim
1602 + sv = sv.reshape(sv.shape[0], h, w, -1)
1603 +
1604 + n, h, w, c = sv.size()
1605 +
1606 + sv = sv.view(n, h, int(w * scale_factor), int(c / scale_factor))
1607 + sv = sv.permute(0, 2, 1, 3).contiguous()
1608 + sv = sv.view(
1609 + n,
1610 + int(w * scale_factor),
1611 + int(h * scale_factor),
1612 + int(c / (scale_factor * scale_factor)),
1613 + )
1614 +
1615 + if self.ps_version == "v2":
1616 + sv = sv.permute(0, 2, 1, 3).contiguous()
1617 +
1618 + sv = sv.reshape(sv.shape[0], -1, sv.shape[-1])
1619 + out.append(sv)
1620 +
1621 + x = torch.cat(out, dim=-2)
1622 +
1623 + return x
1624 +
1625 + def extract_feature_dynamic(
1626 + self, pixel_values: torch.Tensor, imgs_sizes: list[tuple[int, int]]
1627 + ):
1628 + """Dynamic resolution extract_feature for images."""
1629 + _, vit_embeds = self.vision_model(pixel_values, imgs_sizes=imgs_sizes)
1630 + vit_embeds = vit_embeds.to(dtype=torch.bfloat16)
1631 + vit_embeds = self.pixel_shuffle_dynamic_res(vit_embeds, imgs_sizes=imgs_sizes)
1632 + vit_embeds = self.mlp1(vit_embeds)
1633 + return vit_embeds
1634 +
1635 + def extract_feature(self, pixel_values: torch.Tensor):
1636 # Process images in a micro-batch of at most 128 frames per call
1637 # This is done on purpose to ensure peak GPU ram usage of huge batch
1638 # (namely for really long videos with EVS ON) won't cause any problems
@@ -1239,36 +1660,39 @@ class NemotronH_Nano_VL_V2(
1660 def _parse_and_validate_image_input(
1661 self, **kwargs: object
1662 ) -> NanoNemotronVLImageInputs | None:
1242 - pixel_values_flat = kwargs.pop("pixel_values_flat", None)
1243 - image_num_patches = kwargs.pop("image_num_patches", None)
1244 - image_embeds = kwargs.pop("image_embeds", None)
1245 -
1246 - if pixel_values_flat is None and image_embeds is None:
1247 - return None
1248 -
1249 - if image_embeds is not None:
1663 + if image_embeds := kwargs.pop("image_embeds", None):
1664 return NanoNemotronVLImageEmbeddingInputs(
1665 type="image_embeds",
1666 data=image_embeds,
1667 )
1668
1255 - if pixel_values_flat is not None:
1256 - return NanoNemotronVLImagePixelInputs(
1257 - type="pixel_values",
1258 - pixel_values_flat=pixel_values_flat,
1259 - num_patches=image_num_patches,
1669 + if self.dynamic_resolution:
1670 + pixel_values_flat = DynamicResolutionImageTiler.stack(
1671 + kwargs.pop("pixel_values_flat"), self.patch_size
1672 )
1673 + return NanoNemotronVLImagePixelInputsDynamic(
1674 + pixel_values_flat=pixel_values_flat, **kwargs
1675 + )
1676 + else:
1677 + return NanoNemotronVLImagePixelInputs(**kwargs)
1678
1262 - raise AssertionError("This line should be unreachable.")
1263 -
1264 - def _process_image_input(
1265 - self, image_input: NanoNemotronVLImageInputs
1679 + def _process_image_input_dynamic(
1680 + self, image_input: NanoNemotronVLImagePixelInputsDynamic
1681 ) -> tuple[torch.Tensor, ...]:
1267 - if image_input["type"] == "image_embeds":
1268 - return image_input["data"]
1682 + image_embeds = self.extract_feature_dynamic(
1683 + image_input.pixel_values_flat, image_input.imgs_sizes
1684 + )
1685 + num_tokens_per_image = image_input.num_tokens_per_image
1686 +
1687 + if len(num_tokens_per_image) == 1:
1688 + return (image_embeds.view(-1, self.config.text_config.hidden_size),)
1689
1270 - assert self.vision_model is not None
1690 + image_embeds = image_embeds.view(-1, self.config.text_config.hidden_size)
1691 + return image_embeds.split(num_tokens_per_image)
1692
1693 + def _process_image_input(
1694 + self, image_input: NanoNemotronVLImagePixelInputs
1695 + ) -> tuple[torch.Tensor, ...]:
1696 image_embeds = self.extract_feature(image_input["pixel_values_flat"])
1697 num_patches = image_input["num_patches"]
1698
@@ -1470,7 +1894,13 @@ class NemotronH_Nano_VL_V2(
1894 for modality in modalities:
1895 if modality == "images":
1896 image_input = modalities["images"]
1473 - image_embeddings = self._process_image_input(image_input)
1897 + if image_input["type"] == "image_embeds":
1898 + image_embeddings = image_input["data"]
1899 + elif self.dynamic_resolution:
1900 + assert image_input["type"] == "pixel_values_dynamic"
1901 + image_embeddings = self._process_image_input_dynamic(image_input)
1902 + else:
1903 + image_embeddings = self._process_image_input(image_input)
1904 multimodal_embeddings += tuple(image_embeddings)
1905 if modality == "videos":
1906 video_input = modalities["videos"]
@@ -1652,33 +2082,6 @@ class NemotronH_Nano_VL_V2(
2082 if save_to_file and sys.stdout != original_stdout:
2083 sys.stdout = original_stdout
2084
1655 - def get_model_info(self):
1656 - """
1657 - Get basic model information as a dictionary.
1658 - """
1659 - total_params = sum(p.numel() for p in self.parameters())
1660 -
1661 - component_info = {}
1662 - for name, param in self.named_parameters():
1663 - component = name.split(".")[0]
1664 - if component not in component_info:
1665 - component_info[component] = {"params": 0, "size": 0}
1666 - component_info[component]["params"] += 1
1667 - component_info[component]["size"] += param.numel()
1668 -
1669 - return {
1670 - "model_name": "NemotronH_Nano_VL_V2",
1671 - "total_parameters": total_params,
1672 - "memory_estimate_mb": total_params * 2 / (1024**2), # bfloat16
1673 - "components": component_info,
1674 - "config": {
1675 - "image_size": getattr(self.config, "force_image_size", None),
1676 - "patch_size": getattr(self.config, "patch_size", None),
1677 - "num_image_token": self.num_image_token,
1678 - "downsample_ratio": self.downsample_ratio,
1679 - },
1680 - }
1681 -
2085 def get_vit_model_from_radio_config(self, hf_config):
2086 hf_config_vision = hf_config.vision_config
2087 model_name = hf_config_vision.args.get("model")
vllm/model_executor/models/radio.py
+205 -21
@@ -21,7 +21,11 @@ from transformers import PretrainedConfig
21
22 from vllm.model_executor.layers.quantization import QuantizationConfig
23 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
24 -from vllm.model_executor.models.intern_vit import InternVisionEncoder
24 +from vllm.model_executor.models.intern_vit import (
25 + InternParallelAttention,
26 + InternVisionEncoder,
27 + InternVisionEncoderLayer,
28 +)
29
30 input_dim_t: TypeAlias = int | tuple[int, int]
31 norm_t: TypeAlias = tuple[float, float, float] | torch.Tensor
@@ -43,6 +47,15 @@ to_4tuple = _ntuple(4)
47 to_ntuple = _ntuple
48
49
50 +def calc_seq_len(size: tuple[int, int], patch_size: int) -> int:
51 + h, w = size
52 + return (h // patch_size) * (w // patch_size)
53 +
54 +
55 +def calc_seq_lens(sizes: list[tuple[int, int]], patch_size: int) -> list[int]:
56 + return [calc_seq_len(size, patch_size) for size in sizes]
57 +
58 +
59 class ClsToken(nn.Module):
60 def __init__(
61 self,
@@ -164,15 +177,73 @@ class ViTPatchGenerator(nn.Module):
177 nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity()
178 )
179
167 - def forward(self, x: torch.Tensor) -> torch.Tensor:
168 - patches = self.embed_patches(x)
169 - patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:])
170 - patches = self.cls_token(patches)
180 + def forward(
181 + self, x: torch.Tensor, imgs_sizes: list[tuple[int, int]] | None = None
182 + ) -> torch.Tensor:
183 + if imgs_sizes is not None:
184 + patches = self.embedder(x)
185 + patches, pos_enc = self.apply_pos_enc_dynamic(
186 + patches, imgs_sizes=imgs_sizes
187 + )
188 + patches = self.cls_token_dynamic(patches, imgs_sizes=imgs_sizes)
189 + else:
190 + patches = self.embed_patches(x)
191 + patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:])
192 + patches = self.cls_token(patches)
193 patches = self.patch_normalizer(patches)
194 if self.return_pos_enc:
195 return patches, pos_enc
196 return patches
197
198 + def apply_pos_enc_dynamic(
199 + self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]]
200 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
201 + if not self.abs_pos:
202 + return patches, None
203 +
204 + current_length = 0
205 + pos_enc_list = []
206 +
207 + for size in imgs_sizes:
208 + seq_length = calc_seq_len(size, self.patch_size)
209 +
210 + img_patches = patches[:, current_length : current_length + seq_length, :]
211 + pos_enc = self.get_pos_enc(patches.shape[0], input_size=size)
212 + img_patches_with_pos = img_patches + pos_enc
213 +
214 + patches = torch.cat(
215 + [
216 + patches[:, :current_length, :],
217 + img_patches_with_pos,
218 + patches[:, current_length + seq_length :, :],
219 + ],
220 + dim=1,
221 + )
222 + pos_enc_list.append(pos_enc)
223 + current_length += seq_length
224 +
225 + full_pos_enc = torch.cat(pos_enc_list, dim=1) if pos_enc_list else None
226 + return patches, full_pos_enc
227 +
228 + def cls_token_dynamic(
229 + self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]]
230 + ) -> torch.Tensor:
231 + if not self.cls_token.enabled:
232 + return patches
233 +
234 + out = []
235 + current_length = 0
236 +
237 + for seq_len in calc_seq_lens(imgs_sizes, self.patch_size):
238 + class_token = self.cls_token.token.unsqueeze(0).expand(
239 + patches.shape[0], -1, -1
240 + )
241 + out.append(class_token)
242 + out.append(patches[:, current_length : current_length + seq_len, :])
243 + current_length += seq_len
244 +
245 + return torch.cat(out, dim=1)
246 +
247 @property
248 def apply_cls_token(self):
249 return self.cls_token.enabled
@@ -406,6 +477,66 @@ class ViTPatchLinear(nn.Linear):
477 self.patch_size = patch_size
478
479
480 +class RadioParallelAttention(InternParallelAttention):
481 + def forward(
482 + self, x: torch.Tensor, attn_mask: torch.Tensor | None = None
483 + ) -> torch.Tensor:
484 + if attn_mask is None:
485 + return super().forward(x)
486 +
487 + B, N, _ = x.shape
488 + qkv, _ = self.qkv(x)
489 + q, k, v = qkv.chunk(3, dim=-1)
490 +
491 + if self.qk_normalization:
492 + q, k = self._apply_qk_norm(q, k)
493 +
494 + q = q.view(B, N, self.num_heads_per_partition, self.head_dim)
495 + k = k.view(B, N, self.num_heads_per_partition, self.head_dim)
496 + v = v.view(B, N, self.num_heads_per_partition, self.head_dim)
497 + q, k, v = (t.transpose(1, 2) for t in (q, k, v))
498 + out = F.scaled_dot_product_attention(
499 + q, k, v, attn_mask=attn_mask, scale=self.scale
500 + )
501 + out = out.transpose(1, 2).reshape(B, N, -1)
502 + out, _ = self.proj(out)
503 + return out
504 +
505 +
506 +class RadioVisionEncoderLayer(InternVisionEncoderLayer):
507 + def __init__(self, *args, **kwargs) -> None:
508 + super().__init__(*args, attn_cls=RadioParallelAttention, **kwargs)
509 +
510 + def forward(
511 + self,
512 + hidden_states: torch.Tensor,
513 + attn_mask: torch.Tensor | None = None,
514 + ):
515 + hidden_states = (
516 + hidden_states
517 + + self.attn(self.norm1(hidden_states), attn_mask=attn_mask) * self.ls1
518 + )
519 +
520 + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) * self.ls2
521 +
522 + return hidden_states
523 +
524 +
525 +class RadioVisionEncoder(InternVisionEncoder):
526 + def __init__(self, *args, **kwargs) -> None:
527 + super().__init__(*args, layer_cls=RadioVisionEncoderLayer, **kwargs)
528 +
529 + def forward(
530 + self,
531 + inputs_embeds: torch.Tensor,
532 + attn_mask: torch.Tensor | None = None,
533 + ):
534 + hidden_states = inputs_embeds
535 + for encoder_layer in self.layers:
536 + hidden_states = encoder_layer(hidden_states, attn_mask=attn_mask)
537 + return hidden_states
538 +
539 +
540 class RadioInternVisionModel(nn.Module):
541 packed_modules_mapping = {
542 "qkv": ["qkv"],
@@ -440,7 +571,7 @@ class RadioInternVisionModel(nn.Module):
571 register_multiple=config.register_multiple,
572 )
573
443 - self.encoder = InternVisionEncoder(
574 + self.encoder = RadioVisionEncoder(
575 config=config,
576 quant_config=quant_config,
577 num_hidden_layers_override=num_hidden_layers_override,
@@ -459,10 +590,45 @@ class RadioInternVisionModel(nn.Module):
590 def get_input_embeddings(self):
591 return self.embeddings
592
462 - def forward(self, x: torch.Tensor) -> torch.FloatTensor:
593 + def create_inter_image_attention_mask(
594 + self, imgs_sizes: list[tuple[int, int]], device: torch.device
595 + ) -> torch.Tensor:
596 + patch_size = self.patch_generator.patch_size
597 + num_skip = self.patch_generator.num_skip
598 +
599 + seq_lens = calc_seq_lens(imgs_sizes, patch_size)
600 + patch_counts = [seq_len + num_skip for seq_len in seq_lens]
601 + total_patches = sum(patch_counts)
602 +
603 + # Create attention mask - default to False (mask out)
604 + mask = torch.zeros(
605 + total_patches, total_patches, dtype=torch.bool, device=device
606 + )
607 +
608 + # Each image's patches can only attend to patches from the same image
609 + start_idx = 0
610 + for patch_count in patch_counts:
611 + end_idx = start_idx + patch_count
612 + # Allow attention within this image's patches
613 + mask[start_idx:end_idx, start_idx:end_idx] = True
614 + start_idx = end_idx
615 +
616 + return mask
617 +
618 + def forward(
619 + self,
620 + x: torch.Tensor,
621 + imgs_sizes: torch.Tensor | None = None,
622 + ) -> torch.FloatTensor:
623 assert self.patch_generator is not None
464 - hidden_states = self.patch_generator(x)
465 - encoder_outputs = self.encoder(inputs_embeds=hidden_states)
624 + hidden_states = self.patch_generator(x, imgs_sizes=imgs_sizes)
625 + attn_mask = None
626 + if imgs_sizes is not None and len(imgs_sizes) > 1:
627 + # Dynamic Resolution
628 + attn_mask = self.create_inter_image_attention_mask(
629 + imgs_sizes, device=x.device
630 + )
631 + encoder_outputs = self.encoder(inputs_embeds=hidden_states, attn_mask=attn_mask)
632 return encoder_outputs
633
634
@@ -504,9 +670,11 @@ class RadioModel(nn.Module):
670 self,
671 pixel_values: torch.Tensor | None = None,
672 pixel_embeds: torch.Tensor | None = None,
673 + *,
674 + imgs_sizes: torch.Tensor | None = None,
675 ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
508 - y = self.model(pixel_values)
509 - return self._extract_final(y)
676 + y = self.model(pixel_values, imgs_sizes=imgs_sizes)
677 + return self._extract_final(y, imgs_sizes=imgs_sizes)
678
679 def load_weights(self, weights) -> set[str]:
680 loaded_params: set[str] = set()
@@ -558,16 +726,32 @@ class RadioModel(nn.Module):
726 return loaded_params
727
728 def _extract_final(
561 - self, y: torch.Tensor
729 + self, y: torch.Tensor, imgs_sizes: list[tuple[int, int]] | None = None
730 ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
731 # Remove CLS + REGISTERS tokens
564 - patch_gen = getattr(self.model, "patch_generator", None)
565 - if patch_gen is not None:
566 - all_summary = y[:, : patch_gen.num_cls_tokens]
567 - if self.summary_idxs is not None:
568 - bb_summary = all_summary[:, self.summary_idxs]
569 - else:
570 - bb_summary = all_summary
571 - all_feat = y[:, patch_gen.num_skip :]
572 -
732 + num_skip = self.model.patch_generator.num_skip
733 + patch_size = self.model.patch_generator.patch_size
734 + num_cls_tokens = self.model.patch_generator.num_cls_tokens
735 + if imgs_sizes is None:
736 + all_summary = y[:, :num_cls_tokens]
737 + all_feat = y[:, num_skip:]
738 + else:
739 + all_patches = []
740 + summaries = []
741 + current_pos = 0
742 + for num_patches in calc_seq_lens(imgs_sizes, patch_size):
743 + patches = y[
744 + :, current_pos + num_skip : current_pos + num_skip + num_patches, :
745 + ]
746 + all_patches.append(patches)
747 + summary = y[:, current_pos : current_pos + num_cls_tokens, :]
748 + summaries.append(summary)
749 + current_pos += num_skip + num_patches
750 + all_summary = torch.cat(summaries, dim=1)
751 + all_feat = torch.cat(all_patches, dim=1)
752 +
753 + if self.summary_idxs is not None:
754 + bb_summary = all_summary[:, self.summary_idxs]
755 + else:
756 + bb_summary = all_summary
757 return bb_summary.flatten(1), all_feat