# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/video_llama_3/modular_video_llama_3.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_video_llama_3.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2025 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from collections.abc import Callable from dataclasses import dataclass from typing import Any import torch import torch.nn as nn from torch.nn import LayerNorm from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check from ...utils.generic import is_flash_attention_requested, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from ..auto.modeling_auto import AutoModel from .configuration_video_llama_3 import VideoLlama3Config, VideoLlama3VisionConfig class VideoLlama3VisionRotaryEmbedding(nn.Module): inv_freq: torch.Tensor # fix linting for `register_buffer` def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() self.dim = dim self.theta = theta inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, grid_thw, merge_sizes) -> tuple[torch.Tensor, torch.Tensor]: pos_ids = [] for (t, h, w), merge_size in zip(grid_thw, merge_sizes): hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) hpos_ids = hpos_ids.reshape( h // merge_size, merge_size, w // merge_size, merge_size, ) hpos_ids = hpos_ids.permute(0, 2, 1, 3) hpos_ids = hpos_ids.flatten() wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) wpos_ids = wpos_ids.reshape( h // merge_size, merge_size, w // merge_size, merge_size, ) wpos_ids = wpos_ids.permute(0, 2, 1, 3) wpos_ids = wpos_ids.flatten() pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) pos_ids = torch.cat(pos_ids, dim=0) max_grid_thw = grid_thw[:, 1:].max() seq = torch.arange(max_grid_thw, device=self.inv_freq.device, dtype=self.inv_freq.dtype) rotary_pos_emb_full = torch.outer(seq, self.inv_freq) rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) return (emb.cos(), emb.sin()) class VideoLlama3VisionEmbeddings(nn.Module): def __init__(self, config: VideoLlama3VisionConfig) -> None: super().__init__() self.config = config self.embed_dim = config.hidden_size self.patch_size = config.patch_size self.patch_embedding = nn.Conv2d( in_channels=config.num_channels, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size, padding="valid", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states.view(-1, self.config.num_channels, self.patch_size, self.patch_size) patch_embeds = self.patch_embedding(hidden_states) embeddings = patch_embeds.view(-1, self.embed_dim) return embeddings class VideoLlama3VisionMLP(nn.Module): def __init__(self, config): super().__init__() self.config = config self.activation_fn = ACT2FN[config.hidden_act] self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states = self.fc2(hidden_states) return hidden_states def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, scaling: float, dropout: float = 0.0, **kwargs, ): key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def apply_rotary_pos_emb_vision( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: orig_q_dtype = q.dtype orig_k_dtype = k.dtype q, k = q.float(), k.float() cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) q_embed = q_embed.to(orig_q_dtype) k_embed = k_embed.to(orig_k_dtype) return q_embed, k_embed class VideoLlama3VisionAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config): super().__init__() self.config = config self.embed_dim = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.embed_dim // self.num_heads if self.head_dim * self.num_heads != self.embed_dim: raise ValueError( f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" f" {self.num_heads})." ) self.is_causal = False self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) self.num_key_value_groups = 1 self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Args: hidden_states (`torch.Tensor`): Input to the layer of shape `(seq_len, embed_dim)`. cu_seqlens (`torch.Tensor` of shape `(num_images_or_videos + 1,)`): The cumulative sequence lengths of each image or video feature. position_embeddings (`tuple(torch.Tensor, torch.Tensor)` of shape `(num_patches, head_dim // 2)`): The cosine and sine position embeddings for vision attention. """ seq_length = hidden_states.shape[0] query_states = self.q_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim) key_states = self.k_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim) value_states = self.v_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) query_states = query_states.transpose(0, 1).unsqueeze(0) key_states = key_states.transpose(0, 1).unsqueeze(0) value_states = value_states.transpose(0, 1).unsqueeze(0) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) if is_flash_attention_requested(self.config): # Flash Attention 2: Use cu_seqlens for variable length attention max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask=None, scaling=self.scaling, dropout=0.0 if not self.training else self.attention_dropout, cu_seq_lens_q=cu_seqlens, cu_seq_lens_k=cu_seqlens, max_length_q=max_seqlen, max_length_k=max_seqlen, is_causal=False, **kwargs, ) else: # Other implementations: Process each chunk separately lengths = cu_seqlens[1:] - cu_seqlens[:-1] splits = [ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) ] attn_outputs, attn_weights = [], [] for q, k, v in zip(*splits): attn_output, attn_weight = attention_interface( self, q, k, v, attention_mask=None, scaling=self.scaling, dropout=0.0 if not self.training else self.attention_dropout, is_causal=False, **kwargs, ) attn_outputs.append(attn_output) attn_weights.append(attn_weight) attn_output = torch.cat(attn_outputs, dim=1) attn_output = attn_output.reshape(seq_length, -1).contiguous() attn_output = self.out_proj(attn_output) return attn_output, attn_weights class VideoLlama3VisionEncoderLayer(GradientCheckpointingLayer): def __init__(self, config: VideoLlama3VisionConfig): super().__init__() self.embed_dim = config.hidden_size self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.self_attn = VideoLlama3VisionAttention(config=config) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.mlp = VideoLlama3VisionMLP(config=config) @auto_docstring def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: r""" cu_seqlens (`torch.Tensor` of shape `(num_images_or_videos + 1,)`): The cumulative sequence lengths of each image or video feature. position_embeddings (`tuple(torch.Tensor, torch.Tensor)` of shape `(num_patches, head_dim // 2)`): The cosine and sine position embeddings for vision attention. """ residual = hidden_states hidden_states = self.layer_norm1(hidden_states) hidden_states, _ = self.self_attn( hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.layer_norm2(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class VideoLlama3VisionEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`VideoLlama3VisionEncoderLayer`]. Args: config: VideoLlama3VisionConfig """ def __init__(self, config: VideoLlama3VisionConfig): super().__init__() self.config = config self.layers = nn.ModuleList([VideoLlama3VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)]) self.gradient_checkpointing = False # Ignore copy @can_return_tuple @auto_docstring def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutput: r""" cu_seqlens (`torch.Tensor` of shape `(num_images_or_videos + 1,)`): The cumulative sequence lengths of each image or video feature. position_embeddings (`tuple(torch.Tensor, torch.Tensor)` of shape `(num_patches, head_dim // 2)`): The cosine and sine position embeddings for vision attention. """ for encoder_layer in self.layers: hidden_states = encoder_layer( hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, **kwargs, ) return BaseModelOutput(last_hidden_state=hidden_states) @auto_docstring class VideoLlama3PreTrainedModel(PreTrainedModel): config: VideoLlama3Config base_model_prefix = "model" input_modalities = ("image", "video", "text") supports_gradient_checkpointing = True _no_split_modules = ["VideoLlama3VisionEncoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_flash_attn = True _supports_sdpa = True _can_compile_fullgraph = True _supports_attention_backend = True def _init_weights(self, module): super()._init_weights(module) if isinstance(module, VideoLlama3VisionRotaryEmbedding): inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim)) init.copy_(module.inv_freq, inv_freq) class VideoLlama3VisionModel(VideoLlama3PreTrainedModel): config: VideoLlama3VisionConfig main_input_name = "pixel_values" input_modalities = ("image",) _can_record_outputs = { "hidden_states": VideoLlama3VisionEncoderLayer, "attentions": VideoLlama3VisionAttention, } def __init__(self, config: VideoLlama3VisionConfig): super().__init__(config) head_dim = config.hidden_size // config.num_attention_heads self.rotary_pos_emb = VideoLlama3VisionRotaryEmbedding(head_dim // 2) self.embeddings = VideoLlama3VisionEmbeddings(config) self.encoder = VideoLlama3VisionEncoder(config) self.post_layernorm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.post_init() def get_input_embeddings(self) -> VideoLlama3VisionEmbeddings: return self.embeddings.patch_embedding def pixel_unshuffle( self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, merge_sizes: torch.Tensor, ): hidden_states_chunks = hidden_states.split(grid_thw.prod(dim=1).tolist(), dim=0) outputs = [] for hidden_states, (t, h, w), merge_size in zip(hidden_states_chunks, grid_thw, merge_sizes): c = hidden_states.shape[-1] hidden_states = hidden_states.view(t, h // merge_size, w // merge_size, merge_size, merge_size, c).permute( 0, 1, 3, 2, 4, 5 ) hidden_states = hidden_states.reshape(t, h, w, c).permute(0, 3, 1, 2) hidden_states = torch.nn.functional.interpolate( hidden_states, size=(h // merge_size, w // merge_size), mode="bilinear" ) hidden_states = hidden_states.permute(0, 2, 3, 1).view(-1, c) outputs.append(hidden_states) return torch.cat(outputs, dim=0) @merge_with_config_defaults @capture_outputs(tie_last_hidden_states=False) @auto_docstring def forward( self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, merge_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutput: r""" grid_thw (`torch.LongTensor` of shape `(num_images_or_videos, 3)`): The temporal, height and width dimensions of feature shape for each image. Each row contains [t, h, w] values. merge_sizes (`torch.Tensor` of shape `(num_images_or_videos,)`): The spatial downsampling ratio of each image or video feature. """ position_embeddings = self.rotary_pos_emb(grid_thw, merge_sizes) cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( dim=0, # Select dtype based on the following factors: # - FA2 requires that cu_seqlens_q must have dtype int32 # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw # See https://github.com/huggingface/transformers/pull/34852 for more information dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, ) cu_seqlens = torch.nn.functional.pad(cu_seqlens, (1, 0), value=0) hidden_states = self.embeddings(pixel_values.type(self.dtype)) encoder_outputs: BaseModelOutput = self.encoder( hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, **kwargs, ) last_hidden_state = encoder_outputs[0] last_hidden_state = self.post_layernorm(last_hidden_state) last_hidden_state = self.pixel_unshuffle(last_hidden_state, grid_thw, merge_sizes) return BaseModelOutput(last_hidden_state=last_hidden_state) class VideoLlama3Projector(nn.Module): def __init__(self, config: VideoLlama3Config) -> None: super().__init__() in_hidden_size = config.vision_config.hidden_size out_hidden_size = config.text_config.hidden_size self.readout = nn.Sequential( nn.Linear(in_hidden_size, out_hidden_size), nn.GELU(), nn.Linear(out_hidden_size, out_hidden_size), ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.readout(hidden_states) return hidden_states @dataclass @auto_docstring( custom_intro=""" Base class for VideoLLaMA3 outputs, with hidden states and attentions. """ ) class VideoLlama3ModelOutputWithPast(ModelOutput): r""" past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. image_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(num_images_features, hidden_size)`. image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. video_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(num_video_features, hidden_size)`. video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. """ last_hidden_state: torch.FloatTensor = None past_key_values: list[torch.FloatTensor] | None = None hidden_states: tuple[torch.FloatTensor] | None = None attentions: tuple[torch.FloatTensor] | None = None image_hidden_states: torch.FloatTensor | None = None video_hidden_states: torch.FloatTensor | None = None @auto_docstring class VideoLlama3Model(VideoLlama3PreTrainedModel): base_model_prefix = "model" _checkpoint_conversion_mapping = {} # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False _can_compile_fullgraph = False def __init__(self, config: VideoLlama3Config): super().__init__(config) self.vision_model = AutoModel.from_config(config.vision_config) self.projector = VideoLlama3Projector(config) self.language_model = AutoModel.from_config(config.text_config) self.post_init() def get_input_embeddings(self): return self.language_model.get_input_embeddings() def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) @can_return_tuple @auto_docstring def get_video_features( self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor, video_merge_sizes: torch.LongTensor, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): The tensors corresponding to the input videos. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): The temporal, height and width of feature shape of each video in LLM. video_merge_sizes (`torch.Tensor` of shape `(num_videos,)`): The spatial downsampling ratio of each video feature. """ return self.get_image_features( pixel_values=pixel_values_videos, image_grid_thw=video_grid_thw, image_merge_sizes=video_merge_sizes, **kwargs, ) @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor, image_merge_sizes: torch.LongTensor, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): The tensors corresponding to the input images. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. image_merge_sizes (`torch.Tensor` of shape `(num_images,)`): The spatial downsampling ratio of each image feature. """ vision_outputs = self.vision_model( pixel_values=pixel_values, grid_thw=image_grid_thw, merge_sizes=image_merge_sizes, return_dict=True, **kwargs, ) last_hidden_state = vision_outputs.last_hidden_state image_embeds = self.projector(last_hidden_state) split_sizes = image_grid_thw.prod(dim=1) // (image_merge_sizes**2) image_embeds = torch.split(image_embeds, split_sizes.tolist()) vision_outputs.pooler_output = image_embeds return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor | None = None, video_features: torch.FloatTensor | None = None, ): """ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is equal to the length of multimodal features. If the lengths are different, an error is raised. """ if input_ids is None: special_image_mask = inputs_embeds == self.get_input_embeddings()( torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) ) special_image_mask = special_image_mask.all(-1) special_video_mask = inputs_embeds == self.get_input_embeddings()( torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device) ) special_video_mask = special_video_mask.all(-1) else: special_image_mask = input_ids == self.config.image_token_id special_video_mask = input_ids == self.config.video_token_id n_image_tokens = special_image_mask.sum() special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) if image_features is not None: torch_compilable_check( inputs_embeds[special_image_mask].numel() == image_features.numel(), f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}", ) n_video_tokens = special_video_mask.sum() special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) if video_features is not None: torch_compilable_check( inputs_embeds[special_video_mask].numel() == video_features.numel(), f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}", ) return special_image_mask, special_video_mask @can_return_tuple def forward( self, input_ids: torch.LongTensor = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, image_merge_sizes: torch.LongTensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, video_grid_thw: torch.LongTensor | None = None, video_merge_sizes: torch.LongTensor | None = None, video_compression_mask: torch.BoolTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | VideoLlama3ModelOutputWithPast: r""" image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. image_merge_sizes (`torch.Tensor` of shape `(num_images,)`): The spatial downsampling ratio of each image feature. video_grid_thw (`torch.Tensor` of shape `(num_videos, 3)`): The temporal, height and width of feature shape of each video before vision encoder. video_merge_sizes (`torch.Tensor` of shape `(num_videos,)`): The spatial downsampling ratio of each video feature. video_compression_mask (`torch.BoolTensor` of shape `(num_video_features,)`, *optional*): The mask to indicate which video features are kept after token compression. """ if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) image_embeds = None if pixel_values is not None: image_embeds = self.get_image_features( pixel_values, image_grid_thw, image_merge_sizes, return_dict=True ).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds ) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) video_embeds = None if pixel_values_videos is not None: video_embeds = self.get_video_features( pixel_values_videos, video_grid_thw, video_merge_sizes, return_dict=True ).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) if video_compression_mask is not None: video_embeds = video_embeds[video_compression_mask.to(video_embeds.device)] _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds ) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) outputs = self.language_model( input_ids=None, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, cache_position=cache_position, **kwargs, ) return VideoLlama3ModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=image_embeds, video_hidden_states=video_embeds, ) @dataclass @auto_docstring( custom_intro=""" Base class for VideoLLaMA3 causal language model (or autoregressive) outputs. """ ) class VideoLlama3CausalLMOutputWithPast(ModelOutput): r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. image_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(num_images_features, hidden_size)`. image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. video_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(num_video_features, hidden_size)`. video_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. """ loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None past_key_values: list[torch.FloatTensor] | None = None hidden_states: tuple[torch.FloatTensor] | None = None attentions: tuple[torch.FloatTensor] | None = None image_hidden_states: torch.FloatTensor | None = None video_hidden_states: torch.FloatTensor | None = None class VideoLlama3ForConditionalGeneration(VideoLlama3PreTrainedModel, GenerationMixin): _checkpoint_conversion_mapping = {} _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} _can_compile_fullgraph = False def __init__(self, config: VideoLlama3Config): super().__init__(config) self.model = VideoLlama3Model(config) self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) self.post_init() def get_input_embeddings(self): return self.model.get_input_embeddings() def set_input_embeddings(self, value): self.model.set_input_embeddings(value) @auto_docstring def get_video_features( self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): The tensors corresponding to the input videos. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): The temporal, height and width of feature shape of each video in LLM. """ return self.model.get_video_features( pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs ) @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): The tensors corresponding to the input images. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. """ return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, image_merge_sizes: torch.LongTensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, video_grid_thw: torch.LongTensor | None = None, video_merge_sizes: torch.LongTensor | None = None, video_compression_mask: torch.BoolTensor | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | VideoLlama3CausalLMOutputWithPast: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. image_merge_sizes (`torch.Tensor` of shape `(num_images,)`): The spatial downsampling ratio of each image feature. video_grid_thw (`torch.Tensor` of shape `(num_videos, 3)`): The temporal, height and width of feature shape of each video before vision encoder. video_merge_sizes (`torch.Tensor` of shape `(num_videos,)`): The spatial downsampling ratio of each video feature. video_compression_mask (`torch.BoolTensor` of shape `(num_video_features,)`, *optional*): The mask to indicate which video features are kept after token compression. """ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw, image_merge_sizes=image_merge_sizes, pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, video_merge_sizes=video_merge_sizes, video_compression_mask=video_compression_mask, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, return_dict=True, cache_position=cache_position, **kwargs, ) hidden_states = outputs[0] logits = self.lm_head(hidden_states) loss = None if labels is not None: loss = self.loss_function( logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs ) return VideoLlama3CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=outputs.image_hidden_states, video_hidden_states=outputs.video_hidden_states, ) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, position_ids=None, use_cache=True, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, image_merge_sizes: torch.LongTensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, video_grid_thw: torch.LongTensor | None = None, video_merge_sizes: torch.LongTensor | None = None, video_compression_mask: torch.BoolTensor | None = None, is_first_iteration: bool | None = False, **kwargs, ): # Overwritten -- in specific circumstances we don't want to forward image inputs to the model model_inputs = super().prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds, cache_position=cache_position, position_ids=position_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw, image_merge_sizes=image_merge_sizes, pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, video_merge_sizes=video_merge_sizes, video_compression_mask=video_compression_mask, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs, ) if not is_first_iteration and use_cache: model_inputs["pixel_values"] = None model_inputs["pixel_values_videos"] = None return model_inputs def _get_image_nums_and_video_nums( self, input_ids: torch.LongTensor | None, inputs_embeds: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, image_merge_sizes: torch.LongTensor | None = None, video_grid_thw: torch.LongTensor | None = None, video_merge_sizes: torch.LongTensor | None = None, video_compression_mask: torch.BoolTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Get the number of images and videos for each sample to calculate the separation length of the sample tensor. These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications. Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Returns: image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`) video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`) """ image_token_id = self.config.image_token_id video_token_id = self.config.video_token_id if inputs_embeds is not None: image_mask = ( inputs_embeds == self.get_input_embeddings()( torch.tensor(image_token_id, dtype=torch.long, device=inputs_embeds.device) ) )[..., 0] video_mask = ( inputs_embeds == self.get_input_embeddings()( torch.tensor(video_token_id, dtype=torch.long, device=inputs_embeds.device) ) )[..., 0] else: image_mask = input_ids == image_token_id video_mask = input_ids == video_token_id if image_grid_thw is not None: num_image_features = image_grid_thw.prod(dim=1) // (image_merge_sizes**2) else: num_image_features = [] if video_grid_thw is not None: num_video_features = video_grid_thw.prod(dim=1) // (video_merge_sizes**2) if video_compression_mask is not None: num_video_features = video_compression_mask.split(num_video_features.tolist()) num_video_features = [mask.sum() for mask in num_video_features] else: num_video_features = [] image_nums, video_nums = [], [] start_image_idx, start_video_idx = 0, 0 for num_image_tokens, num_video_tokens in zip(image_mask.sum(dim=1), video_mask.sum(dim=1)): cu_num_features = 0 image_idx = start_image_idx while image_idx < len(num_image_features) and cu_num_features < num_image_tokens: cu_num_features += num_image_features[image_idx] image_idx += 1 assert cu_num_features == num_image_tokens, ( "The number of image tokens does not match the number of image features." ) image_nums.append(image_idx - start_image_idx) start_image_idx = image_idx cu_num_features = 0 video_idx = start_video_idx while video_idx < len(num_video_features) and cu_num_features < num_video_tokens: cu_num_features += num_video_features[video_idx] video_idx += 1 assert cu_num_features == num_video_tokens, ( "The number of video tokens does not match the number of video features." ) video_nums.append(video_idx - start_video_idx) start_video_idx = video_idx return image_nums, video_nums def _expand_inputs_for_generation( self, expand_size: int = 1, is_encoder_decoder: bool = False, input_ids: torch.LongTensor | None = None, **model_kwargs, ) -> tuple[torch.LongTensor, dict[str, Any]]: # Overwritten -- Support for expanding tensors without a batch size dimension # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t # pixel_values.shape[0] is sum(seqlen_images for samples) # image_grid_thw.shape[0] is sum(num_images for samples) if expand_size == 1: return input_ids, model_kwargs visual_keys = [ "pixel_values", "image_grid_thw", "image_merge_sizes", "pixel_values_videos", "video_grid_thw", "video_merge_sizes", "video_compression_mask", ] def _repeat_interleave_samples(x, lengths, repeat_times): samples = torch.split(x, lengths) repeat_args = [repeat_times] + [1] * (x.dim() - 1) result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0) return result def _expand_dict_for_generation_visual(dict_to_expand): image_grid_thw = model_kwargs.get("image_grid_thw", None) video_grid_thw = model_kwargs.get("video_grid_thw", None) video_merge_sizes = model_kwargs.get("video_merge_sizes", None) video_compression_mask = model_kwargs.get("video_compression_mask", None) image_nums, video_nums = self._get_image_nums_and_video_nums( input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None), image_grid_thw=image_grid_thw, image_merge_sizes=model_kwargs.get("image_merge_sizes", None), video_grid_thw=video_grid_thw, video_merge_sizes=video_merge_sizes, video_compression_mask=video_compression_mask, ) for key in dict_to_expand: if key == "pixel_values": # split images into samples samples = torch.split(image_grid_thw, list(image_nums)) # compute the sequence length of images for each sample lengths = [torch.prod(sample, dim=1).sum() for sample in samples] dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "image_grid_thw": # get the num of images for each sample lengths = list(image_nums) dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "image_merge_sizes": lengths = list(image_nums) dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "pixel_values_videos": samples = torch.split(video_grid_thw, list(video_nums)) lengths = [torch.prod(sample, dim=1).sum() for sample in samples] dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "video_compression_mask": samples = torch.split(video_grid_thw, list(video_nums)) merge_sizes = torch.split(video_merge_sizes, list(video_nums)) lengths = [ (torch.prod(sample, dim=1) // merge_size**2).sum() for sample, merge_size in zip(samples, merge_sizes) ] dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "video_grid_thw": lengths = list(video_nums) dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "video_merge_sizes": lengths = list(video_nums) dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) return dict_to_expand def _expand_dict_for_generation(dict_to_expand): for key in dict_to_expand: if ( key != "cache_position" and dict_to_expand[key] is not None and isinstance(dict_to_expand[key], torch.Tensor) and key not in visual_keys ): dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) return dict_to_expand model_kwargs = _expand_dict_for_generation_visual(model_kwargs) if input_ids is not None: input_ids = input_ids.repeat_interleave(expand_size, dim=0) model_kwargs = _expand_dict_for_generation(model_kwargs) if is_encoder_decoder: if model_kwargs.get("encoder_outputs") is None: raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) return input_ids, model_kwargs __all__ = [ "VideoLlama3VisionModel", "VideoLlama3PreTrainedModel", "VideoLlama3Model", "VideoLlama3ForConditionalGeneration", ]