import contextlib import sys from pathlib import Path from typing import Optional import einops import numpy as np import torch import torch.nn as nn from torch import Tensor from transformers import Qwen2ForCausalLM from transformers.cache_utils import Cache from transformers.generation.utils import GenerationMixin from transformers.modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel from autogaze.vision_encoders.siglip.modeling_siglip import SiglipVisionModel from .configuration_nvila import NVILAConfig MM_HIDDEN_SIZE = 1152 class TokenShuffle(nn.Module): """Token shuffle module that groups tokens and concatenates their features.""" def __init__(self, shuffle_num: int): super().__init__() self.shuffle_num = shuffle_num def forward(self, x: Tensor) -> Tensor: """ Args: x: (B, N, C) tensor where B is batch size, N is sequence length, C is hidden size Returns: (B, N', C * shuffle_num) tensor where N' = ceil(N / shuffle_num) """ # x: (B, N, C) if x.shape[1] % self.shuffle_num != 0: # Pad with the last token to make sequence length divisible by shuffle_num pad_size = self.shuffle_num - (x.shape[1] % self.shuffle_num) x = torch.cat([x, x[:, -1:].repeat(1, pad_size, 1)], dim=1) # Rearrange: (B, N, C) -> (B, N//k, k*C) where k = shuffle_num return einops.rearrange(x, "b (n k) c -> b n (k c)", k=self.shuffle_num) class NVILAMultiModalProjector(nn.Module): """Multi-modal projector using mlp_shuffle_9 architecture.""" def __init__(self, config: NVILAConfig): super().__init__() self.layers = nn.Sequential( TokenShuffle(9), nn.LayerNorm(MM_HIDDEN_SIZE * 9), nn.Linear(MM_HIDDEN_SIZE * 9, MM_HIDDEN_SIZE * 3), nn.GELU(), nn.LayerNorm(MM_HIDDEN_SIZE * 3), nn.Linear(MM_HIDDEN_SIZE * 3, config.text_config.hidden_size), nn.GELU(), nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size), ) def forward(self, x: Tensor) -> Tensor: return self.layers(x) class NVILAForConditionalGeneration(PreTrainedModel, GenerationMixin): config_class = NVILAConfig base_model_prefix: str = "llm" _auto_class = "AutoModel" _supports_flash_attn_2 = True _supports_sdpa = True def __init__(self, config: NVILAConfig): super().__init__(config) self.config: NVILAConfig @contextlib.contextmanager def default_torch_dtype(dtype): original_dtype = torch.get_default_dtype() torch.set_default_dtype(dtype) try: yield finally: torch.set_default_dtype(original_dtype) with default_torch_dtype(config.torch_dtype): self.vision_tower = SiglipVisionModel(config.vision_config) self.mm_projector = NVILAMultiModalProjector(config) self.llm = Qwen2ForCausalLM(config.text_config) self.post_init() def forward( self, *, input_ids: Tensor | None = None, inputs_embeds: Tensor | None = None, pixel_values: Tensor | None = None, pixel_values_images_tiles: list[Tensor] | None = None, pixel_values_images_thumbnails: list[Tensor] | None = None, num_spatial_tiles_each_image: list[int] | None = None, pixel_values_videos_tiles: list[Tensor] | None = None, pixel_values_videos_thumbnails: list[Tensor] | None = None, gazing_info: dict | None = None, num_spatial_tiles_each_video: list[int] | None = None, **kwargs, ) -> CausalLMOutputWithPast: assert (input_ids is None) != ( inputs_embeds is None ), "Exactly one of `input_ids` or `inputs_embeds` must be specified." # Pop processor-only fields that the LLM should not see kwargs.pop("pixel_values_videos_tiles_autogaze", None) kwargs.pop("pixel_values_videos_thumbnails_autogaze", None) kwargs.pop("pixel_values_videos", None) if input_ids is not None and torch.any( torch.isin( input_ids, torch.tensor( [self.config.image_token_id, self.config.video_token_id], device=input_ids.device, ), ).any() ): # Prefill # Extract fields from kwargs if not passed as explicit args if gazing_info is None: gazing_info = kwargs.pop("gazing_info", None) if pixel_values_images_tiles is None: pixel_values_images_tiles = kwargs.pop("pixel_values_images_tiles", None) if pixel_values_images_thumbnails is None: pixel_values_images_thumbnails = kwargs.pop("pixel_values_images_thumbnails", None) if num_spatial_tiles_each_image is None: num_spatial_tiles_each_image = kwargs.pop("num_spatial_tiles_each_image", None) if pixel_values_videos_tiles is None: pixel_values_videos_tiles = kwargs.pop("pixel_values_videos_tiles", None) if pixel_values_videos_thumbnails is None: pixel_values_videos_thumbnails = kwargs.pop("pixel_values_videos_thumbnails", None) if num_spatial_tiles_each_video is None: num_spatial_tiles_each_video = kwargs.pop("num_spatial_tiles_each_video", None) inputs_embeds = self._embed( input_ids=input_ids, pixel_values=pixel_values, pixel_values_images_tiles=pixel_values_images_tiles, pixel_values_images_thumbnails=pixel_values_images_thumbnails, num_spatial_tiles_each_image=num_spatial_tiles_each_image, pixel_values_videos_tiles=pixel_values_videos_tiles, pixel_values_videos_thumbnails=pixel_values_videos_thumbnails, gazing_info=gazing_info, num_spatial_tiles_each_video=num_spatial_tiles_each_video, ) input_ids = None outputs = self.llm( input_ids=input_ids, inputs_embeds=inputs_embeds, **kwargs, ) return outputs def _embed( self, *, input_ids: Tensor, pixel_values: Tensor | None, pixel_values_images_tiles: list[Tensor] | None, pixel_values_images_thumbnails: list[Tensor] | None, num_spatial_tiles_each_image: list[int] | None, pixel_values_videos_tiles: list[Tensor] | None, pixel_values_videos_thumbnails: list[Tensor] | None, gazing_info: dict | None = None, num_spatial_tiles_each_video: list[int] | None = None, ) -> Tensor: inputs_embeds: Tensor = self.llm.model.embed_tokens(input_ids) # Handle images if pixel_values_images_tiles is not None and len(pixel_values_images_tiles) > 0: per_image_features = self._encode_images( pixel_values_images_tiles=pixel_values_images_tiles, pixel_values_images_thumbnails=pixel_values_images_thumbnails, num_spatial_tiles_each_image=num_spatial_tiles_each_image, ) all_features = torch.cat(per_image_features, dim=0) image_token_mask = input_ids == self.config.image_token_id num_image_tokens = image_token_mask.sum().item() num_image_features = all_features.shape[0] assert num_image_features == num_image_tokens, ( f"Number of image features {num_image_features} does not match " f"number of image tokens {num_image_tokens}" ) inputs_embeds[image_token_mask] = all_features.to(inputs_embeds.dtype) # Handle videos if pixel_values_videos_tiles is not None: per_video_features = self._encode_vision( pixel_values_videos_tiles=pixel_values_videos_tiles, pixel_values_videos_thumbnails=pixel_values_videos_thumbnails, gazing_info=gazing_info, num_spatial_tiles_each_video=num_spatial_tiles_each_video, ) # per_video_features: list of (num_tokens_i, llm_hidden) tensors all_features = torch.cat(per_video_features, dim=0) # Match vision features to video tokens video_token_mask = input_ids == self.config.video_token_id num_video_tokens = video_token_mask.sum().item() num_vision_features = all_features.shape[0] assert num_vision_features == num_video_tokens, ( f"Number of vision features {num_vision_features} does not match " f"number of video tokens {num_video_tokens}" ) inputs_embeds[video_token_mask] = all_features.to(inputs_embeds.dtype) return inputs_embeds def _make_default_gazing_info( self, total_items: int, T: int, device: torch.device, ) -> dict: """Create gazing_info that gazes at every patch (no reduction). Args: total_items: Number of items (tiles or thumbnails) in the batch. T: Temporal frames per item. device: Target torch device. Returns: gazing_info dict with ``gazing_pos``, ``num_gazing_each_frame``, ``if_padded_gazing``. """ image_size = self.vision_tower.config.image_size patch_size = self.vision_tower.config.patch_size scales = sorted( int(s) for s in self.vision_tower.config.scales.split("+") ) num_patches_each_scale = [(s // patch_size) ** 2 for s in scales] total_patches_per_frame = sum(num_patches_each_scale) # Gazing positions: all patches for every frame per_item_pos = [] for t in range(T): start = t * total_patches_per_frame per_item_pos.append( torch.arange(start, start + total_patches_per_frame, device=device, dtype=torch.long) ) per_item_pos = torch.cat(per_item_pos) # (T * total_patches_per_frame,) gazing_pos = per_item_pos.unsqueeze(0).expand(total_items, -1) # (B, N) num_gazing_each_frame = torch.full( (T,), total_patches_per_frame, device=device, dtype=torch.long ) if_padded_gazing = torch.zeros_like(gazing_pos, dtype=torch.bool) return { "gazing_pos": gazing_pos, "num_gazing_each_frame": num_gazing_each_frame, "if_padded_gazing": if_padded_gazing, } def _encode_images( self, pixel_values_images_tiles: list[Tensor], pixel_values_images_thumbnails: list[Tensor] | None, num_spatial_tiles_each_image: list[int], ) -> list[Tensor]: """Encode image tiles + thumbnails and return projected features per image. Each image is a set of spatial tiles plus one thumbnail (T=1 each). All patches are kept (no gazing reduction). For each image the spatial tiles are merged into one effective frame, the thumbnail forms a second effective frame, and both are padded to ``shuffle_num`` before projection through the mm_projector. Args: pixel_values_images_tiles: Per-image tile tensors, each ``(num_tiles_i, 1, C, H, W)``. pixel_values_images_thumbnails: Per-image thumbnail tensors, each ``(1, 1, C, H, W)``. May be ``None``. num_spatial_tiles_each_image: Number of spatial tiles per image. Returns: List of tensors (one per image), each ``(num_tokens_i, llm_hidden)``. """ shuffle_num = 9 device = self.vision_tower.device # --- Run vision tower on all tiles --- all_tiles = torch.cat(pixel_values_images_tiles, dim=0) # (total_tiles, 1, C, H, W) total_tiles = all_tiles.shape[0] gi_tiles = self._make_default_gazing_info(total_tiles, 1, device) tiles_features = self._run_vision_tower_batched(all_tiles, gi_tiles) # (total_tiles, N, H) num_gaze_tiles = gi_tiles["num_gazing_each_frame"] # (1,) if_padded_tiles = gi_tiles["if_padded_gazing"] # (total_tiles, N) frame_lens_tiles = num_gaze_tiles.tolist() tile_feats: list[Tensor] = [] for idx in range(total_tiles): feats = tiles_features[idx] pad_mask = if_padded_tiles[idx] frame_feats = feats.split(frame_lens_tiles, dim=0) frame_pads = pad_mask.split(frame_lens_tiles, dim=0) tile_feats.append( torch.cat([f[~p] for f, p in zip(frame_feats, frame_pads)], dim=0) ) # --- Run vision tower on all thumbnails --- thumb_feats: list[Tensor] | None = None if pixel_values_images_thumbnails is not None and len(pixel_values_images_thumbnails) > 0: all_thumbs = torch.cat(pixel_values_images_thumbnails, dim=0) # (num_images, 1, C, H, W) total_thumbs = all_thumbs.shape[0] gi_thumbs = self._make_default_gazing_info(total_thumbs, 1, device) thumbs_features = self._run_vision_tower_batched(all_thumbs, gi_thumbs) num_gaze_thumbs = gi_thumbs["num_gazing_each_frame"] if_padded_thumbs = gi_thumbs["if_padded_gazing"] frame_lens_thumbs = num_gaze_thumbs.tolist() thumb_feats = [] for idx in range(total_thumbs): feats = thumbs_features[idx] pad_mask = if_padded_thumbs[idx] frame_feats = feats.split(frame_lens_thumbs, dim=0) frame_pads = pad_mask.split(frame_lens_thumbs, dim=0) thumb_feats.append( torch.cat([f[~p] for f, p in zip(frame_feats, frame_pads)], dim=0) ) # --- Build per-image sequences --- tile_offset = 0 per_image_sequences: list[Tensor] = [] per_image_token_counts: list[int] = [] for img_idx, ns in enumerate(num_spatial_tiles_each_image): effective_frames: list[Tensor] = [] # Tiles effective frame: merge all spatial tiles spatial_feats = tile_feats[tile_offset : tile_offset + ns] tile_offset += ns effective_frames.append(torch.cat(spatial_feats, dim=0)) # Thumbnail effective frame if thumb_feats is not None: effective_frames.append(thumb_feats[img_idx]) # Pad each effective frame to divisible by shuffle_num padded_frames: list[Tensor] = [] for frame in effective_frames: n = frame.shape[0] pad = (shuffle_num - n % shuffle_num) % shuffle_num if pad > 0: frame = torch.cat([frame, frame[-1:].expand(pad, -1)], dim=0) padded_frames.append(frame) image_seq = torch.cat(padded_frames, dim=0) per_image_sequences.append(image_seq) per_image_token_counts.append(image_seq.shape[0] // shuffle_num) all_features = torch.cat(per_image_sequences, dim=0).unsqueeze(0) projected = self.mm_projector( all_features.to(device=self.device, dtype=self.dtype) ) projected = projected.squeeze(0) return list(projected.split(per_image_token_counts, dim=0)) def _run_vision_tower_batched( self, all_pixels: Tensor, gazing_info_batch: dict, ) -> Tensor: """Run the vision tower in minibatches and concatenate features. Args: all_pixels: ``(B, T, C, H, W)`` tensor. gazing_info_batch: Dict with ``gazing_pos`` ``(B, N)``, ``if_padded_gazing`` ``(B, N)``, and ``num_gazing_each_frame`` ``(T,)`` (shared across batch). Returns: ``(B, N, H)`` hidden features from the second-to-last layer. """ device = self.vision_tower.device dtype = self.vision_tower.dtype total = all_pixels.shape[0] bs = self.config.max_batch_size_siglip if total <= bs: out: BaseModelOutputWithPooling = self.vision_tower( all_pixels.to(device=device, dtype=dtype), gazing_info=gazing_info_batch, output_hidden_states=True, ) assert out.hidden_states is not None return out.hidden_states[-2] num_gaze_shared = gazing_info_batch["num_gazing_each_frame"] all_pos = gazing_info_batch["gazing_pos"] all_pad = gazing_info_batch["if_padded_gazing"] feature_chunks: list[Tensor] = [] for start in range(0, total, bs): end = min(start + bs, total) mini_gi = { "gazing_pos": all_pos[start:end], "if_padded_gazing": all_pad[start:end], "num_gazing_each_frame": num_gaze_shared, } out = self.vision_tower( all_pixels[start:end].to(device=device, dtype=dtype), gazing_info=mini_gi, output_hidden_states=True, ) assert out.hidden_states is not None feature_chunks.append(out.hidden_states[-2]) return torch.cat(feature_chunks, dim=0) def _encode_vision( self, pixel_values_videos_tiles: list[Tensor], pixel_values_videos_thumbnails: list[Tensor], gazing_info: dict | None, num_spatial_tiles_each_video: list[int], ) -> list[Tensor]: """Encode tiles and thumbnails and return projected features per video. Workflow ------- 1. Batch all tiles / thumbnails across videos and run the vision tower (in minibatches controlled by ``config.max_batch_size_siglip``). 2. Remove padded gazing features. 3. Re-order per video: for each global temporal frame gather all spatial tiles, then append thumbnail frames. 4. Pad each effective frame to be divisible by ``shuffle_num`` (9). 5. Concatenate all videos into a single sequence (batch=1), project through ``mm_projector``, then split back per video. Args: pixel_values_videos_tiles: Per-video tile tensors, each ``(num_tiles_i, T_tile, C, H, W)``. pixel_values_videos_thumbnails: Per-video thumbnail tensors, each ``(T_thumb_i, 1, C, H, W)``. gazing_info: Dict produced by the processor containing per-video gazing data for tiles and thumbnails. ``None`` triggers default "gaze at all patches" behaviour. num_spatial_tiles_each_video: Number of spatial tiles per video. Returns: List of tensors (one per video), each ``(num_tokens_i, llm_hidden)``. """ shuffle_num = 9 # must match TokenShuffle in NVILAMultiModalProjector device = self.vision_tower.device dtype = self.vision_tower.dtype num_videos = len(pixel_values_videos_tiles) num_tiles_per_video = [t.shape[0] for t in pixel_values_videos_tiles] num_thumbs_per_video = [t.shape[0] for t in pixel_values_videos_thumbnails] # ---- 1. Batch & run vision tower on tiles ---- all_tiles = torch.cat(pixel_values_videos_tiles, dim=0) # (total_tiles, T_tile, C, H, W) T_tile = all_tiles.shape[1] if gazing_info is not None: tiles_nge = gazing_info["num_gazing_each_frame_tiles"] ref = tiles_nge[0][0] assert all( torch.equal(t[0], ref) for t in tiles_nge ), "num_gazing_each_frame must be identical across all videos for tiles" tiles_gi = { "gazing_pos": torch.cat(gazing_info["gazing_pos_tiles"], dim=0).to(device), "num_gazing_each_frame": gazing_info["num_gazing_each_frame_tiles"][0][0].to(device), "if_padded_gazing": torch.cat(gazing_info["if_padded_gazing_tiles"], dim=0).to(device), } else: tiles_gi = self._make_default_gazing_info(all_tiles.shape[0], T_tile, device) tiles_features = self._run_vision_tower_batched(all_tiles, tiles_gi) # (total_tiles, N, H) # ---- 2. Batch & run vision tower on thumbnails ---- all_thumbs = torch.cat(pixel_values_videos_thumbnails, dim=0) # (total_thumbs, 1, C, H, W) if gazing_info is not None: thumbs_nge = gazing_info["num_gazing_each_frame_thumbnails"] ref = thumbs_nge[0][0] assert all( torch.equal(t[0], ref) for t in thumbs_nge ), "num_gazing_each_frame must be identical across all videos for thumbnails" thumbs_gi = { "gazing_pos": torch.cat(gazing_info["gazing_pos_thumbnails"], dim=0).to(device), "num_gazing_each_frame": gazing_info["num_gazing_each_frame_thumbnails"][0][0].to(device), "if_padded_gazing": torch.cat(gazing_info["if_padded_gazing_thumbnails"], dim=0).to(device), } else: thumbs_gi = self._make_default_gazing_info(all_thumbs.shape[0], 1, device) thumbs_features = self._run_vision_tower_batched(all_thumbs, thumbs_gi) # (total_thumbs, N', H) # ---- 3. Remove padded features & split by frame ---- # For each tile: list of T_tile tensors, each (n_i, hidden) all_tiles_if_padded = tiles_gi["if_padded_gazing"] all_tiles_num_gaze = tiles_gi["num_gazing_each_frame"] # 1-D (T_tile,) tiles_frame_lens = all_tiles_num_gaze.tolist() all_tiles_frame_feats: list[list[Tensor]] = [] for idx in range(tiles_features.shape[0]): feats = tiles_features[idx] # (N, hidden) pad_mask = all_tiles_if_padded[idx] # (N,) frame_feats = feats.split(tiles_frame_lens, dim=0) frame_pads = pad_mask.split(tiles_frame_lens, dim=0) all_tiles_frame_feats.append( [f[~p] for f, p in zip(frame_feats, frame_pads)] ) # For each thumbnail: list with 1 tensor (n_i, hidden) all_thumbs_if_padded = thumbs_gi["if_padded_gazing"] all_thumbs_num_gaze = thumbs_gi["num_gazing_each_frame"] # 1-D (1,) thumbs_frame_lens = all_thumbs_num_gaze.tolist() all_thumbs_frame_feats: list[list[Tensor]] = [] for idx in range(thumbs_features.shape[0]): feats = thumbs_features[idx] pad_mask = all_thumbs_if_padded[idx] frame_feats = feats.split(thumbs_frame_lens, dim=0) frame_pads = pad_mask.split(thumbs_frame_lens, dim=0) all_thumbs_frame_feats.append( [f[~p] for f, p in zip(frame_feats, frame_pads)] ) # ---- 4. Per-video: reorder, pad frames, build sequences ---- tile_offset = 0 thumb_offset = 0 per_video_sequences: list[Tensor] = [] per_video_token_counts: list[int] = [] for vid_idx in range(num_videos): ns = num_spatial_tiles_each_video[vid_idx] nt = num_tiles_per_video[vid_idx] tc = nt // ns # temporal chunks total_frames = tc * T_tile n_thumbs = num_thumbs_per_video[vid_idx] vid_tile_feats = all_tiles_frame_feats[tile_offset: tile_offset + nt] tile_offset += nt vid_thumb_feats = all_thumbs_frame_feats[thumb_offset: thumb_offset + n_thumbs] thumb_offset += n_thumbs # -- Reorder tile features to frame-first -- # Tiles from processor are ordered: # chunk0: [S0, S1, ..., S_{ns-1}], chunk1: [S0, ...], ... # We want: for each global frame g, cat all spatial tiles. effective_frames: list[Tensor] = [] for g in range(total_frames): chunk = g // T_tile f_in_chunk = g % T_tile spatial_feats = [ vid_tile_feats[chunk * ns + s][f_in_chunk] for s in range(ns) ] effective_frames.append(torch.cat(spatial_feats, dim=0)) # -- Append thumbnail frames -- for thumb in vid_thumb_feats: effective_frames.append(thumb[0]) # single frame # -- Pad each effective frame to divisible by shuffle_num -- padded_frames: list[Tensor] = [] for frame in effective_frames: n = frame.shape[0] pad = (shuffle_num - n % shuffle_num) % shuffle_num if pad > 0: padded_frame = torch.cat( [frame, frame[-1:].expand(pad, -1)], dim=0 ) else: padded_frame = frame padded_frames.append(padded_frame) video_seq = torch.cat(padded_frames, dim=0) # (total_padded, hidden) per_video_sequences.append(video_seq) per_video_token_counts.append(video_seq.shape[0] // shuffle_num) # ---- 5. Concat all videos, project, split back ---- all_features = torch.cat(per_video_sequences, dim=0).unsqueeze(0) # (1, total, hidden) projected = self.mm_projector( all_features.to(device=self.device, dtype=self.dtype) ) # (1, total // shuffle_num, llm_hidden) projected = projected.squeeze(0) # (total // shuffle_num, llm_hidden) per_video_features = list(projected.split(per_video_token_counts, dim=0)) return per_video_features