| 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 | import warnings |
| 12 | from collections.abc import Iterable, Mapping, Sequence |
| 13 | from functools import cached_property |
| 14 | from io import BytesIO |
| 15 | from typing import Annotated, Literal, TypeAlias, TypedDict, cast |
| 16 | |
| 17 | import torch |
| 18 | import torch.nn as nn |
| 19 | from transformers import BatchFeature, PreTrainedConfig |
| 20 | |
| 21 | from vllm import envs |
| 22 | from vllm.config import VllmConfig |
| 23 | from vllm.config.multimodal import ( |
| 24 | MultiModalDummyOptions, |
| 25 | VideoDummyOptions, |
| 26 | ) |
| 27 | from vllm.inputs import MultiModalDataDict, MultiModalInput |
| 28 | from vllm.logger import init_logger |
| 29 | from vllm.lora.layers.base import BaseLayerWithLoRA |
| 30 | from vllm.model_executor.layers.activation import ReLUSquaredActivation |
| 31 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 32 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 33 | from vllm.model_executor.models.interfaces import ( |
| 34 | HasInnerState, |
| 35 | IsHybrid, |
| 36 | MultiModalEmbeddings, |
| 37 | SupportsEagle3, |
| 38 | SupportsLoRA, |
| 39 | SupportsMultiModal, |
| 40 | SupportsMultiModalPruning, |
| 41 | ) |
| 42 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 43 | from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM |
| 44 | from vllm.model_executor.models.parakeet import ParakeetExtractor, ProjectedParakeet |
| 45 | from vllm.model_executor.models.radio import RadioModel, calc_seq_lens |
| 46 | from vllm.model_executor.models.utils import ( |
| 47 | WeightsMapper, |
| 48 | init_vllm_registered_model, |
| 49 | maybe_prefix, |
| 50 | ) |
| 51 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 52 | from vllm.multimodal.inputs import ( |
| 53 | AudioItem, |
| 54 | BatchedTensorInputs, |
| 55 | MultiModalFieldConfig, |
| 56 | MultiModalKwargsItems, |
| 57 | VideoItem, |
| 58 | ) |
| 59 | from vllm.multimodal.media.audio import load_audio_pyav |
| 60 | from vllm.multimodal.parse import ( |
| 61 | AudioProcessorItems, |
| 62 | ImageEmbeddingItems, |
| 63 | ImageProcessorItems, |
| 64 | ImageSize, |
| 65 | MultiModalDataItems, |
| 66 | MultiModalDataParser, |
| 67 | VideoProcessorItems, |
| 68 | ) |
| 69 | from vllm.multimodal.processing import ( |
| 70 | BaseDummyInputsBuilder, |
| 71 | ProcessorInputs, |
| 72 | TimingContext, |
| 73 | ) |
| 74 | from vllm.multimodal.processing.processor import ( |
| 75 | BaseMultiModalProcessor, |
| 76 | BaseProcessingInfo, |
| 77 | HFMultiModalInputs, |
| 78 | PromptReplacement, |
| 79 | PromptUpdate, |
| 80 | cached_encode, |
| 81 | ) |
| 82 | from vllm.multimodal.video_prune.evs import ( |
| 83 | compute_retained_tokens_count, |
| 84 | compute_retention_mask, |
| 85 | ) |
| 86 | from vllm.renderers import TokenizeParams |
| 87 | from vllm.sequence import IntermediateTensors |
| 88 | from vllm.tokenizers import cached_tokenizer_from_config |
| 89 | from vllm.transformers_utils.configs.radio import RadioConfig |
| 90 | from vllm.transformers_utils.processors.internvl import get_internvl_target_ratios |
| 91 | from vllm.transformers_utils.processors.nano_nemotron_vl import ( |
| 92 | AUDIO_CONTEXT, |
| 93 | IMG_CONTEXT, |
| 94 | IMG_END, |
| 95 | IMG_START, |
| 96 | BaseNanoNemotronVLProcessor, |
| 97 | DynamicResolutionImageTiler, |
| 98 | NanoNemotronVLProcessor, |
| 99 | get_video_target_size_and_feature_size, |
| 100 | ) |
| 101 | from vllm.utils.tensor_schema import TensorSchema, TensorShape |
| 102 | |
| 103 | from .utils import _merge_multimodal_embeddings |
| 104 | |
| 105 | logger = init_logger(__name__) |
| 106 | |
| 107 | MAX_AUDIO_LEN_S = 10 * 60 # 10 minutes |
| 108 | |
| 109 | |
| 110 | class NanoNemotronVLAudioFeatureInputs(TensorSchema): |
| 111 | """Dimensions: |
| 112 | - c: Number of audio clips (possibly flattened across audio items) |
| 113 | - b: Number of original audio items |
| 114 | - t: Audio feature length |
| 115 | - f: Feature size (mel bins) |
| 116 | """ |
| 117 | |
| 118 | type: Literal["audio_features"] = "audio_features" |
| 119 | input_audio_features: Annotated[torch.Tensor, TensorShape("c", "t", "f")] |
| 120 | feature_attention_mask: Annotated[torch.Tensor, TensorShape("c", "t")] |
| 121 | audio_num_clips: list[int] | torch.Tensor |
| 122 | |
| 123 | |
| 124 | class NanoNemotronVLImagePixelInputs(TensorSchema): |
| 125 | """Dimensions: |
| 126 | - bn: Batch size * number of images |
| 127 | - bnp: Batch size * number of images * (1 + num_patches) |
| 128 | - c: Number of channels (3) |
| 129 | - h: Height of each image patch |
| 130 | - w: Width of each image patch |
| 131 | """ |
| 132 | |
| 133 | type: Literal["pixel_values"] = "pixel_values" |
| 134 | pixel_values_flat: Annotated[torch.Tensor, TensorShape("bnp", 3, "h", "w")] |
| 135 | num_patches: Annotated[torch.Tensor, TensorShape("bn")] |
| 136 | |
| 137 | |
| 138 | class NanoNemotronVLImagePixelInputsDynamic(TensorSchema): |
| 139 | """Dynamic-resolution image inputs. |
| 140 | |
| 141 | imgs_sizes: per-image (height, width) in pixels. |
| 142 | num_tokens_per_image: per-image number of embedding tokens (post downsample). |
| 143 | """ |
| 144 | |
| 145 | type: Literal["pixel_values_dynamic"] = "pixel_values_dynamic" |
| 146 | pixel_values_flat: Annotated[torch.Tensor, TensorShape("bn", "h", "w")] |
| 147 | imgs_sizes: list[tuple[int, int]] |
| 148 | num_tokens_per_image: list[int] |
| 149 | |
| 150 | |
| 151 | class NanoNemotronVLImageEmbeddingInputs(TensorSchema): |
| 152 | """Dimensions: |
| 153 | - n: Number of images |
| 154 | - f: Total image feature size |
| 155 | - h: Hidden size (must match the hidden size of language model backbone) |
| 156 | """ |
| 157 | |
| 158 | type: Literal["image_embeds"] |
| 159 | data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("n", "f", "h")] |
| 160 | |
| 161 | |
| 162 | NanoNemotronVLImageInputs: TypeAlias = ( |
| 163 | NanoNemotronVLImagePixelInputs |
| 164 | | NanoNemotronVLImagePixelInputsDynamic |
| 165 | | NanoNemotronVLImageEmbeddingInputs |
| 166 | ) |
| 167 | |
| 168 | |
| 169 | class NanoNemotronVLVideoPixelInputs(TensorSchema): |
| 170 | """Dimensions: |
| 171 | - bvf: Batch size * number of videos * num_frames |
| 172 | - bn: Batch size * number of videos |
| 173 | - f: Number of frames |
| 174 | - c: Number of channels (3) |
| 175 | - h: Height of each video frame |
| 176 | - w: Width of each video frame |
| 177 | """ |
| 178 | |
| 179 | type: Literal["pixel_values_videos"] |
| 180 | pixel_values_flat: Annotated[torch.Tensor, TensorShape("bvf", 3, "h", "w")] |
| 181 | num_patches: Annotated[torch.Tensor, TensorShape("bn")] |
| 182 | frames_indices: Annotated[torch.Tensor, TensorShape("bvf")] |
| 183 | frame_duration_ms: Annotated[torch.Tensor, TensorShape("bn")] |
| 184 | |
| 185 | |
| 186 | class NanoNemotronVLVideoEmbeddingInputs(TensorSchema): |
| 187 | """Dimensions: |
| 188 | - n: Number of videos |
| 189 | - f: Total video feature size |
| 190 | - h: Hidden size (must match the hidden size of language model backbone) |
| 191 | """ |
| 192 | |
| 193 | type: Literal["video_embeds"] |
| 194 | data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("n", "f", "h")] |
| 195 | |
| 196 | |
| 197 | NanoNemotronVLVideoInputs: TypeAlias = ( |
| 198 | NanoNemotronVLVideoPixelInputs | NanoNemotronVLVideoEmbeddingInputs |
| 199 | ) |
| 200 | |
| 201 | |
| 202 | class NanoNemotronVLMultiModalInputs(TypedDict, total=False): |
| 203 | images: NanoNemotronVLImageInputs | None |
| 204 | videos: NanoNemotronVLVideoInputs | None |
| 205 | audios: NanoNemotronVLAudioFeatureInputs | None |
| 206 | |
| 207 | |
| 208 | class NanoNemotronVLProcessingInfo(BaseProcessingInfo): |
| 209 | def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor: |
| 210 | # Not accepted by NanoNemotronVLProcessor.__init__ |
| 211 | kwargs.pop("truncation", None) |
| 212 | |
| 213 | return self.ctx.init_processor( |
| 214 | NanoNemotronVLProcessor, |
| 215 | config=self.get_hf_config(), |
| 216 | tokenizer=self.get_tokenizer(), |
| 217 | video_token=self.get_video_token(), |
| 218 | video_pruning_rate=self.get_video_pruning_rate(), |
| 219 | max_model_len=self.ctx.model_config.max_model_len, |
| 220 | **kwargs, |
| 221 | ) |
| 222 | |
| 223 | @cached_property |
| 224 | def is_dynamic_tiler(self) -> bool: |
| 225 | return BaseNanoNemotronVLProcessor.use_dynamic_resolution(self.get_hf_config()) |
| 226 | |
| 227 | @property |
| 228 | def supports_video(self): |
| 229 | return True |
| 230 | |
| 231 | @property |
| 232 | def supports_audio(self) -> bool: |
| 233 | return self.sound_config is not None |
| 234 | |
| 235 | def get_video_token(self) -> str | None: |
| 236 | return IMG_CONTEXT |
| 237 | |
| 238 | def get_video_pruning_rate(self) -> float | None: |
| 239 | return self.ctx.get_mm_config().video_pruning_rate |
| 240 | |
| 241 | @property |
| 242 | def sound_config(self) -> PreTrainedConfig | None: |
| 243 | return getattr(self.get_hf_config(), "sound_config", None) |
| 244 | |
| 245 | def get_default_tok_params(self) -> TokenizeParams: |
| 246 | return super().get_default_tok_params().with_kwargs(add_special_tokens=False) |
| 247 | |
| 248 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 249 | image_limit = {"image": None} |
| 250 | video_limit = {"video": None} if self.supports_video else {} |
| 251 | audio_limit = {"audio": None} if self.supports_audio else {} |
| 252 | return {**image_limit, **video_limit, **audio_limit} |
| 253 | |
| 254 | def get_data_parser(self): |
| 255 | target_sr = None |
| 256 | target_channels = None |
| 257 | if self.sound_config: |
| 258 | target_sr = self.sound_config.sampling_rate |
| 259 | target_channels = 1 |
| 260 | |
| 261 | return MultiModalDataParser( |
| 262 | video_needs_metadata=True, |
| 263 | target_sr=target_sr, |
| 264 | target_channels=target_channels, |
| 265 | expected_hidden_size=self._get_expected_hidden_size(), |
| 266 | ) |
| 267 | |
| 268 | def get_image_size_with_most_features(self, max_num_tiles: int) -> ImageSize: |
| 269 | processor = self.get_hf_processor() |
| 270 | |
| 271 | base_size = processor.image_size |
| 272 | target_ratios = get_internvl_target_ratios(1, max_num_tiles) |
| 273 | |
| 274 | largest_feature_size, largest_feature_pinpoint = 0, None |
| 275 | for wr, hr in target_ratios: |
| 276 | width, height = base_size * wr, base_size * hr |
| 277 | |
| 278 | feat_size = processor.get_num_image_tokens( |
| 279 | image_width=width, image_height=height, max_num_tiles=max_num_tiles |
| 280 | ) |
| 281 | if feat_size > largest_feature_size: |
| 282 | largest_feature_size = feat_size |
| 283 | largest_feature_pinpoint = ImageSize(width=width, height=height) |
| 284 | |
| 285 | if largest_feature_size == 0 or largest_feature_pinpoint is None: |
| 286 | raise ValueError("Cannot have a largest feature size of 0!") |
| 287 | |
| 288 | return largest_feature_pinpoint |
| 289 | |
| 290 | def get_max_image_tokens(self) -> int: |
| 291 | processor = self.get_hf_processor() |
| 292 | # Use default max_num_tiles for max tokens calculation |
| 293 | max_num_tiles = processor.max_num_tiles |
| 294 | target_width, target_height = self.get_image_size_with_most_features( |
| 295 | max_num_tiles |
| 296 | ) |
| 297 | |
| 298 | return processor.get_num_image_tokens( |
| 299 | image_width=target_width, |
| 300 | image_height=target_height, |
| 301 | max_num_tiles=max_num_tiles, |
| 302 | ) |
| 303 | |
| 304 | def get_dummy_image_size_and_max_tokens( |
| 305 | self, mm_counts: Mapping[str, int] |
| 306 | ) -> tuple[tuple[int, int], int]: |
| 307 | processor = self.get_hf_processor() |
| 308 | num_images = mm_counts.get("image", 0) |
| 309 | |
| 310 | if tiler := processor.dynamic_tiler: |
| 311 | budget = tiler.max_num_tokens_available(text_prompt_length=num_images) |
| 312 | target_width, target_height = ( |
| 313 | tiler.width_and_height_for_max_num_tokens_available(budget) |
| 314 | ) |
| 315 | return ( |
| 316 | (target_width, target_height), |
| 317 | tiler._get_num_embeddings(target_width, target_height), |
| 318 | ) |
| 319 | |
| 320 | max_num_tiles = processor.max_num_tiles |
| 321 | target_width, target_height = self.get_image_size_with_most_features( |
| 322 | max_num_tiles |
| 323 | ) |
| 324 | return ( |
| 325 | (target_width, target_height), |
| 326 | processor.get_num_image_tokens( |
| 327 | image_width=target_width, |
| 328 | image_height=target_height, |
| 329 | max_num_tiles=max_num_tiles, |
| 330 | ), |
| 331 | ) |
| 332 | |
| 333 | def get_num_frames_with_most_features( |
| 334 | self, |
| 335 | seq_len: int, |
| 336 | mm_counts: Mapping[str, int], |
| 337 | ) -> int: |
| 338 | max_images = mm_counts.get("image", 0) |
| 339 | max_videos = mm_counts.get("video", 0) |
| 340 | |
| 341 | processor = self.get_hf_processor() # we get the CustomProcessor here |
| 342 | T = processor.video_temporal_patch_size |
| 343 | |
| 344 | max_image_tokens = self.get_max_image_tokens() * max_images |
| 345 | tokens_per_tubelet = processor.num_video_token |
| 346 | max_total_tubelets = (seq_len - max_image_tokens) // tokens_per_tubelet |
| 347 | max_tubelets_per_video = max_total_tubelets // max(max_videos, 1) |
| 348 | max_frames_per_video = max_tubelets_per_video * T |
| 349 | return max(max_frames_per_video, 1) |
| 350 | |
| 351 | def get_mm_max_tokens_per_item( |
| 352 | self, seq_len: int, mm_counts: Mapping[str, int] |
| 353 | ) -> Mapping[str, int]: |
| 354 | mm_max_tokens: dict[str, int] = {} |
| 355 | |
| 356 | if mm_counts.get("image", 0) > 0: |
| 357 | _, mm_max_tokens["image"] = self.get_dummy_image_size_and_max_tokens( |
| 358 | mm_counts |
| 359 | ) |
| 360 | |
| 361 | if mm_counts.get("video", 0) > 0: |
| 362 | assert self.supports_video |
| 363 | mm_max_tokens["video"] = seq_len |
| 364 | |
| 365 | if mm_counts.get("audio", 0) > 0: |
| 366 | assert self.supports_audio |
| 367 | mm_max_tokens["audio"] = seq_len |
| 368 | |
| 369 | return mm_max_tokens |
| 370 | |
| 371 | |
| 372 | class NanoNemotronVLMultiModalProcessor( |
| 373 | BaseMultiModalProcessor[NanoNemotronVLProcessingInfo] |
| 374 | ): |
| 375 | def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: |
| 376 | return self.dummy_inputs.get_dummy_text(mm_counts) |
| 377 | |
| 378 | def _get_hf_mm_inputs( |
| 379 | self, |
| 380 | mm_items: MultiModalDataItems, |
| 381 | hf_kwargs: Mapping[str, object], |
| 382 | ) -> HFMultiModalInputs: |
| 383 | hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs) |
| 384 | |
| 385 | # NanoNemotronVLProcessor expects "audios" instead of "audio" |
| 386 | hf_data = hf_inputs.hf_data |
| 387 | if "audio" in hf_data: |
| 388 | hf_data["audios"] = hf_data.pop("audio") |
| 389 | |
| 390 | return hf_inputs |
| 391 | |
| 392 | def _get_image_fields_config(self, hf_inputs: BatchFeature): |
| 393 | if self.info.is_dynamic_tiler: |
| 394 | pixel_values_flat = MultiModalFieldConfig.batched("image") |
| 395 | else: |
| 396 | image_num_patches = hf_inputs.get("image_num_patches", torch.empty(0)) |
| 397 | pixel_values_flat = MultiModalFieldConfig.flat_from_sizes( |
| 398 | "image", image_num_patches |
| 399 | ) |
| 400 | |
| 401 | return dict( |
| 402 | pixel_values_flat=pixel_values_flat, |
| 403 | image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 404 | image_embeds=MultiModalFieldConfig.batched("image"), |
| 405 | num_tokens_per_image=MultiModalFieldConfig.batched( |
| 406 | "image", keep_on_cpu=True |
| 407 | ), |
| 408 | imgs_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 409 | ) |
| 410 | |
| 411 | def _get_video_fields_config(self, hf_inputs: BatchFeature): |
| 412 | video_num_patches = hf_inputs.get("video_num_patches", torch.empty(0)) |
| 413 | |
| 414 | return dict( |
| 415 | pixel_values_flat_video=MultiModalFieldConfig.flat_from_sizes( |
| 416 | "video", video_num_patches |
| 417 | ), |
| 418 | video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True), |
| 419 | frames_indices=MultiModalFieldConfig.batched("video", keep_on_cpu=True), |
| 420 | frame_duration_ms=MultiModalFieldConfig.batched("video", keep_on_cpu=True), |
| 421 | ) |
| 422 | |
| 423 | def _get_audio_fields_config(self, hf_inputs: BatchFeature): |
| 424 | audio_num_clips = torch.as_tensor( |
| 425 | hf_inputs.get("audio_num_clips", torch.empty(0)) |
| 426 | ) |
| 427 | |
| 428 | return dict( |
| 429 | input_audio_features=MultiModalFieldConfig.flat_from_sizes( |
| 430 | "audio", audio_num_clips |
| 431 | ), |
| 432 | feature_attention_mask=MultiModalFieldConfig.flat_from_sizes( |
| 433 | "audio", audio_num_clips |
| 434 | ), |
| 435 | audio_num_clips=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), |
| 436 | ) |
| 437 | |
| 438 | def _get_mm_fields_config( |
| 439 | self, |
| 440 | hf_inputs: BatchFeature, |
| 441 | hf_processor_mm_kwargs: Mapping[str, object], |
| 442 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 443 | fields = self._get_image_fields_config(hf_inputs) |
| 444 | if self.info.supports_video: |
| 445 | fields |= self._get_video_fields_config(hf_inputs) |
| 446 | if self.info.supports_audio: |
| 447 | fields |= self._get_audio_fields_config(hf_inputs) |
| 448 | |
| 449 | return fields |
| 450 | |
| 451 | def _get_prompt_repl_image( |
| 452 | self, |
| 453 | mm_items: MultiModalDataItems, |
| 454 | hf_processor: NanoNemotronVLProcessor, |
| 455 | out_mm_data: BatchedTensorInputs, |
| 456 | ): |
| 457 | tokenizer = self.info.get_tokenizer() |
| 458 | vocab = tokenizer.get_vocab() |
| 459 | |
| 460 | if "image_num_patches" in out_mm_data: |
| 461 | image_num_patches = out_mm_data["image_num_patches"] |
| 462 | assert isinstance(image_num_patches, torch.Tensor) |
| 463 | image_num_patches = image_num_patches.tolist() |
| 464 | elif "image_embeds" in out_mm_data: |
| 465 | # to compute num_patches (similar to Qwen2-VL) |
| 466 | image_num_patches = [None] * len(out_mm_data["image_embeds"]) |
| 467 | else: |
| 468 | image_num_patches = [] |
| 469 | |
| 470 | def get_image_replacement(item_idx: int): |
| 471 | images = mm_items.get_items( |
| 472 | "image", (ImageEmbeddingItems, ImageProcessorItems) |
| 473 | ) |
| 474 | |
| 475 | if isinstance(images, ImageEmbeddingItems): |
| 476 | feature_size = images.get_feature_size(item_idx) |
| 477 | elif self.info.is_dynamic_tiler: |
| 478 | num_tokens_per_image = out_mm_data["num_tokens_per_image"] |
| 479 | assert isinstance(num_tokens_per_image, list) |
| 480 | feature_size = num_tokens_per_image[item_idx] |
| 481 | else: |
| 482 | assert isinstance(images, ImageProcessorItems) |
| 483 | image_size = images.get_image_size(item_idx) |
| 484 | max_num_tiles = hf_processor.max_num_tiles |
| 485 | feature_size = hf_processor.get_num_image_tokens( |
| 486 | image_width=image_size.width, |
| 487 | image_height=image_size.height, |
| 488 | max_num_tiles=max_num_tiles, |
| 489 | ) |
| 490 | |
| 491 | num_patches = None |
| 492 | local_image_num_patches = image_num_patches |
| 493 | if isinstance(local_image_num_patches, torch.Tensor): |
| 494 | local_image_num_patches = local_image_num_patches.tolist() |
| 495 | if isinstance(local_image_num_patches, (list, tuple)) and item_idx < len( |
| 496 | local_image_num_patches |
| 497 | ): |
| 498 | num_patches = int(local_image_num_patches[item_idx]) |
| 499 | |
| 500 | return hf_processor.get_image_repl(feature_size, num_patches) |
| 501 | |
| 502 | return PromptReplacement( |
| 503 | modality="image", |
| 504 | target=[vocab["<image>"]], |
| 505 | replacement=get_image_replacement, |
| 506 | ) |
| 507 | |
| 508 | def _get_prompt_repl_video( |
| 509 | self, |
| 510 | mm_items: MultiModalDataItems, |
| 511 | hf_processor: NanoNemotronVLProcessor, |
| 512 | out_mm_data: BatchedTensorInputs, |
| 513 | ): |
| 514 | tokenizer = self.info.get_tokenizer() |
| 515 | |
| 516 | if "video_num_patches" in out_mm_data: |
| 517 | video_num_patches = out_mm_data["video_num_patches"] |
| 518 | assert isinstance(video_num_patches, torch.Tensor) |
| 519 | video_num_patches = video_num_patches.tolist() |
| 520 | else: |
| 521 | video_num_patches = [] |
| 522 | |
| 523 | def get_video_replacement(item_idx: int): |
| 524 | video, metadata = mm_items["video"][item_idx] |
| 525 | patch_size = hf_processor.config.patch_size |
| 526 | downsample_ratio = hf_processor.config.downsample_ratio |
| 527 | target_patches = hf_processor.video_target_num_patches |
| 528 | |
| 529 | if target_patches is not None and video is not None and video.shape[0] > 0: |
| 530 | orig_h, orig_w = video.shape[1], video.shape[2] |
| 531 | _, _, feature_size = get_video_target_size_and_feature_size( |
| 532 | orig_w=orig_w, |
| 533 | orig_h=orig_h, |
| 534 | target_patches=target_patches, |
| 535 | maintain_aspect_ratio=hf_processor.video_maintain_aspect_ratio, |
| 536 | patch_size=patch_size, |
| 537 | downsample_ratio=downsample_ratio, |
| 538 | ) |
| 539 | else: |
| 540 | feature_size = hf_processor.num_image_token |
| 541 | num_patches = video_num_patches[item_idx] |
| 542 | if num_patches is not None: |
| 543 | assert isinstance(num_patches, int) |
| 544 | |
| 545 | T = hf_processor.video_temporal_patch_size |
| 546 | num_tubelets: int | None |
| 547 | if T > 1 and num_patches is not None: |
| 548 | num_tubelets = math.ceil(num_patches / T) |
| 549 | else: |
| 550 | num_tubelets = num_patches |
| 551 | assert num_tubelets is not None |
| 552 | |
| 553 | video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate |
| 554 | if video_pruning_rate is not None and video_pruning_rate > 0.0: |
| 555 | # Start of EVS-specific code |
| 556 | num_tokens = compute_retained_tokens_count( |
| 557 | tokens_per_frame=feature_size, |
| 558 | num_frames=num_tubelets, |
| 559 | q=video_pruning_rate, |
| 560 | ) |
| 561 | # Here we just need placeholders that won't actually be replaced - |
| 562 | # we just need to make sure the total number of tokens is correct |
| 563 | # assign all tokens to the first frame |
| 564 | tokens_per_frame = [num_tokens] + [0] * (num_tubelets - 1) |
| 565 | |
| 566 | # End of EVS-specific code |
| 567 | else: |
| 568 | tokens_per_frame = [feature_size] * num_tubelets |
| 569 | |
| 570 | frame_duration_ms = int(1000 / metadata["fps"]) |
| 571 | return hf_processor.get_video_repl( |
| 572 | tokens_per_frame=tokens_per_frame, |
| 573 | frames_indices=metadata["frames_indices"], |
| 574 | frame_duration_ms=frame_duration_ms, |
| 575 | tokenizer=hf_processor.tokenizer, |
| 576 | img_start_token_ids=hf_processor._img_start_token_ids, |
| 577 | img_end_token_ids=hf_processor._img_end_token_ids, |
| 578 | img_context_token_ids=hf_processor._img_context_token_ids, |
| 579 | video_temporal_patch_size=T, |
| 580 | ) |
| 581 | |
| 582 | return PromptReplacement( |
| 583 | modality="video", |
| 584 | target=cached_encode(tokenizer, "<video>", add_special_tokens=False), |
| 585 | replacement=get_video_replacement, |
| 586 | ) |
| 587 | |
| 588 | def _get_prompt_repl_audio( |
| 589 | self, |
| 590 | mm_items: MultiModalDataItems, |
| 591 | hf_processor: NanoNemotronVLProcessor, |
| 592 | out_mm_data: BatchedTensorInputs, |
| 593 | ): |
| 594 | tokenizer = self.info.get_tokenizer() |
| 595 | |
| 596 | def get_audio_replacement(item_idx: int): |
| 597 | audios = mm_items.get_items("audio", AudioProcessorItems) |
| 598 | return hf_processor.get_audio_repl(audios.get(item_idx)) |
| 599 | |
| 600 | return PromptReplacement( |
| 601 | modality="audio", |
| 602 | target=cached_encode(tokenizer, AUDIO_CONTEXT, add_special_tokens=False), |
| 603 | replacement=get_audio_replacement, |
| 604 | ) |
| 605 | |
| 606 | def _get_prompt_updates( |
| 607 | self, |
| 608 | mm_items: MultiModalDataItems, |
| 609 | hf_processor_mm_kwargs: Mapping[str, object], |
| 610 | out_mm_kwargs: MultiModalKwargsItems, |
| 611 | ) -> Sequence[PromptUpdate]: |
| 612 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 613 | out_mm_data = out_mm_kwargs.get_data() |
| 614 | |
| 615 | prompt_repls = [ |
| 616 | self._get_prompt_repl_image(mm_items, hf_processor, out_mm_data), |
| 617 | ] |
| 618 | if self.info.supports_video: |
| 619 | prompt_repls.append( |
| 620 | self._get_prompt_repl_video(mm_items, hf_processor, out_mm_data) |
| 621 | ) |
| 622 | if self.info.supports_audio: |
| 623 | prompt_repls.append( |
| 624 | self._get_prompt_repl_audio(mm_items, hf_processor, out_mm_data) |
| 625 | ) |
| 626 | |
| 627 | return prompt_repls |
| 628 | |
| 629 | def _extract_audio_from_videos( |
| 630 | self, |
| 631 | mm_items: MultiModalDataItems, |
| 632 | ) -> tuple[MultiModalDataItems, list[AudioItem], list[bool]]: |
| 633 | """Extract audio tracks from video bytes in *mm_items*. |
| 634 | |
| 635 | Videos whose bytes are missing or that contain no audio stream are |
| 636 | silently skipped. The returned *has_audio* mask is aligned with |
| 637 | the video list so callers know which ``<video>`` tokens need an |
| 638 | accompanying audio context. |
| 639 | |
| 640 | Returns: |
| 641 | A 3-tuple of (augmented mm_items, extracted audio items, |
| 642 | per-video boolean mask indicating which videos have audio). |
| 643 | |
| 644 | """ |
| 645 | videos = mm_items.get_items("video", VideoProcessorItems) |
| 646 | assert isinstance(videos.metadata, list) |
| 647 | |
| 648 | metadata_list = videos.metadata |
| 649 | |
| 650 | audio_items: list[AudioItem] = [] |
| 651 | has_audio: list[bool] = [] |
| 652 | for idx, metadata in enumerate(metadata_list): |
| 653 | assert metadata is not None |
| 654 | video_bytes = metadata.get("original_video_bytes") |
| 655 | if video_bytes is None or len(video_bytes) == 0: |
| 656 | raise ValueError( |
| 657 | "Cannot extract audio from video: original_video_bytes is " |
| 658 | "missing or empty. When using use_audio_in_video=True, " |
| 659 | "video must be loaded with keep_video_bytes=True (e.g. via " |
| 660 | "the chat API with a model that sets use_audio_in_video)." |
| 661 | ) |
| 662 | try: |
| 663 | audio_items.append( |
| 664 | load_audio_pyav( |
| 665 | BytesIO(video_bytes), |
| 666 | max_duration_s=envs.VLLM_MAX_AUDIO_DECODE_DURATION_S, |
| 667 | ) |
| 668 | ) |
| 669 | has_audio.append(True) |
| 670 | except Exception: |
| 671 | logger.debug( |
| 672 | "Video %d: no audio stream found, skipping audio extraction.", |
| 673 | idx, |
| 674 | exc_info=True, |
| 675 | ) |
| 676 | has_audio.append(False) |
| 677 | |
| 678 | # Create a new VideoProcessorItems with metadata that does not contain |
| 679 | # the large video bytes, to avoid modifying the input `mm_items`. |
| 680 | new_metadata_list = metadata_list.copy() |
| 681 | for idx, metadata in enumerate(metadata_list): |
| 682 | assert metadata is not None |
| 683 | new_metadata_list[idx] = { |
| 684 | key: value |
| 685 | for key, value in metadata.items() |
| 686 | if key != "original_video_bytes" |
| 687 | } |
| 688 | new_videos = VideoProcessorItems(data=videos.data, metadata=new_metadata_list) |
| 689 | |
| 690 | audio_parsed = MultiModalDataItems() |
| 691 | if audio_items: |
| 692 | audio_parsed = self.data_parser.parse_mm_data({"audio": audio_items}) |
| 693 | |
| 694 | # Create a new MultiModalDataItems with the new video and audio items. |
| 695 | new_mm_items_dict = {**mm_items, **audio_parsed, "video": new_videos} |
| 696 | new_mm_items = MultiModalDataItems(new_mm_items_dict) |
| 697 | |
| 698 | return new_mm_items, audio_items, has_audio |
| 699 | |
| 700 | def apply( |
| 701 | self, |
| 702 | inputs: ProcessorInputs, |
| 703 | timing_ctx: TimingContext, |
| 704 | ) -> MultiModalInput: |
| 705 | mm_config = self.info.ctx.model_config.get_multimodal_config() |
| 706 | merged_kwargs = mm_config.merge_mm_processor_kwargs( |
| 707 | inputs.hf_processor_mm_kwargs |
| 708 | ) |
| 709 | use_audio_in_video = bool(merged_kwargs.get("use_audio_in_video", False)) |
| 710 | |
| 711 | inputs.hf_processor_mm_kwargs = { |
| 712 | k: v |
| 713 | for k, v in inputs.hf_processor_mm_kwargs.items() |
| 714 | if k != "use_audio_in_video" |
| 715 | } |
| 716 | |
| 717 | if not (use_audio_in_video and "video" in inputs.mm_data_items): |
| 718 | return super().apply(inputs, timing_ctx) |
| 719 | |
| 720 | mm_items = inputs.mm_data_items |
| 721 | if "audio" in mm_items: |
| 722 | # Audio was pre-populated by upstream (e.g., OpenAI chat endpoint). |
| 723 | # Reuse existing audio items; validate 1:1 correspondence. |
| 724 | videos = mm_items.get_items("video", VideoProcessorItems) |
| 725 | audios = mm_items.get_items("audio", AudioProcessorItems) |
| 726 | if len(audios) != len(videos): |
| 727 | raise ValueError( |
| 728 | "use_audio_in_video requires equal number of audio and " |
| 729 | f"video items, got num_audios={len(audios)}, " |
| 730 | f"num_videos={len(videos)}" |
| 731 | ) |
| 732 | audio_items = audios.get_all() |
| 733 | has_audio = [True] * len(videos) |
| 734 | logger.info( |
| 735 | "Using %d pre-populated audio item(s) from upstream.", |
| 736 | len(audio_items), |
| 737 | ) |
| 738 | else: |
| 739 | # Extract audio from video bytes (library usage path). |
| 740 | mm_items, audio_items, has_audio = self._extract_audio_from_videos(mm_items) |
| 741 | inputs.mm_data_items = mm_items |
| 742 | logger.info( |
| 743 | "Extracted audio from video bytes: %d audio(s), has_audio=%s.", |
| 744 | len(audio_items), |
| 745 | has_audio, |
| 746 | ) |
| 747 | |
| 748 | if not audio_items: |
| 749 | return super().apply(inputs, timing_ctx) |
| 750 | |
| 751 | tokenizer = self.info.get_tokenizer() |
| 752 | prompt = tokenizer.decode(inputs.prompt, skip_special_tokens=False) |
| 753 | |
| 754 | # Inject AUDIO_CONTEXT only after <video> tokens whose video |
| 755 | # actually contained an audio stream (preserving video-audio pairing). |
| 756 | tag = "<video>" |
| 757 | head, *rest = prompt.split(tag) |
| 758 | rebuilt = [head] |
| 759 | for append_audio, part in zip(has_audio, rest, strict=True): |
| 760 | rebuilt.append(tag) |
| 761 | if append_audio: |
| 762 | rebuilt.append(AUDIO_CONTEXT) |
| 763 | rebuilt.append(part) |
| 764 | prompt = "".join(rebuilt) |
| 765 | |
| 766 | inputs.prompt = tokenizer.encode(prompt, add_special_tokens=False) |
| 767 | |
| 768 | # Bypass the cached path: the HF processor must receive the |
| 769 | # prompt (with injected <so_embedding>) and the audio data |
| 770 | # together so it can perform audio-token replacement natively. |
| 771 | mm_res = self._apply_hf_processor(inputs, timing_ctx) |
| 772 | |
| 773 | with timing_ctx.record("apply_prompt_updates"): |
| 774 | prompt_ids, mm_placeholders = self._maybe_apply_prompt_updates( |
| 775 | mm_items=mm_items, |
| 776 | mm_res=mm_res, |
| 777 | ) |
| 778 | |
| 779 | mm_placeholder_ranges = { |
| 780 | modality: [item.to_range() for item in placeholders] |
| 781 | for modality, placeholders in mm_placeholders.items() |
| 782 | } |
| 783 | |
| 784 | return MultiModalInput( |
| 785 | type="multimodal", |
| 786 | prompt_token_ids=prompt_ids, |
| 787 | mm_kwargs=mm_res.kwargs, |
| 788 | mm_hashes=mm_res.hashes, |
| 789 | mm_placeholders=mm_placeholder_ranges, |
| 790 | ) |
| 791 | |
| 792 | |
| 793 | class NanoNemotronVLDummyInputsBuilder( |
| 794 | BaseDummyInputsBuilder[NanoNemotronVLProcessingInfo] |
| 795 | ): |
| 796 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 797 | num_images = mm_counts.get("image", 0) |
| 798 | num_videos = mm_counts.get("video", 0) |
| 799 | num_audios = mm_counts.get("audio", 0) |
| 800 | |
| 801 | return ( |
| 802 | "<image>" * num_images + "<video>" * num_videos + AUDIO_CONTEXT * num_audios |
| 803 | ) |
| 804 | |
| 805 | def _get_dummy_videos( |
| 806 | self, |
| 807 | *, |
| 808 | width: int, |
| 809 | height: int, |
| 810 | num_frames: int, |
| 811 | num_videos: int, |
| 812 | overrides: VideoDummyOptions | None = None, |
| 813 | ) -> list[VideoItem]: |
| 814 | videos = super()._get_dummy_videos( |
| 815 | width=width, |
| 816 | height=height, |
| 817 | num_frames=num_frames, |
| 818 | num_videos=num_videos, |
| 819 | overrides=overrides, |
| 820 | ) |
| 821 | videos = [v.copy() for v in videos] |
| 822 | |
| 823 | video_items: list[VideoItem] = [] |
| 824 | for video in videos: |
| 825 | video_num_frames = video.shape[0] |
| 826 | video_metadata = { |
| 827 | "fps": 2, |
| 828 | "duration": video_num_frames / 2.0, |
| 829 | "total_num_frames": video_num_frames, |
| 830 | "frames_indices": list(range(video_num_frames)), |
| 831 | "video_backend": "opencv_dynamic", |
| 832 | "do_sample_frames": False, |
| 833 | } |
| 834 | video_items.append((video, video_metadata)) |
| 835 | |
| 836 | return video_items |
| 837 | |
| 838 | def get_dummy_mm_data( |
| 839 | self, |
| 840 | seq_len: int, |
| 841 | mm_counts: Mapping[str, int], |
| 842 | mm_options: MultiModalDummyOptions, |
| 843 | ) -> MultiModalDataDict: |
| 844 | (target_width, target_height), _ = ( |
| 845 | self.info.get_dummy_image_size_and_max_tokens(mm_counts) |
| 846 | ) |
| 847 | processor = self.info.get_hf_processor() |
| 848 | |
| 849 | dummy_image = { |
| 850 | "image": self._get_dummy_images( |
| 851 | width=target_width, |
| 852 | height=target_height, |
| 853 | num_images=mm_counts.get("image", 0), |
| 854 | overrides=mm_options.get("image"), |
| 855 | ) |
| 856 | } |
| 857 | |
| 858 | if self.info.supports_video: |
| 859 | config = self.info.get_hf_config() |
| 860 | image_size: int = config.force_image_size |
| 861 | |
| 862 | # When video_target_num_patches is set the per-frame pixel |
| 863 | # resolution can exceed image_size. Use the actual target |
| 864 | # dimensions so that profiling sees the correct upper bound. |
| 865 | if processor.video_target_num_patches is not None: |
| 866 | target_w, target_h, _ = get_video_target_size_and_feature_size( |
| 867 | orig_w=image_size, |
| 868 | orig_h=image_size, |
| 869 | target_patches=processor.video_target_num_patches, |
| 870 | maintain_aspect_ratio=processor.video_maintain_aspect_ratio, |
| 871 | patch_size=config.patch_size, |
| 872 | downsample_ratio=config.downsample_ratio, |
| 873 | ) |
| 874 | video_width, video_height = target_w, target_h |
| 875 | else: |
| 876 | video_width, video_height = image_size, image_size |
| 877 | |
| 878 | target_num_frames = self.info.get_num_frames_with_most_features( |
| 879 | seq_len, mm_counts |
| 880 | ) |
| 881 | mm_config = self.info.ctx.get_mm_config() |
| 882 | if num_frames := mm_config.media_io_kwargs.get("video", {}).get( |
| 883 | "num_frames" |
| 884 | ): |
| 885 | assert num_frames > 0 |
| 886 | target_num_frames = num_frames |
| 887 | num_videos = mm_counts.get("video", 0) |
| 888 | video_overrides = mm_options.get("video") |
| 889 | dummy_video = { |
| 890 | "video": self._get_dummy_videos( |
| 891 | width=video_width, |
| 892 | height=video_height, |
| 893 | num_frames=target_num_frames, |
| 894 | num_videos=num_videos, |
| 895 | overrides=video_overrides, |
| 896 | ) |
| 897 | } |
| 898 | else: |
| 899 | dummy_video = {} |
| 900 | |
| 901 | if sound_config := self.info.sound_config: |
| 902 | num_audios = mm_counts.get("audio", 0) |
| 903 | audio_overrides = mm_options.get("audio") if mm_options else None |
| 904 | tokens_per_audio = max(1, seq_len // max(num_audios, 1)) |
| 905 | max_audio_num_samples = MAX_AUDIO_LEN_S * sound_config.sampling_rate |
| 906 | calculated_max_audio_num_samples = ParakeetExtractor.audio_length( |
| 907 | sound_config, tokens_per_audio |
| 908 | ) |
| 909 | audio_len = min(max_audio_num_samples, calculated_max_audio_num_samples) |
| 910 | dummy_audio = { |
| 911 | "audio": self._get_dummy_audios( |
| 912 | length=audio_len, |
| 913 | num_audios=num_audios, |
| 914 | overrides=audio_overrides, |
| 915 | ) |
| 916 | } |
| 917 | else: |
| 918 | dummy_audio = {} |
| 919 | |
| 920 | return {**dummy_image, **dummy_video, **dummy_audio} |
| 921 | |
| 922 | |
| 923 | @MULTIMODAL_REGISTRY.register_processor( |
| 924 | NanoNemotronVLMultiModalProcessor, |
| 925 | info=NanoNemotronVLProcessingInfo, |
| 926 | dummy_inputs=NanoNemotronVLDummyInputsBuilder, |
| 927 | ) |
| 928 | class NemotronH_Nano_VL_V2( |
| 929 | nn.Module, |
| 930 | HasInnerState, |
| 931 | IsHybrid, |
| 932 | SupportsEagle3, |
| 933 | SupportsMultiModal, |
| 934 | SupportsMultiModalPruning, |
| 935 | SupportsLoRA, |
| 936 | ): |
| 937 | requires_sequential_video_encoding = True |
| 938 | """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet""" |
| 939 | |
| 940 | # LoRA covers the language model only |
| 941 | is_non_gated_moe = NemotronHForCausalLM.is_non_gated_moe |
| 942 | packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping |
| 943 | embedding_modules = NemotronHForCausalLM.embedding_modules |
| 944 | lora_skip_prefixes = NemotronHForCausalLM.lora_skip_prefixes |
| 945 | |
| 946 | hf_to_vllm_mapper = WeightsMapper( |
| 947 | orig_to_new_prefix={ |
| 948 | "language_model.backbone": "language_model.model", |
| 949 | }, |
| 950 | ) |
| 951 | |
| 952 | @classmethod |
| 953 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 954 | if modality.startswith("image"): |
| 955 | return "<image>" |
| 956 | if modality.startswith("video"): |
| 957 | return "<video>" |
| 958 | if modality.startswith("audio"): |
| 959 | return AUDIO_CONTEXT |
| 960 | return None |
| 961 | |
| 962 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 963 | super().__init__() |
| 964 | model_config = vllm_config.model_config |
| 965 | config = model_config.hf_config |
| 966 | multimodal_config = model_config.multimodal_config |
| 967 | assert multimodal_config is not None |
| 968 | image_size = config.force_image_size |
| 969 | patch_size = config.patch_size |
| 970 | self.patch_size = patch_size |
| 971 | self.template = config.template |
| 972 | self.num_image_token = int( |
| 973 | (image_size // patch_size) ** 2 * (config.downsample_ratio**2) |
| 974 | ) |
| 975 | self.downsample_ratio = config.downsample_ratio |
| 976 | self.ps_version = config.ps_version |
| 977 | self.image_tag_type = config.image_tag_type |
| 978 | self.video_pruning_rate = multimodal_config.video_pruning_rate |
| 979 | |
| 980 | vision_config = getattr(config, "vision_config", config) |
| 981 | self.video_temporal_patch_size: int = getattr( |
| 982 | vision_config, "video_temporal_patch_size", 1 |
| 983 | ) |
| 984 | |
| 985 | with self._mark_language_model(vllm_config): |
| 986 | self.language_model = init_vllm_registered_model( |
| 987 | vllm_config=vllm_config, |
| 988 | hf_config=config.text_config, |
| 989 | prefix=maybe_prefix(prefix, "language_model"), |
| 990 | ) |
| 991 | llm_dtype = self.language_model.config.dtype |
| 992 | assert isinstance(llm_dtype, torch.dtype) |
| 993 | self.llm_dtype = llm_dtype |
| 994 | with self._mark_tower_model(vllm_config, {"image", "video", "audio"}): |
| 995 | self.vision_model = self.get_vit_model_from_radio_config(config).to( |
| 996 | llm_dtype |
| 997 | ) |
| 998 | |
| 999 | # Construct the vision projection. |
| 1000 | vit_hidden_size = config.vit_hidden_size |
| 1001 | vision_projection_hidden_size = config.projector_hidden_size |
| 1002 | llm_hidden_size = config.text_config.hidden_size |
| 1003 | |
| 1004 | mlp1 = nn.Sequential( |
| 1005 | RMSNorm( |
| 1006 | hidden_size=vit_hidden_size |
| 1007 | * int(round(1 / self.downsample_ratio)) ** 2, |
| 1008 | eps=1e-5, |
| 1009 | ), |
| 1010 | nn.Linear( |
| 1011 | vit_hidden_size * int(round(1 / self.downsample_ratio)) ** 2, |
| 1012 | vision_projection_hidden_size, |
| 1013 | bias=False, |
| 1014 | ), |
| 1015 | ReLUSquaredActivation(), |
| 1016 | nn.Linear(vision_projection_hidden_size, llm_hidden_size, bias=False), |
| 1017 | ) |
| 1018 | self.mlp1 = mlp1.to(llm_dtype) |
| 1019 | self.sound_encoder: ProjectedParakeet | None = None |
| 1020 | if getattr(config, "sound_config", None) is not None: |
| 1021 | logger.info_once( |
| 1022 | "Found sound config, initializing sound encoder for Nemotron AVLM", |
| 1023 | scope="global", |
| 1024 | ) |
| 1025 | self.sound_encoder = ProjectedParakeet( |
| 1026 | config.sound_config, |
| 1027 | dtype=llm_dtype, |
| 1028 | llm_hidden_size=llm_hidden_size, |
| 1029 | max_model_len=model_config.max_model_len, |
| 1030 | ) |
| 1031 | |
| 1032 | self.config = config |
| 1033 | self.model_config = vllm_config.model_config |
| 1034 | |
| 1035 | # Pre-tokenize special tokens for video processing |
| 1036 | # to avoid repeated tokenization |
| 1037 | tokenizer = cached_tokenizer_from_config(model_config) |
| 1038 | self._img_start_token_ids = tokenizer.encode( |
| 1039 | IMG_START, add_special_tokens=False |
| 1040 | ) |
| 1041 | self._img_end_token_ids = tokenizer.encode(IMG_END, add_special_tokens=False) |
| 1042 | self._img_context_token_ids = tokenizer.encode( |
| 1043 | IMG_CONTEXT, add_special_tokens=False |
| 1044 | ) |
| 1045 | self.dynamic_resolution = BaseNanoNemotronVLProcessor.use_dynamic_resolution( |
| 1046 | config |
| 1047 | ) |
| 1048 | if self.dynamic_resolution: |
| 1049 | logger.info_once( |
| 1050 | "Dynamic resolution is enabled for NanoNemotronVLProcessor", |
| 1051 | scope="global", |
| 1052 | ) |
| 1053 | |
| 1054 | def pixel_shuffle(self, x, scale_factor=0.5): |
| 1055 | n, h, w, c = x.size() |
| 1056 | r = int(1 / scale_factor) |
| 1057 | new_h = h // r |
| 1058 | new_w = w // r |
| 1059 | new_c = c * r * r |
| 1060 | |
| 1061 | x = x.view(n, new_h, r, new_w, r, c) |
| 1062 | if self.ps_version == "v1": |
| 1063 | warnings.warn( |
| 1064 | "In ps_version 'v1', the height and width have not " |
| 1065 | "been swapped back, which results in a transposed image.", |
| 1066 | stacklevel=2, |
| 1067 | ) |
| 1068 | x = x.permute(0, 3, 1, 2, 4, 5).reshape(n, new_w, new_h, new_c) |
| 1069 | else: |
| 1070 | x = x.permute(0, 1, 3, 2, 4, 5).reshape(n, new_h, new_w, new_c) |
| 1071 | return x |
| 1072 | |
| 1073 | def pixel_shuffle_dynamic_res( |
| 1074 | self, x: torch.Tensor, *, imgs_sizes: list[tuple[int, int]] |
| 1075 | ) -> torch.Tensor: |
| 1076 | patch_dim = self.patch_size |
| 1077 | seq_lens = calc_seq_lens(imgs_sizes, patch_dim) |
| 1078 | splits = torch.split(x, seq_lens, dim=-2) |
| 1079 | out = [] |
| 1080 | for i, sv in enumerate(splits): |
| 1081 | h = imgs_sizes[i][0] // patch_dim |
| 1082 | w = imgs_sizes[i][1] // patch_dim |
| 1083 | sv = sv.reshape(sv.shape[0], h, w, -1) |
| 1084 | sv = self.pixel_shuffle(sv, scale_factor=self.downsample_ratio) |
| 1085 | sv = sv.flatten(1, 2) |
| 1086 | out.append(sv) |
| 1087 | |
| 1088 | x = torch.cat(out, dim=-2) |
| 1089 | |
| 1090 | return x |
| 1091 | |
| 1092 | def extract_feature_dynamic( |
| 1093 | self, pixel_values: torch.Tensor, imgs_sizes: list[tuple[int, int]] |
| 1094 | ): |
| 1095 | """Dynamic resolution extract_feature for images.""" |
| 1096 | _, vit_embeds = self.vision_model(pixel_values, imgs_sizes=imgs_sizes) |
| 1097 | vit_embeds = vit_embeds.to(dtype=torch.bfloat16) |
| 1098 | vit_embeds = self.pixel_shuffle_dynamic_res(vit_embeds, imgs_sizes=imgs_sizes) |
| 1099 | vit_embeds = self.mlp1(vit_embeds) |
| 1100 | return vit_embeds |
| 1101 | |
| 1102 | def extract_feature( |
| 1103 | self, |
| 1104 | pixel_values: torch.Tensor, |
| 1105 | num_frames: int | None = None, |
| 1106 | ) -> torch.Tensor: |
| 1107 | # Process images in a micro-batch of at most 128 frames per call |
| 1108 | # This is done on purpose to ensure peak GPU ram usage of huge batch |
| 1109 | # (namely for really long videos with EVS ON) won't cause any problems |
| 1110 | # as we don't support chunked prefill for video media |
| 1111 | # When num_frames is provided and temporal_patch_size > 1, consecutive |
| 1112 | # frames are grouped into tubelets — the batch size must be a multiple |
| 1113 | # of T so chunk boundaries don't split a tubelet. |
| 1114 | N, _C, H, W = pixel_values.shape |
| 1115 | |
| 1116 | T = self.video_temporal_patch_size if num_frames is not None else 1 |
| 1117 | micro_batch_size = 128 - (128 % T) |
| 1118 | patch_size = self.patch_size |
| 1119 | H_patches = H // patch_size |
| 1120 | W_patches = W // patch_size |
| 1121 | |
| 1122 | vit_embeds_list = [] |
| 1123 | for i in range(0, N, micro_batch_size): |
| 1124 | chunk = pixel_values[i : i + micro_batch_size] |
| 1125 | if num_frames is not None and T > 1: |
| 1126 | _, vit_embeds = self.vision_model(chunk, num_frames=chunk.shape[0]) |
| 1127 | else: |
| 1128 | _, vit_embeds = self.vision_model(chunk) |
| 1129 | vit_embeds = vit_embeds.to(dtype=torch.bfloat16) |
| 1130 | vit_embeds = vit_embeds.reshape( |
| 1131 | vit_embeds.shape[0], H_patches, W_patches, -1 |
| 1132 | ) |
| 1133 | vit_embeds = self.pixel_shuffle( |
| 1134 | vit_embeds, scale_factor=self.downsample_ratio |
| 1135 | ) |
| 1136 | vit_embeds = vit_embeds.reshape( |
| 1137 | vit_embeds.shape[0], -1, vit_embeds.shape[-1] |
| 1138 | ) |
| 1139 | vit_embeds = self.mlp1(vit_embeds) |
| 1140 | vit_embeds_list.append(vit_embeds) |
| 1141 | |
| 1142 | vit_embeds = torch.cat(vit_embeds_list, dim=0) |
| 1143 | return vit_embeds |
| 1144 | |
| 1145 | def _parse_and_validate_image_input( |
| 1146 | self, **kwargs: object |
| 1147 | ) -> NanoNemotronVLImageInputs | None: |
| 1148 | image_embeds = kwargs.pop("image_embeds", None) |
| 1149 | if image_embeds is not None: |
| 1150 | return NanoNemotronVLImageEmbeddingInputs( |
| 1151 | type="image_embeds", |
| 1152 | data=image_embeds, |
| 1153 | ) |
| 1154 | |
| 1155 | pixel_values_flat = kwargs.pop("pixel_values_flat", None) |
| 1156 | if pixel_values_flat is None: |
| 1157 | return None |
| 1158 | |
| 1159 | if self.dynamic_resolution: |
| 1160 | assert isinstance(pixel_values_flat, (list, torch.Tensor)) |
| 1161 | pixel_values_flat = DynamicResolutionImageTiler.stack( |
| 1162 | pixel_values_flat, self.patch_size |
| 1163 | ) |
| 1164 | imgs_sizes = kwargs.pop("imgs_sizes") |
| 1165 | num_tokens_per_image = kwargs.pop("num_tokens_per_image") |
| 1166 | assert isinstance(imgs_sizes, list) |
| 1167 | assert isinstance(num_tokens_per_image, list) |
| 1168 | return NanoNemotronVLImagePixelInputsDynamic( |
| 1169 | pixel_values_flat=pixel_values_flat, |
| 1170 | imgs_sizes=imgs_sizes, |
| 1171 | num_tokens_per_image=num_tokens_per_image, |
| 1172 | ) |
| 1173 | else: |
| 1174 | assert isinstance(pixel_values_flat, torch.Tensor) |
| 1175 | image_num_patches = kwargs.pop("image_num_patches") |
| 1176 | assert isinstance(image_num_patches, torch.Tensor) |
| 1177 | return NanoNemotronVLImagePixelInputs( |
| 1178 | pixel_values_flat=pixel_values_flat, |
| 1179 | num_patches=image_num_patches, |
| 1180 | ) |
| 1181 | |
| 1182 | def _process_image_input_dynamic( |
| 1183 | self, image_input: NanoNemotronVLImagePixelInputsDynamic |
| 1184 | ) -> tuple[torch.Tensor, ...]: |
| 1185 | image_embeds = self.extract_feature_dynamic( |
| 1186 | image_input.pixel_values_flat, image_input.imgs_sizes |
| 1187 | ) |
| 1188 | num_tokens_per_image = image_input.num_tokens_per_image |
| 1189 | |
| 1190 | if len(num_tokens_per_image) == 1: |
| 1191 | return (image_embeds.view(-1, self.config.text_config.hidden_size),) |
| 1192 | |
| 1193 | image_embeds = image_embeds.view(-1, self.config.text_config.hidden_size) |
| 1194 | return image_embeds.split(num_tokens_per_image) |
| 1195 | |
| 1196 | def _process_image_input( |
| 1197 | self, |
| 1198 | image_input: NanoNemotronVLImagePixelInputs | NanoNemotronVLVideoPixelInputs, |
| 1199 | ) -> tuple[torch.Tensor, ...]: |
| 1200 | image_embeds = self.extract_feature(image_input["pixel_values_flat"]) |
| 1201 | num_patches = image_input["num_patches"] |
| 1202 | |
| 1203 | # Only one image in the current batch |
| 1204 | if len(num_patches) == 1: |
| 1205 | return (image_embeds.view(-1, self.config.text_config.hidden_size),) |
| 1206 | |
| 1207 | # NOTE: Image embeddings are split into separate tensors for each image |
| 1208 | # by the size of each embedding. |
| 1209 | feature_size = image_embeds.shape[1] |
| 1210 | image_embeds = image_embeds.view(-1, self.config.text_config.hidden_size) |
| 1211 | image_feature_sizes = [ |
| 1212 | num_patches * feature_size for num_patches in num_patches |
| 1213 | ] |
| 1214 | return image_embeds.split(image_feature_sizes) |
| 1215 | |
| 1216 | def _process_video_input( |
| 1217 | self, video_input: NanoNemotronVLVideoPixelInputs |
| 1218 | ) -> tuple[torch.Tensor, ...]: |
| 1219 | """Process video input and create final embeddings with video content |
| 1220 | and indicator tokens.""" |
| 1221 | T = self.video_temporal_patch_size |
| 1222 | |
| 1223 | if T > 1: |
| 1224 | video_embeddings = self._extract_video_embeddings_temporal(video_input) |
| 1225 | else: |
| 1226 | video_embeddings = self._process_image_input(video_input) |
| 1227 | |
| 1228 | final_video_embeddings: tuple[torch.Tensor, ...] = () |
| 1229 | |
| 1230 | downsample_ratio = self.config.downsample_ratio |
| 1231 | patch_size = self.config.patch_size |
| 1232 | pixel_values = video_input["pixel_values_flat"] |
| 1233 | frame_h, frame_w = pixel_values.shape[-2], pixel_values.shape[-1] |
| 1234 | rows = int(frame_h * downsample_ratio // patch_size) |
| 1235 | cols = int(frame_w * downsample_ratio // patch_size) |
| 1236 | video_pruning_rate = self.video_pruning_rate |
| 1237 | video_num_frames = video_input["num_patches"].tolist() |
| 1238 | video_frames_indices = video_input["frames_indices"].split(video_num_frames) |
| 1239 | # Calculate video feature dimensions (number of frames and |
| 1240 | # their feature size (AKA tokens per frame)) |
| 1241 | # TODO: Maybe this can be optimized to avoid the loop? |
| 1242 | for i, single_video_embeddings in enumerate(video_embeddings): |
| 1243 | num_frames = video_num_frames[i] |
| 1244 | frames_indices = video_frames_indices[i].tolist() |
| 1245 | frame_duration_ms = video_input["frame_duration_ms"][i].item() |
| 1246 | num_tubelets = math.ceil(num_frames / T) if T > 1 else num_frames |
| 1247 | assert single_video_embeddings.shape[0] % num_tubelets == 0 |
| 1248 | |
| 1249 | if video_pruning_rate is not None and video_pruning_rate > 0.0: |
| 1250 | # Start of EVS-specific code |
| 1251 | retention_mask = compute_retention_mask( |
| 1252 | single_video_embeddings, |
| 1253 | video_size_thw=(num_tubelets, rows, cols), |
| 1254 | spatial_merge_size=1, |
| 1255 | q=video_pruning_rate, |
| 1256 | ) |
| 1257 | |
| 1258 | # apply retention mask |
| 1259 | single_video_embeddings = single_video_embeddings[retention_mask] |
| 1260 | |
| 1261 | # calculate the actual number of retained tokens per frame |
| 1262 | retention_mask_thw = retention_mask.reshape(num_tubelets, rows, cols) |
| 1263 | num_tokens_per_frame = ( |
| 1264 | retention_mask_thw.sum(dim=(1, 2)).long().tolist() |
| 1265 | ) |
| 1266 | # End of EVS-specific code |
| 1267 | else: |
| 1268 | feature_size = single_video_embeddings.shape[0] // num_tubelets |
| 1269 | num_tokens_per_frame = [feature_size] * num_tubelets |
| 1270 | |
| 1271 | final_video_embeddings += ( |
| 1272 | self._create_final_video_embeddings( |
| 1273 | single_video_embeddings, |
| 1274 | num_tokens_per_frame, |
| 1275 | frames_indices, |
| 1276 | frame_duration_ms, |
| 1277 | video_temporal_patch_size=T, |
| 1278 | ), |
| 1279 | ) |
| 1280 | |
| 1281 | return final_video_embeddings |
| 1282 | |
| 1283 | def _extract_video_embeddings_temporal( |
| 1284 | self, video_input: NanoNemotronVLVideoPixelInputs |
| 1285 | ) -> tuple[torch.Tensor, ...]: |
| 1286 | """Extract per-video embeddings with temporal compression. |
| 1287 | |
| 1288 | Each video is processed separately through extract_feature with |
| 1289 | num_frames, which uses the fixed-resolution temporal path in RADIO |
| 1290 | (no attention mask, flash attention). |
| 1291 | """ |
| 1292 | pixel_values = video_input["pixel_values_flat"] |
| 1293 | num_frames_per_video = video_input["num_patches"].tolist() |
| 1294 | hidden_size = self.config.text_config.hidden_size |
| 1295 | |
| 1296 | results: list[torch.Tensor] = [] |
| 1297 | frame_offset = 0 |
| 1298 | for nf in num_frames_per_video: |
| 1299 | video_frames = pixel_values[frame_offset : frame_offset + nf] |
| 1300 | frame_offset += nf |
| 1301 | |
| 1302 | vit_embeds = self.extract_feature(video_frames, num_frames=nf) |
| 1303 | results.append(vit_embeds.view(-1, hidden_size)) |
| 1304 | |
| 1305 | return tuple(results) |
| 1306 | |
| 1307 | def _process_audio_input( |
| 1308 | self, audio_input: NanoNemotronVLAudioFeatureInputs |
| 1309 | ) -> tuple[torch.Tensor, ...]: |
| 1310 | assert self.sound_encoder is not None |
| 1311 | input_audio_features = audio_input.input_audio_features |
| 1312 | feature_attention_mask = audio_input.feature_attention_mask |
| 1313 | audio_num_clips = audio_input.audio_num_clips |
| 1314 | target_device = next(self.sound_encoder.parameters()).device |
| 1315 | |
| 1316 | input_audio_features = input_audio_features.to( |
| 1317 | dtype=self.llm_dtype, device=target_device |
| 1318 | ) |
| 1319 | feature_attention_mask = feature_attention_mask.to(device=target_device) |
| 1320 | sound_embeds = self.sound_encoder(input_audio_features, feature_attention_mask) |
| 1321 | |
| 1322 | valid_input_lens = feature_attention_mask.sum(dim=1) |
| 1323 | valid_output_lens = self.sound_encoder.encoder._get_subsampling_output_length( |
| 1324 | valid_input_lens |
| 1325 | ).tolist() |
| 1326 | grouped_embeds = [] |
| 1327 | clip_offset = 0 |
| 1328 | for num_clips in audio_num_clips: |
| 1329 | embeds = [] |
| 1330 | for clip_idx in range(clip_offset, clip_offset + num_clips): |
| 1331 | valid_len = valid_output_lens[clip_idx] |
| 1332 | embeds.append(sound_embeds[clip_idx, :valid_len]) |
| 1333 | grouped_embeds.append(torch.cat(embeds, dim=0)) |
| 1334 | clip_offset += num_clips |
| 1335 | |
| 1336 | return tuple(grouped_embeds) |
| 1337 | |
| 1338 | def _create_final_video_embeddings( |
| 1339 | self, |
| 1340 | video_embeddings: torch.Tensor, |
| 1341 | num_tokens_per_frame: list[int], |
| 1342 | frames_indices: list[int], |
| 1343 | frame_duration_ms: int, |
| 1344 | video_temporal_patch_size: int = 1, |
| 1345 | ) -> torch.Tensor: |
| 1346 | """Create final embeddings that combine video embeddings with |
| 1347 | text embeddings of indicator tokens. |
| 1348 | |
| 1349 | These final embeddings contain: |
| 1350 | - Actual video embeddings in positions corresponding to video content |
| 1351 | - Text embeddings for indicator tokens (<img>, </img>, and |
| 1352 | frame separation text) in their respective positions |
| 1353 | |
| 1354 | These embeddings will replace the placeholder embeddings to create |
| 1355 | input_embeds for the LLM. |
| 1356 | """ |
| 1357 | tokenizer = cached_tokenizer_from_config(self.model_config) |
| 1358 | |
| 1359 | # Generate video replacement token IDs using get_video_repl |
| 1360 | # This tokenizes each frame separator independently, then uses pre-tokenized |
| 1361 | # special tokens to ensure consistent tokenization regardless of |
| 1362 | # num_tokens_per_frame values. |
| 1363 | video_repl = NanoNemotronVLProcessor.get_video_repl( |
| 1364 | tokens_per_frame=num_tokens_per_frame, |
| 1365 | frames_indices=frames_indices, |
| 1366 | frame_duration_ms=frame_duration_ms, |
| 1367 | tokenizer=tokenizer, |
| 1368 | img_start_token_ids=self._img_start_token_ids, |
| 1369 | img_end_token_ids=self._img_end_token_ids, |
| 1370 | img_context_token_ids=self._img_context_token_ids, |
| 1371 | video_temporal_patch_size=video_temporal_patch_size, |
| 1372 | ) |
| 1373 | device = video_embeddings.device |
| 1374 | |
| 1375 | # video_repl.full is a list of token IDs |
| 1376 | repl_token_ids = torch.tensor(video_repl.full, device=device) |
| 1377 | |
| 1378 | # Get embedding token IDs for image context (use pre-tokenized version) |
| 1379 | embed_token_ids = torch.tensor(self._img_context_token_ids, device=device) |
| 1380 | |
| 1381 | # Create mask for video embedding positions |
| 1382 | is_video_embed = torch.isin(repl_token_ids, embed_token_ids) |
| 1383 | |
| 1384 | # Create final video embeddings, merging text embeddings for indicator |
| 1385 | # tokens with video embeddings |
| 1386 | |
| 1387 | # LoRA support - |
| 1388 | # These replacement tokens are produced inside the encoder, so they are |
| 1389 | # absent from the request-token batch that the LoRA adapter index |
| 1390 | # mapping is built from, and there can be more |
| 1391 | # of them than max_num_batched_tokens. |
| 1392 | # Embed them with the base weights - |
| 1393 | # the LoRA delta is undefined for tokens with no mapping entry. |
| 1394 | language_model = cast(NemotronHForCausalLM, self.get_language_model()) |
| 1395 | embed_tokens = language_model.model.embed_tokens |
| 1396 | if isinstance(embed_tokens, BaseLayerWithLoRA): |
| 1397 | embed_tokens = embed_tokens.base_layer |
| 1398 | text_embeddings = embed_tokens(repl_token_ids) |
| 1399 | |
| 1400 | final_video_embeddings = _merge_multimodal_embeddings( |
| 1401 | inputs_embeds=text_embeddings, |
| 1402 | multimodal_embeddings=video_embeddings, |
| 1403 | is_multimodal=is_video_embed, |
| 1404 | ) |
| 1405 | |
| 1406 | return final_video_embeddings |
| 1407 | |
| 1408 | def _parse_and_validate_video_input( |
| 1409 | self, **kwargs: object |
| 1410 | ) -> NanoNemotronVLVideoInputs | None: |
| 1411 | pixel_values_flat_video = kwargs.pop("pixel_values_flat_video", None) |
| 1412 | video_num_patches = kwargs.pop("video_num_patches", None) |
| 1413 | video_embeds = kwargs.pop("video_embeds", None) |
| 1414 | frames_indices = kwargs.pop("frames_indices", None) |
| 1415 | frame_duration_ms = kwargs.pop("frame_duration_ms", None) |
| 1416 | |
| 1417 | if pixel_values_flat_video is None and video_embeds is None: |
| 1418 | return None |
| 1419 | |
| 1420 | if video_embeds is not None: |
| 1421 | assert isinstance(video_embeds, (torch.Tensor, list)) |
| 1422 | return NanoNemotronVLVideoEmbeddingInputs( |
| 1423 | type="video_embeds", |
| 1424 | data=video_embeds, |
| 1425 | ) |
| 1426 | |
| 1427 | if pixel_values_flat_video is not None: |
| 1428 | if isinstance(frames_indices, torch.Tensor): |
| 1429 | frames_indices = frames_indices.flatten() |
| 1430 | else: |
| 1431 | assert isinstance(frames_indices, (list, tuple)) |
| 1432 | frames_indices = torch.cat([f.flatten() for f in frames_indices], dim=0) |
| 1433 | |
| 1434 | if isinstance(frame_duration_ms, torch.Tensor): |
| 1435 | frame_duration_ms = frame_duration_ms.flatten() |
| 1436 | else: |
| 1437 | assert isinstance(frame_duration_ms, (list, tuple)) |
| 1438 | frame_duration_ms = torch.cat( |
| 1439 | [f.flatten() for f in frame_duration_ms], dim=0 |
| 1440 | ) |
| 1441 | |
| 1442 | if ( |
| 1443 | isinstance(pixel_values_flat_video, torch.Tensor) |
| 1444 | and pixel_values_flat_video.ndim == 5 |
| 1445 | ): |
| 1446 | # batched._reduce_data stacked same-shape videos into |
| 1447 | # [num_videos, nf, 3, H, W]; unstack back to a list so the |
| 1448 | # same-H,W cat path below handles it uniformly. |
| 1449 | pixel_values_flat_video = list(pixel_values_flat_video) |
| 1450 | |
| 1451 | if not isinstance(pixel_values_flat_video, torch.Tensor): |
| 1452 | assert isinstance(pixel_values_flat_video, (list, tuple)) |
| 1453 | pixel_values_flat_video = torch.cat(pixel_values_flat_video, dim=0) |
| 1454 | |
| 1455 | assert isinstance(video_num_patches, torch.Tensor) |
| 1456 | expected_h = pixel_values_flat_video.shape[-2] |
| 1457 | expected_w = pixel_values_flat_video.shape[-1] |
| 1458 | num_frames = video_num_patches[0].item() |
| 1459 | resolve_bindings = {"h": expected_h, "w": expected_w, "f": num_frames} |
| 1460 | |
| 1461 | return NanoNemotronVLVideoPixelInputs( |
| 1462 | type="pixel_values_videos", |
| 1463 | pixel_values_flat=pixel_values_flat_video, |
| 1464 | num_patches=video_num_patches, |
| 1465 | frames_indices=frames_indices, |
| 1466 | frame_duration_ms=frame_duration_ms, |
| 1467 | resolve_bindings=resolve_bindings, |
| 1468 | ) |
| 1469 | |
| 1470 | raise AssertionError("This line should be unreachable.") |
| 1471 | |
| 1472 | def _parse_and_validate_multimodal_inputs( |
| 1473 | self, **kwargs: object |
| 1474 | ) -> NanoNemotronVLMultiModalInputs: |
| 1475 | modalities: NanoNemotronVLMultiModalInputs = {} |
| 1476 | # Preserve the order of modalities if there are multiple of them |
| 1477 | # from the order of kwargs. |
| 1478 | for input_key in kwargs: |
| 1479 | if ( |
| 1480 | input_key in ("pixel_values_flat", "image_embeds") |
| 1481 | and "images" not in modalities |
| 1482 | ): |
| 1483 | modalities["images"] = self._parse_and_validate_image_input(**kwargs) |
| 1484 | if ( |
| 1485 | input_key in ("pixel_values_flat_video", "video_embeds") |
| 1486 | and "videos" not in modalities |
| 1487 | ): |
| 1488 | modalities["videos"] = self._parse_and_validate_video_input(**kwargs) |
| 1489 | if ( |
| 1490 | input_key |
| 1491 | in ( |
| 1492 | "input_audio_features", |
| 1493 | "feature_attention_mask", |
| 1494 | "audio_num_clips", |
| 1495 | ) |
| 1496 | and "audios" not in modalities |
| 1497 | ): |
| 1498 | input_audio_features = kwargs.get("input_audio_features") |
| 1499 | feature_attention_mask = kwargs.get("feature_attention_mask") |
| 1500 | audio_num_clips = kwargs.get("audio_num_clips") |
| 1501 | assert isinstance(input_audio_features, torch.Tensor) |
| 1502 | assert isinstance(feature_attention_mask, torch.Tensor) |
| 1503 | assert isinstance(audio_num_clips, (list, torch.Tensor)) |
| 1504 | modalities["audios"] = NanoNemotronVLAudioFeatureInputs( |
| 1505 | input_audio_features=input_audio_features, |
| 1506 | feature_attention_mask=feature_attention_mask, |
| 1507 | audio_num_clips=audio_num_clips, |
| 1508 | validate=False, |
| 1509 | ) |
| 1510 | |
| 1511 | return modalities |
| 1512 | |
| 1513 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 1514 | # Validate the multimodal input keyword arguments |
| 1515 | modalities = self._parse_and_validate_multimodal_inputs(**kwargs) |
| 1516 | if modalities is None: |
| 1517 | return [] |
| 1518 | |
| 1519 | # # The result multimodal_embeddings is tuple of tensors, with each |
| 1520 | # tensor corresponding to a multimodal data item (image or video). |
| 1521 | multimodal_embeddings: tuple[torch.Tensor, ...] = () |
| 1522 | |
| 1523 | # NOTE: It is important to iterate over the keys in this dictionary |
| 1524 | # to preserve the order of the modalities. |
| 1525 | for modality in modalities: |
| 1526 | if modality == "images": |
| 1527 | image_input = modalities["images"] |
| 1528 | assert image_input is not None |
| 1529 | if isinstance(image_input, NanoNemotronVLImageEmbeddingInputs): |
| 1530 | image_embeddings = image_input["data"] |
| 1531 | elif isinstance(image_input, NanoNemotronVLImagePixelInputsDynamic): |
| 1532 | image_embeddings = self._process_image_input_dynamic(image_input) |
| 1533 | else: |
| 1534 | image_embeddings = self._process_image_input(image_input) |
| 1535 | multimodal_embeddings += tuple(image_embeddings) |
| 1536 | if modality == "videos": |
| 1537 | video_input = modalities["videos"] |
| 1538 | assert video_input is not None |
| 1539 | if video_input["type"] == "video_embeds": |
| 1540 | video_embeddings = video_input["data"] |
| 1541 | else: |
| 1542 | assert isinstance(video_input, NanoNemotronVLVideoPixelInputs) |
| 1543 | video_embeddings = self._process_video_input(video_input) |
| 1544 | multimodal_embeddings += tuple(video_embeddings) |
| 1545 | if modality == "audios": |
| 1546 | audio_input = modalities["audios"] |
| 1547 | assert audio_input is not None |
| 1548 | audio_embeddings = self._process_audio_input(audio_input) |
| 1549 | multimodal_embeddings += tuple(audio_embeddings) |
| 1550 | |
| 1551 | return multimodal_embeddings |
| 1552 | |
| 1553 | def forward( |
| 1554 | self, |
| 1555 | input_ids: torch.Tensor | None, |
| 1556 | positions: torch.Tensor, |
| 1557 | intermediate_tensors: IntermediateTensors | None = None, |
| 1558 | inputs_embeds: torch.Tensor | None = None, |
| 1559 | **kwargs: object, |
| 1560 | ) -> torch.Tensor | IntermediateTensors: |
| 1561 | if intermediate_tensors is not None: |
| 1562 | inputs_embeds = None |
| 1563 | |
| 1564 | hidden_states = self.language_model( |
| 1565 | input_ids=input_ids, |
| 1566 | positions=positions, |
| 1567 | intermediate_tensors=intermediate_tensors, |
| 1568 | inputs_embeds=inputs_embeds, |
| 1569 | **kwargs, |
| 1570 | ) |
| 1571 | |
| 1572 | return hidden_states |
| 1573 | |
| 1574 | def get_mm_mapping(self) -> MultiModelKeys: |
| 1575 | """Get the module prefix in multimodal models""" |
| 1576 | return MultiModelKeys.from_string_field( |
| 1577 | language_model="language_model", |
| 1578 | connector=["mlp1", "sound_encoder.projection"], |
| 1579 | tower_model=["vision_model", "sound_encoder.encoder"], |
| 1580 | ) |
| 1581 | |
| 1582 | def compute_logits( |
| 1583 | self, |
| 1584 | hidden_states: torch.Tensor, |
| 1585 | ) -> torch.Tensor | None: |
| 1586 | return self.language_model.compute_logits(hidden_states) |
| 1587 | |
| 1588 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): |
| 1589 | mm_config = self.model_config.multimodal_config |
| 1590 | assert mm_config is not None |
| 1591 | load_multimodal_weights = not all( |
| 1592 | mm_config.get_limit_per_prompt(modality) == 0 |
| 1593 | for modality in ("image", "video", "audio") |
| 1594 | ) |
| 1595 | adapter_dict = dict(self.mlp1.named_parameters()) |
| 1596 | |
| 1597 | def is_llm(name: str) -> bool: |
| 1598 | return name.startswith("language_model") |
| 1599 | |
| 1600 | def is_adapter_weights(weight: tuple[str, torch.Tensor]): |
| 1601 | return weight[0].startswith("mlp1") |
| 1602 | |
| 1603 | def is_vision_weights(name: str) -> bool: |
| 1604 | return name.startswith("vision_model.radio_model.") |
| 1605 | |
| 1606 | def is_sound_weights(name: str) -> bool: |
| 1607 | return name.startswith("sound") |
| 1608 | |
| 1609 | # LLM weights (the bulk of the model) are streamed lazily through a |
| 1610 | # generator so each tensor is copied into its parameter before the |
| 1611 | # iterator advances, avoiding stale-reference corruption with |
| 1612 | # reusable-buffer streamers. The smaller mm components (mlp1, vision, |
| 1613 | # sound) are detach+cloned on append so they are independent of any |
| 1614 | # reusable buffer the streamer may use, then loaded after the LLM. |
| 1615 | adapter_weights: list[tuple[str, torch.Tensor]] = [] |
| 1616 | vision_weights: list[tuple[str, torch.Tensor]] = [] |
| 1617 | sound_weights: list[tuple[str, torch.Tensor]] = [] |
| 1618 | |
| 1619 | def llm_weights_gen(): |
| 1620 | for name, w in weights: |
| 1621 | if is_llm(name): |
| 1622 | # Strip 'language_model.' prefix for LLM weights |
| 1623 | yield ".".join(name.split(".")[1:]), w |
| 1624 | elif is_adapter_weights((name, w)): |
| 1625 | if not load_multimodal_weights: |
| 1626 | continue |
| 1627 | trimmed_name = ".".join(name.split(".")[1:]) |
| 1628 | adapter_weights.append((trimmed_name, w.detach().clone())) |
| 1629 | elif is_vision_weights(name): |
| 1630 | if not load_multimodal_weights: |
| 1631 | continue |
| 1632 | # Convert: vision_model.radio_model.* → radio_model.* |
| 1633 | hf_key = name[len("vision_model.") :] |
| 1634 | vision_weights.append((hf_key, w.detach().clone())) |
| 1635 | elif is_sound_weights(name): |
| 1636 | if not load_multimodal_weights: |
| 1637 | continue |
| 1638 | assert self.sound_encoder is not None |
| 1639 | sound_weights.append((name, w.detach().clone())) |
| 1640 | |
| 1641 | # Fully drain the generator so every mm tensor is buffered, even if |
| 1642 | # the LLM loader stops iterating early. |
| 1643 | llm_weights_iter = llm_weights_gen() |
| 1644 | self.language_model.load_weights(llm_weights_iter) |
| 1645 | for _ in llm_weights_iter: |
| 1646 | pass |
| 1647 | |
| 1648 | if load_multimodal_weights: |
| 1649 | for trimmed_name, w in adapter_weights: |
| 1650 | # Load vision-language adapter weights directly |
| 1651 | param = adapter_dict[trimmed_name] |
| 1652 | with torch.no_grad(): |
| 1653 | default_weight_loader(param, w) |
| 1654 | self.vision_model.load_weights(vision_weights) |
| 1655 | if self.sound_encoder is not None and len(sound_weights) > 0: |
| 1656 | self.sound_encoder.load_weights(sound_weights) |
| 1657 | |
| 1658 | def get_vit_model_from_radio_config(self, hf_config): |
| 1659 | hf_config_vision = hf_config.vision_config |
| 1660 | model_name = hf_config_vision.args.get("model") |
| 1661 | if model_name is None: |
| 1662 | raise ValueError(f"Unsupported vit model type: {model_name}") |
| 1663 | |
| 1664 | preferred_resolution = getattr(hf_config_vision, "preferred_resolution", None) |
| 1665 | image_size = preferred_resolution[0] if preferred_resolution else 224 |
| 1666 | patch_size = getattr(hf_config_vision, "patch_size", 16) |
| 1667 | |
| 1668 | # video_temporal_patch_size and separate_video_embedder are |
| 1669 | # top-level vision_config attributes, not inside args. |
| 1670 | video_temporal_patch_size = getattr( |
| 1671 | hf_config_vision, "video_temporal_patch_size", 1 |
| 1672 | ) |
| 1673 | separate_video_embedder = getattr( |
| 1674 | hf_config_vision, "separate_video_embedder", True |
| 1675 | ) |
| 1676 | |
| 1677 | radio_config = RadioConfig( |
| 1678 | model_name=model_name, |
| 1679 | image_size=image_size, |
| 1680 | patch_size=patch_size, |
| 1681 | norm_mean=hf_config.norm_mean, |
| 1682 | norm_std=hf_config.norm_std, |
| 1683 | video_temporal_patch_size=video_temporal_patch_size, |
| 1684 | separate_video_embedder=separate_video_embedder, |
| 1685 | **hf_config_vision.args, |
| 1686 | ) |
| 1687 | |
| 1688 | return RadioModel(config=radio_config) |
| 1689 | |
| 1690 | def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): |
| 1691 | return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs( |
| 1692 | input_buffers, **kwargs |
| 1693 | ) |
| 1694 | |
| 1695 | def get_seqlen_agnostic_capture_inputs(self, batch_size: int): |
| 1696 | return self.language_model.mamba_cache.get_seqlen_agnostic_capture_inputs( |
| 1697 | batch_size |
| 1698 | ) |
| 1699 | |
| 1700 | @classmethod |
| 1701 | def get_mamba_state_shape_from_config(cls, vllm_config: "VllmConfig"): |
| 1702 | text_config = vllm_config.model_config.hf_config.text_config |
| 1703 | temp_vllm_config = vllm_config.with_hf_config(text_config) |
| 1704 | return NemotronHForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) |
| 1705 | |
| 1706 | @classmethod |
| 1707 | def get_mamba_state_dtype_from_config(cls, vllm_config: "VllmConfig"): |
| 1708 | text_config = vllm_config.model_config.hf_config.text_config |
| 1709 | temp_vllm_config = vllm_config.with_hf_config(text_config) |
| 1710 | return NemotronHForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) |
| 1711 | |
| 1712 | @classmethod |
| 1713 | def get_mamba_state_copy_func(cls): |
| 1714 | return NemotronHForCausalLM.get_mamba_state_copy_func() |