# coding=utf-8 # Copyright 2024 the HuggingFace Inc. 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. """PyTorch VideoMllama model.""" import os import math import time import inspect import warnings import importlib from packaging import version from typing import List, Optional, Tuple, Union, Callable, TYPE_CHECKING import queue from collections import deque import gc import torch from torch import nn import torch.utils.checkpoint import torch.nn.functional as F import torch.distributed as dist from transformers import PreTrainedModel from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.generation import GenerationMixin, GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput, GenerateBeamDecoderOnlyOutput, GenerateBeamEncoderDecoderOutput from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS from transformers.utils import ( add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings, is_torchdynamo_compiling, ) from transformers.utils.import_utils import _is_package_available from transformers.modeling_flash_attention_utils import _flash_attention_forward from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled from transformers.integrations.fsdp import is_fsdp_managed_module from transformers.generation.configuration_utils import ( GenerationConfig, GenerationMode, ) from transformers.generation.logits_process import LogitsProcessorList from transformers.generation.stopping_criteria import StoppingCriteriaList GenerateNonBeamOutput = Union[GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput] GenerateBeamOutput = Union[GenerateBeamDecoderOnlyOutput, GenerateBeamEncoderDecoderOutput] GenerateOutput = Union[GenerateNonBeamOutput, GenerateBeamOutput] if TYPE_CHECKING: from transformers.generation.streamers import BaseStreamer from .configuration_video_mllama import VideoMllamaConfig, VideoMllamaTextConfig, VideoMllamaVisionConfig from .processing_video_mllama import VIDEO_MLLAMA_PROCESSOR_PAD_POSITION_ID, VIDEO_MLLAMA_PROCESSOR_CROSS_ATTENTION_TOKEN_MASK_PAD_TOKEN_ID logger = logging.get_logger(__name__) def is_flash_attn_greater_or_equal_2_10(): if not _is_package_available("flash_attn"): return False return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.1.0") def _apply_rotary_pos_emb(states, cos, sin, position_ids=None, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) states_embed = (states * cos) + (rotate_half(states) * sin) return states_embed def _prepare_cross_attention_mask( cross_attention_mask: torch.Tensor, num_vision_tokens: int, dtype: str, ) -> Tuple[torch.Tensor, torch.Tensor]: # reshape so it can be used by attn module # shape = (batch_size, length, max_num_images, max_num_tiles) batch_size, text_total_length, *_ = cross_attention_mask.shape cross_attention_mask = cross_attention_mask.repeat_interleave(num_vision_tokens, dim=3) cross_attention_mask = cross_attention_mask.view(batch_size, text_total_length, -1) cross_attention_mask = cross_attention_mask.unsqueeze(1) # invert the mask inverted_cross_attn_mask = (1.0 - cross_attention_mask).to(dtype) cross_attention_mask = inverted_cross_attn_mask.masked_fill( inverted_cross_attn_mask.to(torch.bool), torch.finfo(dtype).min ) # apply full-row bias, which return 4D tensor of shape [B, H, S1, 1] where value is 0 if the a full row in cross attn mask's # last dimension contains negative infinity values, otherwise it's 1 negative_inf_value = torch.finfo(dtype).min full_text_row_masked_out_mask = ( (cross_attention_mask != negative_inf_value).any(dim=-1).type_as(cross_attention_mask)[..., None] ) cross_attention_mask *= full_text_row_masked_out_mask # cross_attention_mask.shape is (batch_size, 1, seq_len, num_concurrent_media * num_tiles * num_patches) # full_text_row_masked_out_mask.shape is (batch_size, 1, seq_len, 1) return cross_attention_mask, full_text_row_masked_out_mask def _prepare_aspect_ratio_attention_mask( aspect_ratio_mask: torch.Tensor, num_patches: int, target_length: int, dtype: torch.dtype, ) -> torch.Tensor: # Expand aspect ratio mask to target_length batch_size, max_num_tiles = aspect_ratio_mask.shape attention_mask = aspect_ratio_mask.view(batch_size, max_num_tiles, 1, 1).to(dtype) attention_mask = attention_mask.repeat(1, 1, target_length, 1) # Mask padding patches pad_patches = target_length - num_patches attention_mask[:, :, -pad_patches:] = 0 # Invert the mask (0 -> 1, 1 -> 0) attention_mask = 1 - attention_mask # Reshape to 2D and create 4D attention mask # (batch_size, 1, max_num_tiles * target_length, max_num_tiles * target_length) attention_mask = attention_mask.reshape(batch_size, max_num_tiles * target_length, 1) attention_mask = attention_mask @ attention_mask.transpose(-1, -2) * torch.finfo(dtype).min attention_mask = attention_mask.unsqueeze(1) return attention_mask class VideoMllamaPrecomputedAspectRatioEmbedding(nn.Module): def __init__(self, config: VideoMllamaVisionConfig, is_gated: bool = True): super().__init__() self.max_num_tiles = config.max_num_tiles self.hidden_size = config.hidden_size self.max_aspect_ratio_id = config.max_aspect_ratio_id self.is_gated = is_gated self.embedding = nn.Embedding(self.max_aspect_ratio_id + 1, self.max_num_tiles * self.hidden_size) if is_gated: self.gate = nn.Parameter(torch.zeros(1)) def forward(self, hidden_state: torch.Tensor, aspect_ratio_ids: torch.Tensor) -> torch.Tensor: """ Args: hidden_state: (batch_size * num_concurrent_media, num_tiles(video_num_tiles), num_patches, hidden_size) aspect_ratio_ids: (batch_size * num_concurrent_media, 1) """ embeddings = self.embedding(aspect_ratio_ids) embeddings = embeddings.reshape(-1, self.max_num_tiles, 1, self.hidden_size) num_tiles = hidden_state.shape[1] embeddings = embeddings[:, :num_tiles, :, :] if self.is_gated: embeddings = embeddings * self.gate.tanh() hidden_state = hidden_state + embeddings return hidden_state class VideoMllamaPrecomputedPositionEmbedding(nn.Module): def __init__(self, config: VideoMllamaVisionConfig): super().__init__() self.max_num_tiles = config.max_num_tiles self.max_aspect_ratio_id = config.max_aspect_ratio_id self.num_patches = (config.image_size // config.patch_size) ** 2 + 1 self.hidden_size = config.hidden_size self.scale = config.hidden_size**-0.5 self.gate = nn.Parameter(torch.zeros(1)) # position embedding position_embedding = torch.randn(self.num_patches, self.hidden_size) self.embedding = nn.Parameter(self.scale * position_embedding) # tile position embedding self.tile_embedding = nn.Embedding( self.max_aspect_ratio_id + 1, self.max_num_tiles * self.num_patches * self.hidden_size ) def forward(self, hidden_state: torch.Tensor, aspect_ratio_ids: torch.Tensor) -> torch.Tensor: """ Args: hidden_state: hidden_state.shape is (batch_size * num_concurrent_media, num_tiles(video_num_tiles), num_patches, hidden_size) aspect_ratio_ids: (batch_size * num_concurrent_media, 1) """ # position embeddings gated_position_embedding = (1 - self.gate.tanh()) * self.embedding hidden_state = hidden_state + gated_position_embedding.view(1, 1, self.num_patches, self.hidden_size) # precomputed tile position embeddings tile_position_embedding = self.tile_embedding(aspect_ratio_ids) batch_size = hidden_state.shape[0] tile_position_embedding = tile_position_embedding.reshape( batch_size, self.max_num_tiles, self.num_patches, self.hidden_size ) num_tiles = hidden_state.shape[1] tile_position_embedding = tile_position_embedding[:, :num_tiles, :, :] gated_tile_position_embedding = self.gate.tanh() * tile_position_embedding hidden_state = hidden_state + gated_tile_position_embedding return hidden_state # Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->VideoMllamaVision class VideoMllamaVisionMLP(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 class VideoMllamaVisionAttention(nn.Module): def __init__(self, config: VideoMllamaVisionConfig): super().__init__() self.embed_dim = config.hidden_size self.num_heads = config.attention_heads self.head_dim = config.hidden_size // config.attention_heads self.q_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.embed_dim, bias=False) def forward( self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = None, ) -> torch.Tensor: query = self.q_proj(hidden_state) key = self.k_proj(hidden_state) value = self.v_proj(hidden_state) batch_size, q_seq_len, _ = query.shape _, kv_seq_len, _ = key.shape query = query.view(batch_size, q_seq_len, self.num_heads, self.head_dim).transpose(1, 2) key = key.view(batch_size, kv_seq_len, self.num_heads, self.head_dim).transpose(1, 2) value = value.view(batch_size, kv_seq_len, self.num_heads, self.head_dim).transpose(1, 2) attn_weights = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(self.head_dim) if attention_mask is not None: # no matter the length, we just slice it causal_mask = attention_mask[:, :, :, : key.shape[-2]] attn_weights = attn_weights + causal_mask # upcast attention to fp32 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_output = torch.matmul(attn_weights, value) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(batch_size, q_seq_len, -1) output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return output, attn_weights class VideoMllamaVisionFlashAttention2(VideoMllamaVisionAttention): # def __init__(self, config: MllamaVisionConfig): def __init__(self, config): super().__init__(config) self.config = config self._softmax_scale = 1 / math.sqrt(self.head_dim) self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() def forward( self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = None, **kwargs ) -> torch.Tensor: if attention_mask is not None: raise NotImplementedError("The Flash-attention-2 implementation of MllamaVisionModel can only be used when attention_mask is None.") output_attentions = False query = self.q_proj(hidden_state) key = self.k_proj(hidden_state) value = self.v_proj(hidden_state) batch_size, q_seq_len, _ = query.shape _, kv_seq_len, _ = key.shape query = query.view(batch_size, q_seq_len, self.num_heads, self.head_dim).transpose(1, 2) key = key.view(batch_size, kv_seq_len, self.num_heads, self.head_dim).transpose(1, 2) value = value.view(batch_size, kv_seq_len, self.num_heads, self.head_dim).transpose(1, 2) query_states = query.transpose(1, 2) key_states = key.transpose(1, 2) value_states = value.transpose(1, 2) is_causal = False input_dtype = query_states.dtype if input_dtype == torch.float32: if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_gpu_dtype() # Handle the case where the model is quantized elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: target_dtype = self.q_proj.weight.dtype logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" f" {target_dtype}." ) query_states = query_states.to(target_dtype) key_states = key_states.to(target_dtype) value_states = value_states.to(target_dtype) attn_output = _flash_attention_forward( query_states, key_states, value_states, attention_mask, q_seq_len, sliding_window=getattr(self, "sliding_window", None), use_top_left_mask=self._flash_attn_uses_top_left_mask, is_causal=is_causal, ) # attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(batch_size, q_seq_len, -1).contiguous() output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return output, attn_weights class VideoMllamaVisionSdpaAttention(VideoMllamaVisionAttention): # Adapted from VideoMllamaVisionAttention def forward( self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = None, ) -> torch.Tensor: # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. if output_attentions: logger.warning_once( "VideoMllamaModel is using VideoMllamaVisionSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' ) return super().forward( hidden_state=hidden_state, attention_mask=attention_mask, output_attentions=output_attentions, ) query = self.q_proj(hidden_state) key = self.k_proj(hidden_state) value = self.v_proj(hidden_state) batch_size, q_seq_len, _ = query.shape _, kv_seq_len, _ = key.shape query = query.view(batch_size, q_seq_len, self.num_heads, self.head_dim) key = key.view(batch_size, kv_seq_len, self.num_heads, self.head_dim) value = value.view(batch_size, kv_seq_len, self.num_heads, self.head_dim) query = query.transpose(1, 2) key = key.transpose(1, 2) value = value.transpose(1, 2) attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(batch_size, q_seq_len, -1) output = self.o_proj(attn_output) return output, None VideoMllama_VISION_ATTENTION_CLASSES = {"eager": VideoMllamaVisionAttention, "sdpa": VideoMllamaVisionSdpaAttention, "flash_attention_2": VideoMllamaVisionFlashAttention2} class VideoMllamaVisionEncoderLayer(nn.Module): def __init__(self, config: VideoMllamaVisionConfig, is_gated: bool = False): super().__init__() self.hidden_size = config.hidden_size self.num_attention_heads = config.attention_heads self.is_gated = is_gated self.intermediate_size = config.intermediate_size self.self_attn = VideoMllama_VISION_ATTENTION_CLASSES[config._attn_implementation](config) self.mlp = VideoMllamaVisionMLP(config) self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps) self.post_attention_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps) if is_gated: self.gate_attn = nn.Parameter(torch.ones(1) * math.pi / 4) self.gate_ffn = nn.Parameter(torch.ones(1) * math.pi / 4) def forward( self, hidden_state: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = None, ): # Self Attention residual = hidden_state hidden_state = self.input_layernorm(hidden_state) hidden_state, attn_weights = self.self_attn(hidden_state, attention_mask=attention_mask) if self.is_gated: hidden_state = self.gate_attn.tanh() * hidden_state hidden_state = residual + hidden_state # Feed forward residual = hidden_state hidden_state = self.post_attention_layernorm(hidden_state) hidden_state = self.mlp(hidden_state) if self.is_gated: hidden_state = self.gate_ffn.tanh() * hidden_state hidden_state = residual + hidden_state outputs = (hidden_state,) if output_attentions: outputs += (attn_weights,) return outputs class VideoMllamaVisionEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`VideoMllamaEncoderLayer`]. Args: config: VideoMllamaConfig """ def __init__(self, config: VideoMllamaVisionConfig, num_layers=32, is_gated=False): super().__init__() self.config = config self.layers = nn.ModuleList([VideoMllamaVisionEncoderLayer(config, is_gated) for _ in range(num_layers)]) self.gradient_checkpointing = False self.config = config def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, BaseModelOutput]: r""" Args: inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None for encoder_layer in self.layers: if output_hidden_states: encoder_states = encoder_states + (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( encoder_layer.__call__, hidden_states, attention_mask, output_attentions, ) else: layer_outputs = encoder_layer( hidden_state=hidden_states, attention_mask=attention_mask, output_attentions=output_attentions, ) if output_attentions: all_attentions = all_attentions + (layer_outputs[1],) hidden_states = layer_outputs[0] if output_hidden_states: encoder_states = encoder_states + (hidden_states,) if not return_dict: return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) return BaseModelOutput( last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions ) # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->VideoMllamaText class VideoMllamaTextRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): """ VideoMllamaTextRMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class VideoMllamaTextCrossAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__( self, config: Optional[VideoMllamaTextConfig] = None, layer_idx: Optional[int] = None, ): super().__init__() self.config = config self.num_heads = self.config.num_attention_heads self.num_key_value_heads = self.config.num_key_value_heads self.dropout = config.dropout self.hidden_size = config.hidden_size self.head_dim = config.hidden_size // self.num_heads self.layer_idx = layer_idx self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) self.q_norm = VideoMllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = VideoMllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, cross_attention_states: Optional[torch.Tensor] = None, past_key_value: Optional[Cache] = None, attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, use_cache: bool = None, position_ids: Optional[torch.LongTensor] = None, # vision_position_ids cache_position: Optional[torch.LongTensor] = None, # vision_cache_position position_embeddings: Optional[torch.Tensor] = None, # vision_position_embeddings query_position_embeddings: Optional[torch.Tensor] = None, # position_embeddings ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """Input shape: Batch x Time x Channel""" bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) query_states = self.q_norm(query_states) if cross_attention_states is not None: key_states = self.k_proj(cross_attention_states) value_states = self.v_proj(cross_attention_states) key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) # add rope embedding for query and key # Note that position embedding is not the same for query and key cos, sin = query_position_embeddings query_states= _apply_rotary_pos_emb(query_states, cos, sin) vision_cos, vision_sin = position_embeddings # vision_cos.shape is (batch_size, num_concurrent_media * num_tiles * num_patches, head_dim) # key_states.shape is (batch_size, num_key_value_heads, num_concurrent_media * num_tiles * num_patches, head_dim) key_states = _apply_rotary_pos_emb(key_states, vision_cos, vision_sin) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) key_states = self.k_norm(key_states) if past_key_value is not None: # if we have a new image + new tokens, we only computed key_states on that new image # we still update the cross key states, past_image, new_image. And use it! head_num = key_states.shape[1] num_concurrent_media = len(cache_position) key_states = key_states.view(bsz, head_num, num_concurrent_media, -1, self.head_dim).transpose(2, 3) # key_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) key_states = key_states.reshape(bsz, -1, num_concurrent_media, self.head_dim) # key_states.shape is (batch_size, head_num * num_tiles * num_patches, num_concurrent_media, head_dim) value_states = value_states.view(bsz, head_num, num_concurrent_media, -1, self.head_dim).transpose(2, 3) # value_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) value_states = value_states.reshape(bsz, -1, num_concurrent_media, self.head_dim) # value_states.shape is (batch_size, head_num * num_tiles * num_patches, num_concurrent_media, head_dim) key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, {"cache_position": cache_position} ) # Restore the shape of key_states and value_states num_concurrent_media = key_states.shape[-2] key_states = key_states.reshape(bsz, head_num, -1, num_concurrent_media, self.head_dim) # key_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) key_states = key_states.transpose(2, 3).reshape(bsz, head_num, -1, self.head_dim) value_states = value_states.reshape(bsz, head_num, -1, num_concurrent_media, self.head_dim) # value_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) value_states = value_states.transpose(2, 3).reshape(bsz, head_num, -1, self.head_dim) # key_states.shape is (batch_size, head_num, num_concurrent_media * num_tiles * num_patches, head_dim) elif past_key_value is not None: key_states, value_states = ( past_key_value.key_cache[self.layer_idx], past_key_value.value_cache[self.layer_idx], ) # Restore the shape of key_states and value_states num_concurrent_media = key_states.shape[-2] head_num = self.num_heads # num_key_value_heads * num_key_value_groups key_states = key_states.reshape(bsz, head_num, -1, num_concurrent_media, self.head_dim) # key_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) key_states = key_states.transpose(2, 3).reshape(bsz, head_num, -1, self.head_dim) value_states = value_states.reshape(bsz, head_num, -1, num_concurrent_media, self.head_dim) # value_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) value_states = value_states.transpose(2, 3).reshape(bsz, head_num, -1, self.head_dim) # key_states.shape is (batch_size, head_num, num_concurrent_media * num_tiles * num_patches, head_dim) else: raise ValueError( "Cross attention layer can't find neither `cross_attn_states` nor cached values for key/values!" ) # key_states = repeat_kv(key_states, self.num_key_value_groups) # value_states = repeat_kv(value_states, self.num_key_value_groups) # key_states = self.k_norm(key_states) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attention_mask is not None: # no matter the length, we just slice it causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class VideoMllamaTextCrossSdpaAttention(VideoMllamaTextCrossAttention): """ VideoMllama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from `VideoMllamaTextCrossAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to SDPA API. """ def forward( self, hidden_states: torch.Tensor, cross_attention_states: Optional[torch.Tensor] = None, past_key_value: Optional[Cache] = None, attention_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, use_cache: bool = None, position_ids: Optional[torch.LongTensor] = None, # vision_position_ids cache_position: Optional[torch.LongTensor] = None, # vision_cache_position position_embeddings: Optional[torch.Tensor] = None, # vision_position_embeddings query_position_embeddings: Optional[torch.Tensor] = None, # position_embeddings ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """Input shape: Batch x Time x Channel""" if output_attentions: # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. logger.warning_once( "MllamaModel is using MllamaTextCrossSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' ) return super().forward( hidden_states=hidden_states, cross_attention_states=cross_attention_states, attention_mask=attention_mask, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, ) bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) query_states = self.q_norm(query_states) if cross_attention_states is not None: key_states = self.k_proj(cross_attention_states) value_states = self.v_proj(cross_attention_states) key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2) # add rope embedding for query and key # Note that position embedding is not the same for query and key cos, sin = query_position_embeddings query_states= _apply_rotary_pos_emb(query_states, cos, sin) vision_cos, vision_sin = position_embeddings key_states = _apply_rotary_pos_emb(key_states, vision_cos, vision_sin) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) # key_states.shape is (batch_size, head_num, num_concurrent_media * num_tiles * num_patches, head_dim) key_states = self.k_norm(key_states) if past_key_value is not None: # if we have a new image + new tokens, we only computed key_states on that new image # we still update the cross key states, past_image, new_image. And use it! head_num = key_states.shape[1] num_concurrent_media = len(cache_position) key_states = key_states.view(bsz, head_num, num_concurrent_media, -1, self.head_dim).transpose(2, 3) # key_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) key_states = key_states.reshape(bsz, -1, num_concurrent_media, self.head_dim) # key_states.shape is (batch_size, head_num * num_tiles * num_patches, num_concurrent_media, head_dim) value_states = value_states.view(bsz, head_num, num_concurrent_media, -1, self.head_dim).transpose(2, 3) # value_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) value_states = value_states.reshape(bsz, -1, num_concurrent_media, self.head_dim) # value_states.shape is (batch_size, head_num * num_tiles * num_patches, num_concurrent_media, head_dim) key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, {"cache_position": cache_position} ) # Restore the shape of key_states and value_states num_concurrent_media = key_states.shape[-2] key_states = key_states.reshape(bsz, head_num, -1, num_concurrent_media, self.head_dim) # key_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) key_states = key_states.transpose(2, 3).reshape(bsz, head_num, -1, self.head_dim) value_states = value_states.reshape(bsz, head_num, -1, num_concurrent_media, self.head_dim) # value_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) value_states = value_states.transpose(2, 3).reshape(bsz, head_num, -1, self.head_dim) # key_states.shape is (batch_size, head_num, num_concurrent_media * num_tiles * num_patches, head_dim) elif past_key_value is not None: key_states, value_states = ( past_key_value.key_cache[self.layer_idx], past_key_value.value_cache[self.layer_idx], ) # Restore the shape of key_states and value_states num_concurrent_media = key_states.shape[-2] head_num = self.num_heads # num_key_value_heads * num_key_value_groups key_states = key_states.reshape(bsz, head_num, -1, num_concurrent_media, self.head_dim) # key_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) key_states = key_states.transpose(2, 3).reshape(bsz, head_num, -1, self.head_dim) value_states = value_states.reshape(bsz, head_num, -1, num_concurrent_media, self.head_dim) # value_states.shape is (batch_size, head_num, num_tiles * num_patches, num_concurrent_media, head_dim) value_states = value_states.transpose(2, 3).reshape(bsz, head_num, -1, self.head_dim) # key_states.shape is (batch_size, head_num, num_concurrent_media * num_tiles * num_patches, head_dim) else: raise ValueError( "Cross attention layer can't find neither `cross_attn_states` nor cached values for key/values!" ) # key_states = repeat_kv(key_states, self.num_key_value_groups) # value_states = repeat_kv(value_states, self.num_key_value_groups) # key_states = self.k_norm(key_states) # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, # Reference: https://github.com/pytorch/pytorch/issues/112577. if query_states.device.type == "cuda" and attention_mask is not None: query_states = query_states.contiguous() key_states = key_states.contiguous() value_states = value_states.contiguous() # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. is_causal = True if attention_mask is None and q_len > 1 else False attn_output = torch.nn.functional.scaled_dot_product_attention( query_states, key_states, value_states, attn_mask=attention_mask, dropout_p=self.dropout if self.training else 0.0, is_causal=is_causal, ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.o_proj(attn_output) return attn_output, None, past_key_value # Copied from transformers.models.llama.modeling_llama.rotate_half 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) # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. position_ids (`torch.Tensor`, *optional*): Deprecated and unused. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed # Copied from transformers.models.llama.modeling_llama.repeat_kv 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) class VideoMllamaTextSelfAttention(nn.Module): def __init__(self, config: VideoMllamaTextConfig, layer_idx: int): super().__init__() self.config = config self.num_heads = config.num_attention_heads self.dropout = config.dropout self.hidden_size = config.hidden_size self.num_key_value_heads = config.num_key_value_heads self.head_dim = config.hidden_size // self.num_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.rope_theta = config.rope_theta self.layer_idx = layer_idx self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_embeddings: torch.Tensor, output_attentions: bool = False, use_cache: bool = False, past_key_value=None, cache_position=None, **kwargs, ): bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attention_mask is not None: # no matter the length, we just slice it causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask # upcast attention to fp32 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, -1) attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class VideoMllamaTextFlashAttention2(VideoMllamaTextSelfAttention): def __init__(self, config: VideoMllamaTextConfig,layer_idx:int): super().__init__(config,layer_idx) self._softmax_scale = 1 / math.sqrt(self.head_dim) self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_embeddings: torch.Tensor, output_attentions: bool = False, use_cache: bool = False, past_key_value=None, cache_position=None, **kwargs ): if isinstance(past_key_value, StaticCache): raise ValueError( "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" ) output_attentions = False bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) ## LlamaFlashAttention2.forward() also not do this # key_states = repeat_kv(key_states, self.num_key_value_groups) # value_states = repeat_kv(value_states, self.num_key_value_groups) # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache # to be able to avoid many of these transpose/reshape/view. query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) is_causal = True if q_len > 1 else False dropout_rate = self.dropout if self.training else 0.0 input_dtype = query_states.dtype if input_dtype == torch.float32: if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_gpu_dtype() # Handle the case where the model is quantized elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: target_dtype = self.q_proj.weight.dtype logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" f" {target_dtype}." ) query_states = query_states.to(target_dtype) key_states = key_states.to(target_dtype) value_states = value_states.to(target_dtype) attn_output = _flash_attention_forward( query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate, sliding_window=getattr(self, "sliding_window", None), use_top_left_mask=self._flash_attn_uses_top_left_mask, is_causal=is_causal, **kwargs, ) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class VideoMllamaTextSelfSdpaAttention(VideoMllamaTextSelfAttention): # Adapted from VideoMllamaTextSelfAttention def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, position_embeddings: torch.Tensor, output_attentions: bool = False, use_cache: bool = False, past_key_value=None, cache_position=None, **kwargs, ): if output_attentions: # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. logger.warning_once( "VideoMllamaModel is using VideoMllamaTextSelfSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' ) return super().forward( hidden_states=hidden_states, attention_mask=attention_mask, position_embeddings=position_embeddings, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, **kwargs, ) bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) causal_mask = attention_mask if attention_mask is not None: causal_mask = causal_mask[:, :, :, : key_states.shape[-2]] # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, # Reference: https://github.com/pytorch/pytorch/issues/112577. if query_states.device.type == "cuda" and causal_mask is not None: query_states = query_states.contiguous() key_states = key_states.contiguous() value_states = value_states.contiguous() # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. is_causal = True if causal_mask is None and q_len > 1 else False attn_output = torch.nn.functional.scaled_dot_product_attention( query_states, key_states, value_states, attn_mask=causal_mask, dropout_p=self.dropout if self.training else 0.0, is_causal=is_causal, ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, -1) attn_output = self.o_proj(attn_output) return attn_output, None, past_key_value VideoMllama_TEXT_CROSS_ATTENTION_CLASSES = {"eager": VideoMllamaTextCrossAttention, "sdpa": VideoMllamaTextCrossSdpaAttention, "flash_attention_2": VideoMllamaTextCrossSdpaAttention} VideoMllama_TEXT_ATTENTION_CLASSES = {"eager": VideoMllamaTextSelfAttention, "sdpa": VideoMllamaTextSelfSdpaAttention, "flash_attention_2": VideoMllamaTextFlashAttention2} # Copied from transformers.models.gemma2.modeling_gemma2.Gemma2MLP with Gemma2->VideoMllamaText class VideoMllamaTextMLP(nn.Module): def __init__(self, config): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) # Ignore copy self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) # Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer class VideoMllamaSelfAttentionDecoderLayer(nn.Module): def __init__(self, config: VideoMllamaTextConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = VideoMllama_TEXT_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) self.mlp = VideoMllamaTextMLP(config) self.input_layernorm = VideoMllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = VideoMllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.layer_idx = layer_idx def forward( self, hidden_states: torch.Tensor, cross_attention_states: torch.Tensor, cross_attention_mask: torch.Tensor, attention_mask: torch.Tensor, full_text_row_masked_out_mask: Tuple[torch.Tensor, torch.Tensor], position_ids: Optional[torch.LongTensor] = None, vision_position_ids: Optional[torch.LongTensor] = None, # vision_position_ids past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, vision_cache_position: Optional[torch.LongTensor] = None, # vision_cache_position position_embeddings: Optional[torch.Tensor] = None, vision_position_embeddings: Optional[torch.Tensor] = None, # vision_position_embeddings ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, query_sequence_length, key_sequence_length)` if default attention is used. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, with `head_dim` being the embedding dimension of each attention head. kwargs (`dict`, *optional*): Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code into the model """ residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) return outputs class VideoMllamaCrossAttentionDecoderLayer(torch.nn.Module): """Cross-attention transformer block with tanh-gated attention and feedforward.""" def __init__(self, config: VideoMllamaTextConfig, layer_idx: int) -> None: super().__init__() self.layer_idx = layer_idx self.cross_attn = VideoMllama_TEXT_CROSS_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx) self.input_layernorm = VideoMllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.cross_attn_attn_gate = torch.nn.Parameter(torch.zeros(1)) self.mlp = VideoMllamaTextMLP(config) self.post_attention_layernorm = VideoMllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.cross_attn_mlp_gate = torch.nn.Parameter(torch.zeros(1)) def forward( self, hidden_states: torch.Tensor, cross_attention_states: torch.Tensor, cross_attention_mask: torch.Tensor, attention_mask: torch.Tensor, full_text_row_masked_out_mask: Tuple[torch.Tensor, torch.Tensor], position_ids: Optional[torch.LongTensor] = None, vision_position_ids: Optional[torch.LongTensor] = None, # vision_position_ids past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, vision_cache_position: Optional[torch.LongTensor] = None, # vision_cache_position position_embeddings: Optional[torch.Tensor] = None, vision_position_embeddings: Optional[torch.Tensor] = None, # vision_position_embeddings ) -> Tuple[torch.Tensor]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, attn_weights, past_key_value = self.cross_attn( hidden_states=hidden_states, attention_mask=cross_attention_mask, cross_attention_states=cross_attention_states, past_key_value=past_key_value, output_attentions=output_attentions, position_ids=vision_position_ids, cache_position=vision_cache_position, position_embeddings=vision_position_embeddings, query_position_embeddings=position_embeddings, ) hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) if full_text_row_masked_out_mask is not None: hidden_states = full_text_row_masked_out_mask[:, 0] * hidden_states # type: ignore hidden_states = residual + self.cross_attn_mlp_gate.tanh() * hidden_states outputs = (hidden_states,) if output_attentions: outputs += (attn_weights,) if use_cache: outputs += (past_key_value,) return outputs class VideoMllamaRotaryEmbedding(nn.Module): def __init__(self, config: VideoMllamaTextConfig, device=None): super().__init__() self.rope_type = config.rope_scaling["rope_type"] self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq def _dynamic_frequency_update(self, position_ids, device): """ dynamic RoPE layers should recompute `inv_freq` in the following situations: 1 - growing beyond the cached sequence length (allow scaling) 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) """ seq_len = torch.max(position_ids) + 1 if seq_len > self.max_seq_len_cached: # growth inv_freq, self.attention_scaling = self.rope_init_fn( self.config, device, seq_len=seq_len, **self.rope_kwargs ) self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation self.max_seq_len_cached = seq_len if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) self.max_seq_len_cached = self.original_max_seq_len @torch.no_grad() def forward(self, x, position_ids): if "dynamic" in self.rope_type: self._dynamic_frequency_update(position_ids, device=x.device) # Core RoPE block inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() # Force float32 (see https://github.com/huggingface/transformers/pull/29285) device_type = x.device.type device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" with torch.autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention cos = cos * self.attention_scaling sin = sin * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) class VideoMllamaPreTrainedModel(PreTrainedModel): config_class = VideoMllamaConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = [ "VideoMllamaVisionEncoderLayer", "VideoMllamaCrossAttentionDecoderLayer", "VideoMllamaSelfAttentionDecoderLayer", ] _supports_cache_class = True _supports_static_cache = False # static cache cannot have different shapes for each layer _supports_sdpa = True _supports_flash_attn_2 = True _supports_quantized_cache = True def _init_weights(self, module): std = self.config.get_text_config().initializer_range if isinstance(module, (nn.Linear, nn.Conv2d)): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.Parameter): module.data.normal_(mean=0.0, std=std) elif isinstance(module, VideoMllamaVisionModel): nn.init.normal_(module.class_embedding.data, std=std) elif isinstance(module, VideoMllamaPrecomputedPositionEmbedding): nn.init.normal_(module.embedding.data, std=std) elif isinstance(module, VideoMllamaVisionEncoderLayer) and module.is_gated: nn.init.normal_(module.gate_attn.data, std=std) nn.init.normal_(module.gate_ffn.data, std=std) # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask def _update_causal_mask( self, attention_mask: torch.Tensor, input_tensor: torch.Tensor, cache_position: torch.Tensor, past_key_values: Cache, output_attentions: bool, ): if self.config._attn_implementation == "flash_attention_2": if attention_mask is not None and 0.0 in attention_mask: return attention_mask return None # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail # to infer the attention mask. past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 using_static_cache = isinstance(past_key_values, StaticCache) # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: if AttentionMaskConverter._ignore_causal_mask_sdpa( attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens, is_training=self.training, ): return None dtype, device = input_tensor.dtype, input_tensor.device sequence_length = input_tensor.shape[1] if using_static_cache: target_length = past_key_values.get_max_cache_shape() else: target_length = ( attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else past_seen_tokens + sequence_length + 1 ) # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( attention_mask, sequence_length=sequence_length, target_length=target_length, dtype=dtype, device=device, cache_position=cache_position, batch_size=input_tensor.shape[0], ) if ( self.config._attn_implementation == "sdpa" and attention_mask is not None and attention_mask.device.type == "cuda" and not output_attentions ): # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. # Details: https://github.com/pytorch/pytorch/issues/110213 min_dtype = torch.finfo(dtype).min causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) return causal_mask @staticmethod # Copied from transformers.models.llama.modeling_llama.LlamaModel._prepare_4d_causal_attention_mask_with_cache_position def _prepare_4d_causal_attention_mask_with_cache_position( attention_mask: torch.Tensor, sequence_length: int, target_length: int, dtype: torch.dtype, device: torch.device, cache_position: torch.Tensor, batch_size: int, **kwargs, ): """ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. Args: attention_mask (`torch.Tensor`): A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. sequence_length (`int`): The sequence length being processed. target_length (`int`): The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. dtype (`torch.dtype`): The dtype to use for the 4D attention mask. device (`torch.device`): The device to plcae the 4D attention mask on. cache_position (`torch.Tensor`): Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): Batch size. """ if attention_mask is not None and attention_mask.dim() == 4: # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. causal_mask = attention_mask else: min_dtype = torch.finfo(dtype).min causal_mask = torch.full( (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device ) if sequence_length != 1: causal_mask = torch.triu(causal_mask, diagonal=1) causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) if attention_mask is not None: causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit mask_length = attention_mask.shape[-1] padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] padding_mask = padding_mask == 0 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( padding_mask, min_dtype ) return causal_mask VideoMllama_START_DOCSTRING = r""" This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`VideoMllamaConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. """ VideoMllama_VISION_INPUTS_DOCSTRING = r""" Args: pixel_values (`torch.FloatTensor` of shape `(batch_size, max_num_images, max_num_tiles, channels, image_size, image_size)): The tensors corresponding to the input images. Pixel values can be obtained using [`AutoImageProcessor`]. See [`VideoMllamaImageProcessor.__call__`] for details ([]`VideoMllamaProcessor`] uses [`VideoMllamaImageProcessor`] for processing images). aspect_ratio_mask (`torch.Tensor` of shape `(batch_size, max_num_images, max_num_tiles)`, *optional*): Mask to avoid performing attention on padding tiles. Mask values selected in `[0, 1]`: - 1 for tiles that are **not masked**, - 0 for tiles that are **masked**. aspect_ratio_ids (`torch.Tensor` of shape `(batch_size, max_num_images)`, *optional*): Aspect ratio ids used to select the appropriate precomputed tile embeddings based on the aspect ratio of each input image. These ids correspond to indices in the model's list of supported aspect ratios, offset by 1. For example, if the model supports aspect ratios [[1, 1], [1, 2], [2, 1]]: - An image with aspect ratio [1, 1] would have ID 1 - An image with aspect ratio [1, 2] would have ID 2 - An image with aspect ratio [2, 1] would have ID 3 The id 0 is reserved for padding (i.e., no image). If an image has aspect ratio [1, 2], that means it was split into 2 tiles horizontally, and its `aspect_ratio_id` would be 2. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. """ VideoMllama_TEXT_INPUTS_DOCSTRING = r""" Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. If `past_key_values` is used, optionally only the last `input_ids` have to be input (see `past_key_values`). If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): Cross-attention mask to control the interaction between text tokens and image tiles. This 4D tensor defines which image tiles each text token should attend to. For each text token (in seq_length): - 1 indicates the token **should attend** to the corresponding image tile - 0 indicates the token **should not attend** to the corresponding image tile cross_attention_states (`torch.FloatTensor`, *optional*): Output of the vision model, used for cross-attention. This tensor contains the processed image features that the language model will attend to. position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. Two formats are allowed: - a [`~cache_utils.Cache`] instance, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache); - 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)`). This is also known as the legacy cache format. The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the legacy cache format will be returned. If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, this tensor is not affected by padding. It is used to update the cache in the correct position and to infer the complete sequence length. """ VIDEOMLLAMA_INPUTS_DOCSTRING = r""" Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) pixel_values (`torch.FloatTensor` of shape `(batch_size, max_num_images, max_num_tiles, channels, image_size, image_size)): The tensors corresponding to the input images. Pixel values can be obtained using [`AutoImageProcessor`]. See [`VideoMllamaImageProcessor.__call__`] for details ([]`VideoMllamaProcessor`] uses [`VideoMllamaImageProcessor`] for processing images). aspect_ratio_mask (`torch.Tensor` of shape `(batch_size, max_num_images, max_num_tiles)`, *optional*): Mask to avoid performing attention on padding tiles. Mask values selected in `[0, 1]`: - 1 for tiles that are **not masked**, - 0 for tiles that are **masked**. aspect_ratio_ids (`torch.Tensor` of shape `(batch_size, max_num_images)`, *optional*): Aspect ratio ids used to select the appropriate precomputed tile embeddings based on the aspect ratio of each input image. These ids correspond to indices in the model's list of supported aspect ratios, offset by 1. For example, if the model supports aspect ratios [[1, 1], [1, 2], [2, 1]]: - An image with aspect ratio [1, 1] would have ID 1 - An image with aspect ratio [1, 2] would have ID 2 - An image with aspect ratio [2, 1] would have ID 3 The id 0 is reserved for padding (i.e., no image). If an image has aspect ratio [1, 2], that means it was split into 2 tiles horizontally, and its `aspect_ratio_id` would be 2. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. If `past_key_values` is used, optionally only the last `input_ids` have to be input (see `past_key_values`). If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*): Cross-attention mask to control the interaction between text tokens and image tiles. This 4D tensor defines which image tiles each text token should attend to. For each text token (in seq_length): - 1 indicates the token **should attend** to the corresponding image tile - 0 indicates the token **should not attend** to the corresponding image tile cross_attention_states (`torch.FloatTensor`, *optional*): Output of the vision model, used for cross-attention. This tensor contains the processed image features that the language model will attend to. position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids) past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. Two formats are allowed: - a [`~cache_utils.Cache`] instance, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache); - 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)`). This is also known as the legacy cache format. The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the legacy cache format will be returned. If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, this tensor is not affected by padding. It is used to update the cache in the correct position and to infer the complete sequence length. """ @add_start_docstrings( """The VideoMllama Vision Model which consists of two vision encoders.""", VideoMllama_START_DOCSTRING, ) class VideoMllamaVisionModel(VideoMllamaPreTrainedModel): config_class = VideoMllamaVisionConfig base_model_prefix = "vision_model" def __init__(self, config: VideoMllamaVisionConfig): super().__init__(config) self.image_size = config.image_size self.patch_size = config.patch_size self.max_num_tiles = config.max_num_tiles self.hidden_size = config.hidden_size self.num_channels = config.num_channels self.intermediate_layers_indices = config.intermediate_layers_indices self.num_patches = (self.image_size // self.patch_size) ** 2 + 1 self.scale = config.hidden_size**-0.5 self.patch_embedding = nn.Conv2d( in_channels=config.num_channels, out_channels=self.hidden_size, kernel_size=self.patch_size, stride=self.patch_size, padding="valid", bias=False, ) self.class_embedding = nn.Parameter(self.scale * torch.randn(self.hidden_size)) self.gated_positional_embedding = VideoMllamaPrecomputedPositionEmbedding(config) self.pre_tile_positional_embedding = VideoMllamaPrecomputedAspectRatioEmbedding(config, is_gated=True) self.post_tile_positional_embedding = VideoMllamaPrecomputedAspectRatioEmbedding(config, is_gated=True) # layer norms self.layernorm_pre = nn.LayerNorm(self.hidden_size) self.layernorm_post = nn.LayerNorm(self.hidden_size) # encoders self.transformer = VideoMllamaVisionEncoder(config, config.num_hidden_layers, is_gated=False) self.global_transformer = VideoMllamaVisionEncoder(config, config.num_global_layers, is_gated=True) self.vision_ignore_attention_mask = config.vision_ignore_attention_mask self.post_init() def get_input_embeddings(self): """ This function is used to fetch the first embedding layer to activate grads on inputs. """ return self.patch_embedding def apply_class_embedding(self, hidden_state: torch.Tensor) -> torch.Tensor: batch_size, _, hidden_size = hidden_state.shape class_embedding = self.class_embedding.expand(batch_size, 1, hidden_size) hidden_state = torch.cat([class_embedding, hidden_state], dim=1) return hidden_state @add_start_docstrings_to_model_forward(VideoMllama_VISION_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=BaseModelOutput, config_class="VideoMllamaVisionConfig") def forward( self, pixel_values: torch.Tensor, aspect_ratio_ids: torch.Tensor, aspect_ratio_mask: torch.Tensor, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]: r""" Returns: Example: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, MllamaVisionModel >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision" >>> model = MllamaVisionModel.from_pretrained(checkpoint) >>> processor = AutoProcessor.from_pretrained(checkpoint) >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(images=image, return_tensors="pt") >>> output = model(**inputs) >>> print(output.last_hidden_state.shape) torch.Size([1, 1, 4, 1025, 7680]) ``` """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict batch_size, num_concurrent_media, num_tiles, num_channels, height, width = pixel_values.shape # pixel_values.shape is (batch_size, num_concurrent_media, num_tiles(video_num_tiles), num_channels, height, width) # aspect_ratio_mask.shape is (batch_size, num_concurrent_media, num_tiles(video_num_tiles)) # DONE pixel_values = pixel_values.reshape(batch_size * num_concurrent_media * num_tiles, num_channels, height, width) # pixel_values.shape is (batch_size * num_concurrent_media * num_tiles(video_num_tiles), num_channels, height, width) # aspect_ratio_ids.shape is (batch_size, num_concurrent_media) aspect_ratio_ids = aspect_ratio_ids.reshape(batch_size * num_concurrent_media, -1) # aspect_ratio_ids.shape is (batch_size * num_concurrent_media, 1) # DONE # Patch embedding patch_embeds = self.patch_embedding(pixel_values.to(self.dtype).to(self.device)) # patch_embeds.shape is (batch_size * num_concurrent_media * num_tiles(video_num_tiles), hidden_size, image_size // patch_size, image_size // patch_size) hidden_state = patch_embeds.flatten(2).transpose(1, 2) # hidden_state.shape is (batch_size * num_concurrent_media * num_tiles(video_num_tiles), num_patches, hidden_size) # num_patches = (image_size // patch_size) ** 2 # DONE # Tile embeddings _, num_patches, dim = hidden_state.shape hidden_state = hidden_state.reshape(batch_size * num_concurrent_media, num_tiles, -1, dim) # hidden_state.shape is (batch_size * num_concurrent_media, num_tiles(video_num_tiles), num_patches, hidden_size) hidden_state = self.pre_tile_positional_embedding(hidden_state, aspect_ratio_ids) # hidden_state.shape is kept the same # DONE # Add cls token hidden_state = hidden_state.reshape(batch_size * num_concurrent_media * num_tiles, num_patches, dim) # hidden_state.shape is (batch_size * num_concurrent_media * num_tiles(video_num_tiles), num_patches, hidden_size) hidden_state = self.apply_class_embedding(hidden_state) num_patches += 1 # !!! NOTICE !!! num_patches is now increased by 1 # DONE # Position embeddings hidden_state = hidden_state.reshape(batch_size * num_concurrent_media, num_tiles, num_patches, dim) # hidden_state.shape is (batch_size * num_concurrent_media, num_tiles(video_num_tiles), num_patches, hidden_size) hidden_state = self.gated_positional_embedding(hidden_state, aspect_ratio_ids) # DONE hidden_state = self.layernorm_pre(hidden_state) # DONE: Whether or not to padd the hidden state ????????? # Compute the number of tokens to pad num_padding_patches = (8 - (hidden_state.shape[-2] % 8)) % 8 # Compute padding tuple for pad function padding = (0, 0, 0, num_padding_patches) # (pad_left, pad_right, pad_left for dim -2, pad_right for dim -2) # Pad the tensor hidden_state = F.pad(hidden_state, padding, mode="constant", value=0) slice_index = -num_padding_patches if num_padding_patches > 0 else None # DONE # Prepare attention mask if self.vision_ignore_attention_mask: # when using full attention, we need to create attn mask attention_mask = None else: # aspect_ratio_mask.shape is (batch_size, num_concurrent_media, num_tiles(video_num_tiles)) attention_mask = aspect_ratio_mask.reshape(batch_size * num_concurrent_media, -1) # aspect_ratio_mask.shape is (batch_size * num_concurrent_media, num_tiles(video_num_tiles)) attention_mask = _prepare_aspect_ratio_attention_mask( aspect_ratio_mask=attention_mask, num_patches=self.num_patches, target_length=hidden_state.shape[2], dtype=self.dtype, ) # attention_mask.shape is (batch_size * num_concurrent_media, num_tiles(video_num_tiles) * num_patches, num_tiles(video_num_tiles) * num_patches) # NOTICE: the num_patches here is the padded hidden_state's num_patches, not the original one # DONE # Apply encoder hidden_state = hidden_state.view(batch_size * num_concurrent_media, -1, dim) # hidden_state.shape is (batch_size * num_concurrent_media, num_tiles(video_num_tiles) * num_patches, dim) output = self.transformer( hidden_state, attention_mask=attention_mask, output_hidden_states=True, output_attentions=output_attentions, ) hidden_state = output[0] hidden_state = self.layernorm_post(hidden_state) # DONE # Apply global encoder hidden_state = hidden_state.reshape( batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, dim ) # hidden_state.shape is (batch_size * num_concurrent_media, num_tiles(video_num_tiles), num_patches, dim) hidden_state = self.post_tile_positional_embedding(hidden_state, aspect_ratio_ids) hidden_state = hidden_state.reshape( batch_size * num_concurrent_media, num_tiles * (num_patches + num_padding_patches), dim ) global_output = self.global_transformer( hidden_state, attention_mask=attention_mask, output_hidden_states=output_hidden_states, output_attentions=output_attentions, ) hidden_state = global_output[0] # DONE # Remove padding form hidden state hidden_state = hidden_state.reshape( batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, dim ) hidden_state = hidden_state[:, :, :slice_index] hidden_state = hidden_state.reshape(batch_size, num_concurrent_media, num_tiles, num_patches, dim) # DONE # Collect intermediate layer outputs from encoder output # all_intermediate_hidden_states = output[1] # intermediate_hidden_states = torch.stack(all_intermediate_hidden_states, dim=-1) # intermediate_hidden_states = intermediate_hidden_states[..., self.intermediate_layers_indices] # Collect intermediate layer outputs from encoder output all_intermediate_hidden_states = [output[1][i] for i in self.intermediate_layers_indices] intermediate_hidden_states = torch.stack(all_intermediate_hidden_states, dim=-1) # DONE # Remove padding from intermediate hidden states intermediate_hidden_states = intermediate_hidden_states.reshape( batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, -1 ) intermediate_hidden_states = intermediate_hidden_states[:, :, :slice_index] intermediate_hidden_states = intermediate_hidden_states.reshape( batch_size, num_concurrent_media, num_tiles, num_patches, -1 ) # DONE # Concatenate final hidden state and intermediate hidden states hidden_state = torch.cat([hidden_state, intermediate_hidden_states], dim=-1) # DONE if output_hidden_states: hidden_states = tuple(all_intermediate_hidden_states) + tuple(global_output[1]) else: hidden_states = None if output_attentions: # global transformer in contrast to `self.transformer` doesn't always return hidden states so we might go index out-of-range global_attn = tuple(global_output[2]) if output_hidden_states else tuple(global_output[1]) attentions = tuple(output[2]) + global_attn else: attentions = None if not return_dict: return tuple(v for v in [hidden_state, hidden_states, attentions] if v is not None) return BaseModelOutput( last_hidden_state=hidden_state, hidden_states=hidden_states, attentions=attentions, ) @add_start_docstrings( """The VideoMllama Text Model which consists of transformer with self and cross attention layers.""", VideoMllama_START_DOCSTRING, ) class VideoMllamaTextModel(VideoMllamaPreTrainedModel): config_class = VideoMllamaTextConfig base_model_prefix = "language_model.model" def __init__(self, config: VideoMllamaTextConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size + 8, config.hidden_size, self.padding_idx) self.cross_attention_layers = config.cross_attention_layers layers = [] for layer_idx in range(config.num_hidden_layers): if layer_idx in self.cross_attention_layers: layers.append(VideoMllamaCrossAttentionDecoderLayer(config, layer_idx)) else: layers.append(VideoMllamaSelfAttentionDecoderLayer(config, layer_idx)) self.layers = nn.ModuleList(layers) self.norm = VideoMllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = VideoMllamaRotaryEmbedding(config=config) self.gradient_checkpointing = False self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value @add_start_docstrings_to_model_forward(VideoMllama_TEXT_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=BaseModelOutputWithPast, config_class="VideoMllamaTextConfig") def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, vision_position_ids: Optional[torch.LongTensor] = None, # vision_position_ids cross_attention_states: Optional[torch.FloatTensor] = None, cross_attention_mask: Optional[torch.Tensor] = None, full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, vision_cache_position: Optional[torch.LongTensor] = None, # vision_cache_position ) -> Union[Tuple, BaseModelOutputWithPast]: """ Returns: Example: ```python >>> from transformers import AutoProcessor, MllamaTextModel >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision" >>> model = MllamaTextModel.from_pretrained(checkpoint) >>> processor = AutoProcessor.from_pretrained(checkpoint) >>> text = "<|image|>If I had to write a haiku for this one" >>> inputs = processor(text=text, return_tensors="pt") >>> output = model(**inputs) >>> print(output.last_hidden_state.shape) torch.Size([1, 13, 4096]) ``` """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if self.gradient_checkpointing and self.training and use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." ) use_cache = False if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) hidden_states = inputs_embeds if use_cache and past_key_values is None: past_key_values = DynamicCache() if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) if position_ids is None: position_ids = cache_position.unsqueeze(0) causal_mask = self._update_causal_mask( attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions ) # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) if cross_attention_states is not None: # cross_attention_states.shape is (batch_size, num_concurrent_media, num_tiles, num_patches, dim) batch_size, num_concurrent_media, num_tiles, num_patches, hidden_size = cross_attention_states.shape ## delay this reshape operation from VideoMllamaForConditionalGeneration.forward() cross_attention_states = cross_attention_states.reshape( -1, cross_attention_states.shape[-2], hidden_size ) # cross_attention_states.shape is (batch_size * num_concurrent_media * num_tiles, num_patches, dim) # Vision Position: calculate vision_cache_position and vision_position_ids if vision_cache_position is None: # cross_attention_states.shape is (batch_size * num_concurrent_media * num_tiles, num_patches, dim) past_seen_media = past_key_values.get_seq_length(self.cross_attention_layers[0]) if past_key_values is not None else 0 vision_cache_position = torch.arange( past_seen_media, past_seen_media + num_concurrent_media, device=cross_attention_states.device ) # vision_cache_position.shape be: (num_concurrent_media, ) if vision_position_ids is None: vision_position_ids = vision_cache_position.unsqueeze(0) # vision_position_ids.shape be: (1, num_concurrent_media) assert num_concurrent_media == len(vision_cache_position), "num_concurrent_media is not equal to len(vision_cache_position)" # Vision Position: calculate vision_position_embeddings # create vision position embeddings to be shared across the decoder layers vision_position_embeddings = self.rotary_emb(cross_attention_states, vision_position_ids) vision_cos, vision_sin = vision_position_embeddings # vision_cos.shape is (batch_size, num_concurrent_media, head_dim) # vision_cos.shape should be (batch_size, num_concurrent_media * num_tiles * num_patches, head_dim) vision_cos = torch.repeat_interleave(vision_cos, repeats=num_tiles * num_patches, dim=1) vision_sin = torch.repeat_interleave(vision_sin, repeats=num_tiles * num_patches, dim=1) vision_position_embeddings = (vision_cos, vision_sin) # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None next_decoder_cache = None for idx, decoder_layer in enumerate(self.layers): if output_hidden_states: all_hidden_states += (hidden_states,) # For text-only path we should skip cross attention layers. # Let's check if the layer is cross attention layer and if we have cross attention states # or cached cross attention states. is_cross_attention_layer = idx in self.cross_attention_layers is_cross_attention_cache_empty = past_key_values is None or ( past_key_values is not None and past_key_values.get_seq_length(idx) == 0 ) if is_cross_attention_layer and cross_attention_states is None and is_cross_attention_cache_empty: continue if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, cross_attention_states, cross_attention_mask, causal_mask, full_text_row_masked_out_mask, position_ids, vision_position_ids, # vision_position_ids past_key_values, output_attentions, use_cache, cache_position, vision_cache_position, # vision_cache_position position_embeddings, vision_position_embeddings if cross_attention_states is not None else None, # vision_position_embeddings ) else: layer_outputs = decoder_layer( hidden_states, cross_attention_states=cross_attention_states, cross_attention_mask=cross_attention_mask, attention_mask=causal_mask, full_text_row_masked_out_mask=full_text_row_masked_out_mask, position_ids=position_ids, vision_position_ids=vision_position_ids, # vision_position_ids past_key_value=past_key_values, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, vision_cache_position=vision_cache_position, # vision_cache_position position_embeddings=position_embeddings, vision_position_embeddings=vision_position_embeddings if cross_attention_states is not None else None, # vision_position_embeddings ) hidden_states = layer_outputs[0] if use_cache: next_decoder_cache = layer_outputs[2 if output_attentions else 1] if output_attentions: all_self_attns += (layer_outputs[1],) hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = next_decoder_cache if use_cache else None if not return_dict: return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, ) @add_start_docstrings( """The VideoMllama Text Model with a language modeling head on top.""", VideoMllama_START_DOCSTRING, ) class VideoMllamaForCausalLM(VideoMllamaPreTrainedModel, GenerationMixin): config_class = VideoMllamaTextConfig _supports_static_cache = True # only the LLM without cross attn can do compile base_model_prefix = "language_model" _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): super().__init__(config.get_text_config()) self.text_config = config.get_text_config() self.vocab_size = self.text_config.vocab_size self.model = VideoMllamaTextModel._from_config(self.text_config) self.lm_head = nn.Linear(self.text_config.hidden_size, self.vocab_size, bias=False) self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def set_decoder(self, decoder): self.model = decoder def get_decoder(self): return self.model @add_start_docstrings_to_model_forward(VIDEOMLLAMA_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class="VideoMllamaTextConfig") def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, vision_position_ids: Optional[torch.LongTensor] = None, # vision_position_ids cross_attention_states: Optional[torch.LongTensor] = None, cross_attention_mask: Optional[torch.LongTensor] = None, full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, vision_cache_position: Optional[torch.LongTensor] = None, # vision_cache_position num_logits_to_keep: int = 0, **loss_kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: r""" Args: 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]`. num_logits_to_keep (`int`, *optional*): Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that token can save memory, which becomes pretty significant for long sequences or large vocabulary size. Returns: Example: ```python >>> from transformers import AutoTokenizer, MllamaForCausalLM >>> model = MllamaForCausalLM.from_pretrained("Llama-3.2-11B-Vision") >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision") >>> prompt = "If I had to write a haiku, it would be:" >>> inputs = tokenizer(prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6) >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] >>> print(result) If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful. I love the idea of snowflakes gently falling, each one ``` """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, cross_attention_states=cross_attention_states, attention_mask=attention_mask, position_ids=position_ids, vision_position_ids=vision_position_ids, # vision_position_ids cross_attention_mask=cross_attention_mask, full_text_row_masked_out_mask=full_text_row_masked_out_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, vision_cache_position=vision_cache_position, # vision_cache_position ) hidden_states = outputs[0] logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float() loss = None if labels is not None: loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs) if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) @add_start_docstrings( """The VideoMllama model which consists of a vision encoder and a language model.""", VideoMllama_START_DOCSTRING, ) class VideoMllamaForConditionalGeneration(VideoMllamaPreTrainedModel, GenerationMixin): _supports_quantized_cache = False # quant cache not supported in encoder-decoder setting def __init__(self, config: VideoMllamaConfig): super().__init__(config) self.vocab_size = config.text_config.vocab_size self.hidden_size = config.text_config.hidden_size self.max_num_tiles = config.vision_config.max_num_tiles self.vision_output_dim = config.vision_config.vision_output_dim self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1 self.vision_model = VideoMllamaVisionModel._from_config(config.vision_config) self.language_model = VideoMllamaForCausalLM._from_config(config.text_config) self.multi_modal_projector = nn.Linear( config.vision_config.vision_output_dim, config.text_config.hidden_size, bias=True, ) self.merge_size = config.vision_config.merge_size self.merge_mode = config.vision_config.merge_mode if self.merge_mode is not None and self.merge_size is None: raise ValueError("merge_size must be a positive number when merge_mode is not None") if self.merge_mode == "channelFusion": self.patch_merger = nn.Conv2d( in_channels=self.vision_output_dim, out_channels=self.vision_output_dim, kernel_size=self.merge_size, stride=self.merge_size, bias=False, ) self.num_vision_tokens = 0 # A variable to record the number of vision tokens in each frame, primarily used in the generate() function. self.continue_generating = False 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) def get_output_embeddings(self): return self.language_model.get_output_embeddings() def set_output_embeddings(self, new_embeddings): self.language_model.set_output_embeddings(new_embeddings) def set_decoder(self, decoder): self.language_model.set_decoder(decoder) def get_decoder(self): return self.language_model.get_decoder() def tie_weights(self): return self.language_model.tie_weights() def start_real_time_generate(self): self.continue_generating = True def stop_real_time_generate(self): # clear cache gc.collect() torch.cuda.empty_cache() # set continue_generating to False self.continue_generating = False @add_start_docstrings_to_model_forward(VIDEOMLLAMA_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class="VideoMllamaConfig") def forward( self, input_ids: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, aspect_ratio_mask: Optional[torch.Tensor] = None, aspect_ratio_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, cross_attention_mask: Optional[torch.Tensor] = None, cross_attention_states: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, vision_position_ids: Optional[torch.LongTensor] = None, # vision_position_ids is used to specify video position_id past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, vision_cache_position: Optional[torch.LongTensor] = None, # vision_cache_position is used to specify video cache_position_id num_logits_to_keep: int = 0, ) -> Union[Tuple, CausalLMOutputWithPast]: r""" Args: 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]`. num_logits_to_keep (`int`, *optional*): Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that token can save memory, which becomes pretty significant for long sequences or large vocabulary size. Returns: Example: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, MllamaForConditionalGeneration >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision" >>> model = MllamaForConditionalGeneration.from_pretrained(checkpoint) >>> processor = AutoProcessor.from_pretrained(checkpoint) >>> prompt = "<|image|>If I had to write a haiku for this one" >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> inputs = processor(text=prompt, images=image, return_tensors="pt") >>> # Generate >>> output = model.generate(**inputs, max_new_tokens=15) >>> prompt_len = inputs.input_ids.shape[-1] >>> generated_ids = output[:, prompt_len:] >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) >>> print(generated_text) [', it would be:.\\nA stop sign in Chinatown.\\n'] ``` """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if pixel_values is not None and inputs_embeds is not None: raise ValueError( "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" ) if pixel_values is not None and cross_attention_states is not None: raise ValueError("`pixel_values` and `cross_attention_states` cannot be provided simultaneously") if pixel_values is not None: if aspect_ratio_ids is None: raise ValueError("`aspect_ratio_ids` must be provided if `pixel_values` is provided") # get vision tokens from vision model vision_outputs = self.vision_model( pixel_values=pixel_values, aspect_ratio_ids=aspect_ratio_ids, aspect_ratio_mask=aspect_ratio_mask, output_hidden_states=output_hidden_states, output_attentions=output_attentions, return_dict=return_dict, ) cross_attention_states = vision_outputs[0] if self.merge_mode is not None: cross_attention_states = self.get_2dPool(cross_attention_states) if self.num_vision_tokens == 0: self.num_vision_tokens = cross_attention_states.shape[-2] # cross_attention_states.shape is (batch_size, num_concurrent_media, num_tiles, num_patches, dim) cross_attention_states = self.multi_modal_projector(cross_attention_states) ## delay this reshape operation to VideoMllamaTextModel.forward() # cross_attention_states = cross_attention_states.reshape( # -1, cross_attention_states.shape[-2], self.hidden_size # ) # After reshape(), cross_attention_states.shape is (batch_size * num_concurrent_media * num_tiles, num_patches, dim) if cross_attention_mask is not None: assert self.num_vision_tokens != 0, "Please correctly calculate or pass the value of 'num_vision_tokens'." cross_attention_mask, full_text_row_masked_out_mask = _prepare_cross_attention_mask( cross_attention_mask, num_vision_tokens=self.num_vision_tokens, dtype=self.dtype, ) else: full_text_row_masked_out_mask = None # TODO: Streaming mode, we can just pass the current_new_inputs' cross_attention_mask to the model if cross_attention_mask is not None and cache_position is not None and cross_attention_mask.shape[2] != cache_position.shape[0]: # cross_attention_mask.shape is (batch_size, 1, seq_len, num_concurrent_media * num_tiles * num_patches) cross_attention_mask = cross_attention_mask[:, :, cache_position] full_text_row_masked_out_mask = full_text_row_masked_out_mask[:, :, cache_position] outputs = self.language_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, vision_position_ids=vision_position_ids, # vision_position_ids cross_attention_states=cross_attention_states, cross_attention_mask=cross_attention_mask, full_text_row_masked_out_mask=full_text_row_masked_out_mask, past_key_values=past_key_values, use_cache=use_cache, inputs_embeds=inputs_embeds, labels=labels, output_hidden_states=output_hidden_states, output_attentions=output_attentions, return_dict=return_dict, cache_position=cache_position, vision_cache_position=vision_cache_position, # vision_cache_position num_logits_to_keep=num_logits_to_keep, ) return outputs @add_start_docstrings_to_model_forward(VIDEOMLLAMA_INPUTS_DOCSTRING) @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class="VideoMllamaConfig") def vision_chunked_forward( self, input_ids: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, aspect_ratio_mask: Optional[torch.Tensor] = None, aspect_ratio_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, cross_attention_mask: Optional[torch.Tensor] = None, cross_attention_states: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, vision_position_ids: Optional[torch.LongTensor] = None, # vision_position_ids is used to specify video position_id past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, vision_cache_position: Optional[torch.LongTensor] = None, # vision_cache_position is used to specify video cache_position_id num_logits_to_keep: int = 0, vision_chunked_length: int = 64, ) -> Union[Tuple, CausalLMOutputWithPast]: r""" Args: The only additional parameter is `vision_chunked_length`, while the other parameters are exactly the same as `forward`. Explanation: Images are fed into self.vision_model in batches based on vision_chunked_length, while the rest is identical to the original `forward` function. Returns: """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if pixel_values is not None and inputs_embeds is not None: raise ValueError( "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one" ) if pixel_values is not None and cross_attention_states is not None: raise ValueError("`pixel_values` and `cross_attention_states` cannot be provided simultaneously") if pixel_values is not None: if aspect_ratio_ids is None: raise ValueError("`aspect_ratio_ids` must be provided if `pixel_values` is provided") # Get vision tokens from vision model pixel_chunks = pixel_values.chunk(math.ceil(pixel_values.shape[1] / vision_chunked_length), dim=1) aspect_ratio_id_chunks = aspect_ratio_ids.chunk(math.ceil(aspect_ratio_ids.shape[1] / vision_chunked_length), dim=1) aspect_ratio_mask_chunks = aspect_ratio_mask.chunk(math.ceil(aspect_ratio_mask.shape[1] / vision_chunked_length), dim=1) # 批量处理并拼接结果 cross_attention_states = torch.cat([ self.vision_model( pixel_values=pv, aspect_ratio_ids=ari, aspect_ratio_mask=arm, output_hidden_states=output_hidden_states, output_attentions=output_attentions, return_dict=return_dict, )[0] for pv, ari, arm in zip(pixel_chunks, aspect_ratio_id_chunks, aspect_ratio_mask_chunks) ], dim=1) if self.merge_mode is not None: cross_attention_states = self.get_2dPool(cross_attention_states) if self.num_vision_tokens == 0: self.num_vision_tokens = cross_attention_states.shape[-2] # cross_attention_states.shape is (batch_size, num_concurrent_media, num_tiles, num_patches, dim) cross_attention_states = self.multi_modal_projector(cross_attention_states) ## delay this reshape operation to VideoMllamaTextModel.forward() # cross_attention_states = cross_attention_states.reshape( # -1, cross_attention_states.shape[-2], self.hidden_size # ) # After reshape(), cross_attention_states.shape is (batch_size * num_concurrent_media * num_tiles, num_patches, dim) if cross_attention_mask is not None: assert self.num_vision_tokens != 0, "Please correctly calculate or pass the value of 'num_vision_tokens'." cross_attention_mask, full_text_row_masked_out_mask = _prepare_cross_attention_mask( cross_attention_mask, num_vision_tokens=self.num_vision_tokens, dtype=self.dtype, ) else: full_text_row_masked_out_mask = None # TODO: Streaming mode, we can just pass the current_new_inputs' cross_attention_mask to the model if cross_attention_mask is not None and cache_position is not None and cross_attention_mask.shape[2] != cache_position.shape[0]: # cross_attention_mask.shape is (batch_size, 1, seq_len, num_concurrent_media * num_tiles * num_patches) cross_attention_mask = cross_attention_mask[:, :, cache_position] full_text_row_masked_out_mask = full_text_row_masked_out_mask[:, :, cache_position] outputs = self.language_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, vision_position_ids=vision_position_ids, # vision_position_ids cross_attention_states=cross_attention_states, cross_attention_mask=cross_attention_mask, full_text_row_masked_out_mask=full_text_row_masked_out_mask, past_key_values=past_key_values, use_cache=use_cache, inputs_embeds=inputs_embeds, labels=labels, output_hidden_states=output_hidden_states, output_attentions=output_attentions, return_dict=return_dict, cache_position=cache_position, vision_cache_position=vision_cache_position, # vision_cache_position num_logits_to_keep=num_logits_to_keep, ) return outputs def get_2dPool(self, hidden_state): """ Args: hidden_state: (batch_size, num_concurrent_media, num_tiles, num_patches, dim) """ height = width = self.vision_model.image_size // self.vision_model.patch_size batch_size, num_concurrent_media, num_tiles, num_patches, dim = hidden_state.shape # num_patches = 1 + height * width cls_hidden_state = hidden_state[:, :, :, :1, :] img_hidden_state = hidden_state[:, :, :, 1:, :] image_feature = img_hidden_state.reshape(batch_size * num_concurrent_media * num_tiles, height, width, dim) # these codes are referenced from llava-next.llava.model.llava_arch.py.LlavaMetaForCausalLM.get_2dPool() image_feature = image_feature.permute(0, 3, 1, 2).contiguous() if self.merge_mode == "average": image_feature = nn.functional.avg_pool2d(image_feature, self.merge_size) elif self.merge_mode == "max": image_feature = nn.functional.max_pool2d(image_feature, self.merge_size) elif self.merge_mode == "bilinear": scaled_shape = [math.ceil(height / self.merge_size), math.ceil(width / self.merge_size)] image_feature = nn.functional.interpolate(image_feature, size=scaled_shape, mode='bilinear') elif self.merge_mode == "channelFusion": image_feature = self.patch_merger(image_feature) else: raise ValueError(f"Unexpected merge_mode: {self.merge_mode}") image_feature = image_feature.permute(0, 2, 3, 1) image_feature = image_feature.view(batch_size * num_concurrent_media * num_tiles, -1, dim) img_hidden_state = image_feature.reshape(batch_size, num_concurrent_media, num_tiles, -1, dim) hidden_state = torch.cat([cls_hidden_state, img_hidden_state], dim=-2) return hidden_state def prepare_inputs_for_generation( self, input_ids=None, inputs_embeds=None, attention_mask=None, position_ids=None, vision_position_ids=None, # vision_position_ids is used to specify video position_id pixel_values=None, aspect_ratio_ids=None, aspect_ratio_mask=None, cross_attention_mask=None, past_key_values=None, use_cache=False, cache_position=None, vision_cache_position=None, # vision_cache_position is used to specify video cache_position_id num_logits_to_keep=None, **kwargs, ): # Overwritten -- in specific circumstances we don't want to forward image inputs to the model # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens # Exception 1: when passing input_embeds, input_ids may be missing entries # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here if past_key_values is not None: # If `inputs_embeds` are passed, `input_ids` will be ignored # But notice that we only want to use them in the ** 1st ** generation step if inputs_embeds is not None and cache_position[0] == 0: assert input_ids.shape[1] == 0, "input_ids.shape[1] != 0" # Otherwise, we need to slice `input_ids` through `cache_position` elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) input_ids = input_ids[:, cache_position] assert input_ids.shape[1] == cache_position.shape[0], "input_ids.shape[1] != cache_position.shape[0]" # TODO: we have no attention_mask so this won't work, check if we really won't need attention mask and find another way if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, VIDEO_MLLAMA_PROCESSOR_PAD_POSITION_ID) if past_key_values and position_ids is not None: position_ids = position_ids[:, -cache_position.shape[0] :] # Clone to avoid recapturing cuda graphs with torch.compile's reduce-overhead mode position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and cache_position[0] == 0: model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} else: # The clone here is for the same reason as for `position_ids`. model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None} if num_logits_to_keep is not None: model_inputs["num_logits_to_keep"] = num_logits_to_keep if cross_attention_mask is not None: cross_attention_mask = cross_attention_mask[:, -cache_position.shape[0] :] model_inputs.update( { "position_ids": position_ids, "vision_position_ids": vision_position_ids, "cache_position": cache_position, "vision_cache_position": vision_cache_position, "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, "cross_attention_mask": cross_attention_mask, } ) # If we're in pre-fill or cacheless decoding step, then we need pixel_values and aspect ratios # to compute image hidden states, otherwise they are cached within each cross attn layer image_token_id = self.config.image_token_index final_input_ids = model_inputs["input_ids"] if "input_ids" in model_inputs else None if final_input_ids is not None and image_token_id in final_input_ids: # count the number of image tokens in the final_input_ids if cache_position[0] == 0: image_token_count = pixel_values.shape[1] else: image_token_count = (final_input_ids == image_token_id).sum(dim=-1).max().item() # truncate the `vision_position_ids`, `pixel_values`, `aspect_ratio_ids`, `aspect_ratio_mask` model_inputs["vision_position_ids"] = vision_position_ids[:, -image_token_count:] model_inputs["pixel_values"] = pixel_values[:, -image_token_count:] model_inputs["aspect_ratio_ids"] = aspect_ratio_ids[:, -image_token_count:] model_inputs["aspect_ratio_mask"] = aspect_ratio_mask[:, -image_token_count:] return model_inputs # raise NotImplementedError("`prepare_inputs_for_generation` has not been implemented.") def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder, **kwargs): # TODO Currently, our generation code only supports OFFLINE visual understanding, # TODO meaning that all multimodal content and questions must be provided before generating the answer. # TODO NEED to implement new inference code to support real-time streaming visual understanding. # Notice: `vision_position_ids` and `vision_cache_position` will remain unchanged in **offline** scenarios. ### Notice ### # The following code was directly copied from transformers.modeling_mllama.MllamaForConditionalGeneration._update_model_kwargs_for_generation() cross_attention_mask_prev = model_kwargs.get("cross_attention_mask", None) position_ids_prev = model_kwargs.get("position_ids", None) model_kwargs = super()._update_model_kwargs_for_generation( outputs=outputs, model_kwargs=model_kwargs, is_encoder_decoder=is_encoder_decoder, **kwargs, ) # add cross-attn mask for new token if cross_attention_mask_prev is not None: # each image or video frame has its own token in text token list # we just need to copy last cross_attention_mask for the current newly generated token. model_kwargs["cross_attention_mask"] = torch.cat( [cross_attention_mask_prev, cross_attention_mask_prev[:, -1:, ...]], dim=1 ) # add position_id for new token if position_ids_prev is not None: last_position_ids = position_ids_prev[:, -1:] cur_position_ids = last_position_ids + 1 model_kwargs["position_ids"] = torch.cat( [position_ids_prev, cur_position_ids], dim=1 ) return model_kwargs # raise NotImplementedError("`_update_model_kwargs_for_generation` has not been implemented.") def offline_generate( self, processor, new_queries: queue.Queue, output_text_queue: queue.Queue, vision_chunked_length: int = 64, ): """ Args: processor: The processor for preparing model inputs. new_queries (queue.Queue): A queue to which new query are added. A query includes: ``` { "prompt": "", "images": [], "videos": [], "media_kwargs": { "video_fps": 1.0, "video_minlen": 8, "video_maxlen": 256 }, "system_prompt_type": "text_image / video", "thinking_mode": "no_thinking / deep_thinking", "generate_kwargs": { "temperature": 1.0, "top_k": 50, "top_p": 1.0, "max_new_tokens": 1024, "repetition_penalty": 1.0 }, "stop_offline_generate": False, } ``` output_text_queue (queue.Queue): A queue to which generated text is added. """ system_prompts = { # no_thinking system_prompt "no_thinking": { # text or image "text_image": "You are a helpful AI assistant. Respond to the user's request based on the provided text and/or images.", # video "video": "You are a helpful AI assistant specializing in video analysis. Respond to the user's request based on the provided video content.", }, # deep_thinking system_prompt "deep_thinking": { # text or image "text_image": "A conversation between User and Assistant. The user makes a request, and the assistant responds to it based on the provided text and/or images. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within and tags, respectively, i.e., reasoning process here answer here .", # video "video": "A conversation between User and Assistant specializing in video analysis. The user makes a request, and the assistant responds to it based on the provided video content. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within and tags, respectively, i.e., reasoning process here answer here .", } } new_video_frames = queue.Queue() new_prompts = queue.Queue() max_tokens_per_turn = 86400 images, videos = [], [] past_key_values = None is_first_turn = True while True: self.start_real_time_generate() # Get a new request current_query = new_queries.get() print("A new query has been passed.") # Determine whether to end the conversation if current_query["stop_offline_generate"]: print("Current conversation has been ended.") break if is_first_turn: thinking_mode, system_prompt_type = current_query["thinking_mode"], current_query["system_prompt_type"] messages = [ {"role": "system", "content": system_prompts[thinking_mode][system_prompt_type]}, ] is_first_turn = False cur_generate_kwargs = current_query["generate_kwargs"] if cur_generate_kwargs["temperature"] == 0.0: cur_generate_kwargs["do_sample"] = False cur_generate_kwargs["temperature"] = 1.0 cur_images = current_query["images"] cur_videos = current_query["videos"] images.extend(cur_images) videos.extend(cur_videos) media_kwargs = current_query["media_kwargs"] prompt = current_query["prompt"] prompt = prompt.replace(processor.image_token, "[image]").replace(processor.image_placeholder, "[image]") prompt = prompt.replace(processor.video_token, "[video]").replace(processor.video_placeholder, "[video]") prompt = processor.image_token * len(cur_images) + processor.video_token * len(cur_videos) + prompt messages.append({"role": "user", "content": prompt}) current_input_text = processor.apply_chat_template(messages, add_generation_prompt=True) current_inputs = processor(text=current_input_text, images=images, videos=videos, **media_kwargs, add_special_tokens=False, return_tensors="pt").to(self.device) current_input_text = processor.decode(current_inputs.input_ids[0]) try: outputs = self._real_time_generate( new_video_frames=new_video_frames, new_prompts=new_prompts, output_text_queue=output_text_queue, processor=processor, **current_inputs, **cur_generate_kwargs, max_tokens_per_turn=max_tokens_per_turn, is_in_offline_generate=True, return_dict_in_generate=True, past_key_values =past_key_values, vision_chunked_length = vision_chunked_length, ) past_key_values = outputs.past_key_values output_content = processor.decode(outputs.sequences[0]) output_content = output_content[len(current_input_text):] if output_content.endswith("<|eot_id|>"): output_content = output_content[:-len("<|eot_id|>")] messages.append({"role": "user", "content": output_content}) output_text_queue.put("<|round_end|>") except Exception as e: print(f"Error occurred: {e}") print(f"Error type: {type(e).__name__}") self.stop_real_time_generate() return finally: self.stop_real_time_generate() def prepare_inputs_for_real_time_generation( self, input_ids=None, inputs_embeds=None, attention_mask=None, position_ids=None, vision_position_ids=None, pixel_values=None, aspect_ratio_ids=None, aspect_ratio_mask=None, cross_attention_mask=None, past_key_values=None, use_cache=False, cache_position=None, vision_cache_position=None, num_logits_to_keep=None, **kwargs, ): # Verify that the signatures of both functions are identical. sig_prepare = inspect.signature(self.prepare_inputs_for_generation) sig_real_time_prepare = inspect.signature(self.prepare_inputs_for_real_time_generation) if sig_prepare != sig_real_time_prepare: raise ValueError( "`prepare_inputs_for_real_time_generation` should have the same signature as `prepare_inputs_for_generation`." ) # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens if past_key_values is not None: # If `inputs_embeds` are passed, `input_ids` will be ignored # But notice that we only want to use them in the ** 1st ** generation step if inputs_embeds is not None and cache_position[0] == 0: assert input_ids.shape[1] == 0, "input_ids.shape[1] != 0" # Otherwise, we need to slice `input_ids` through `cache_position` elif input_ids.shape[1] != cache_position.shape[0]: input_ids = input_ids[:, cache_position] assert input_ids.shape[1] == cache_position.shape[0], "input_ids.shape[1] != cache_position.shape[0]" # If `position_ids` is not provided, we need to create it on the fly if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, VIDEO_MLLAMA_PROCESSOR_PAD_POSITION_ID) # If we have cache: let's slice `position_ids` through `cache_position`, to keep only the unprocessed tokens if past_key_values and position_ids is not None: position_ids = position_ids[:, -cache_position.shape[0] :] # Clone to avoid recapturing cuda graphs with torch.compile's reduce-overhead mode position_ids = position_ids.clone(memory_format=torch.contiguous_format) # if `inputs_embeds` are passed, `input_ids` will be ignored # But notice that we only want to use them in the ** 1st ** generation step if inputs_embeds is not None and cache_position[0] == 0: model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} # Otherwise, we need to use `input_ids` else: # The clone here is for the same reason as for `position_ids`. model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None} # pass the `attention_mask`, `position_ids`, `cache_position` to the model model_inputs["attention_mask"] = attention_mask model_inputs["position_ids"] = position_ids model_inputs["cache_position"] = cache_position # pass the `use_cache` and `past_key_values` to the model model_inputs["use_cache"] = use_cache model_inputs["past_key_values"] = past_key_values # If `num_logits_to_keep` is provided, we need to pass it to the model if num_logits_to_keep is not None: model_inputs["num_logits_to_keep"] = num_logits_to_keep # Note that: # The `vision_cache_position` actually retains the complete cache_position_ids for all images. # Therefore, when using it, we need to truncate it according to the actual num_images in `pixel_values`. if pixel_values is not None and vision_position_ids is not None: # pixel_values.shape: (batch_size, num_images, max_num_tiles, channel, height, width) # calculate cur_image_num cur_image_num = pixel_values.shape[1] # calculate image_token_count final_input_ids = model_inputs["input_ids"] if "input_ids" in model_inputs else None if final_input_ids is not None: image_token_count = (final_input_ids == self.config.image_token_index).sum(dim=-1).max().item() else: image_token_count = pixel_values.shape[1] # Case 1: typically in multi-turn [offline generate mode] if image_token_count < cur_image_num: if image_token_count <= 0: # truncate `vision_cache_position` and `vision_cache_position` vision_position_ids, vision_cache_position = None, None # truncate `pixel_values`, `aspect_ratio_ids` and `aspect_ratio_mask` pixel_values, aspect_ratio_ids, aspect_ratio_mask = None, None, None else: # truncate `vision_cache_position` and `vision_cache_position` vision_position_ids = vision_position_ids[:, -image_token_count:] if vision_cache_position is not None: vision_cache_position = vision_cache_position[-image_token_count:] # truncate `pixel_values`, `aspect_ratio_ids` and `aspect_ratio_mask` pixel_values = pixel_values[:, -image_token_count:] aspect_ratio_ids = aspect_ratio_ids[:, -image_token_count:] aspect_ratio_mask = aspect_ratio_mask[:, -image_token_count:] # Case 2: Invalid case elif image_token_count > cur_image_num: raise ValueError( f"image_token_count is {image_token_count}, cur_image_num is {cur_image_num}. " f"This is invalid. image_token_count must be less than or equal to cur_image_num." ) # Case 3: first turn in [offline generate mode] or in [realtime generate mode] else: # truncate the `vision_cache_position` if vision_cache_position is not None: vision_cache_position = vision_cache_position[-cur_image_num:] vision_cache_position = vision_cache_position.clone(memory_format=torch.contiguous_format) # truncate the `vision_position_ids`, and clone to avoid recapturing cuda graphs with torch.compile's reduce-overhead mode vision_position_ids = vision_position_ids[:, -cur_image_num:] vision_position_ids = vision_position_ids.clone(memory_format=torch.contiguous_format) else: vision_cache_position = None vision_position_ids = None # pass the `pixel_values`, `aspect_ratio_ids`, `aspect_ratio_mask`, `vision_position_ids`, `vision_cache_position`, `cross_attention_mask` to the model model_inputs.update( { "pixel_values": pixel_values, "aspect_ratio_ids": aspect_ratio_ids, "aspect_ratio_mask": aspect_ratio_mask, "vision_position_ids": vision_position_ids, "vision_cache_position": vision_cache_position, "cross_attention_mask": cross_attention_mask } ) return model_inputs def _update_model_kwargs_for_real_time_generation( self, outputs, input_ids, model_kwargs, should_wait_for_new_input: bool, is_encoder_decoder, new_video_frames: queue.Queue, new_prompts: queue.Queue, output_text_queue: queue.Queue, token_buffer: deque, processor, **kwargs, ): # Real-time input processing frames_to_process = [] prompts_to_process = [] # Drain the video frames queue completely while True: while not new_video_frames.empty(): try: frames_to_process.append(new_video_frames.get_nowait()) except queue.Empty: # This should not happen due to the `while` condition, but it's a safeguard. break # Drain the prompts queue completely while not new_prompts.empty(): try: prompts_to_process.append(new_prompts.get_nowait()) output_text_queue.put("<|round_start|>") token_buffer.clear() except queue.Empty: # This should not happen, but it's a safeguard. break if self.continue_generating and should_wait_for_new_input and not frames_to_process and not prompts_to_process: continue else: break # This variable will hold the new text to be tokenized and appended to the input_ids. text_to_append = "" # prefix and suffix for real-time chat real_time_chat_prefix = "<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n" real_time_chat_suffix = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" # If there are new prompts, format them and append to the text to be tokenized if prompts_to_process: for prompt in prompts_to_process: formatted_prompt = f"{real_time_chat_prefix}{prompt}{real_time_chat_suffix}" text_to_append += formatted_prompt # Then, if there are new video frames, prepare their text representation. if frames_to_process: # If a prompt was also added in this turn, add a separator token. if prompts_to_process: text_to_append += "<|silence|>" # Create a string of <|image|> tokens, separated by <|silence|>. image_tokens_str = "<|silence|>".join(["<|image|>"] * len(frames_to_process)) # Append the image token string. text_to_append += image_tokens_str # Get new inputs from `new_prompts` and `new_video_frames` # - If new_prompts is not empty and new_video_frames is empty: # `input_ids`, `attention_mask`, `position_ids` # - If new_prompts is not empty and new_video_frames is not empty: # `input_ids`, `attention_mask`, `position_ids`, # `pixel_values`, `aspect_ratio_ids`, `aspect_ratio_mask`, `vision_position_ids`, `cross_attention_mask` # - If new_prompts is empty and new_video_frames is empty: # empty # Note that: if new_video_frames is not empty, new_prompts must be not empty. current_new_inputs = None if text_to_append or frames_to_process: # if both text_to_append and frames_to_process are empty, processor will raise an error current_new_inputs = processor(text=text_to_append, images=frames_to_process, add_special_tokens=False, return_tensors="pt") # [update use_cache] # [update num_logits_to_keep] # `use_cache` and `num_logits_to_keep` always keep the same value as the previous turn. # [update past_key_values] keeping its naming used in model code cache_name, cache = self._extract_past_from_model_output(outputs) model_kwargs[cache_name] = cache # NOTE: Currently, we do NOT have `state` in the outputs if getattr(outputs, "state", None) is not None: model_kwargs["state"] = outputs.state # NOTE: Currently, we do NOT have `token_type_ids` in the outputs if "token_type_ids" in model_kwargs: token_type_ids = model_kwargs["token_type_ids"] model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1) # get the device of input_ids inputs_device = input_ids.device # NOTE: Currently, we only support batch size 1 for real-time generation. assert input_ids.shape[0] == 1, "Currently, we only support batch size 1 for real-time generation." # [update input_ids] if current_new_inputs is not None and "input_ids" in current_new_inputs: input_ids = torch.cat([input_ids, current_new_inputs["input_ids"].to(inputs_device)], dim=1) # [calculate num_new_tokens] num_new_tokens = 1 # generated_token's + current_new_inputs's if current_new_inputs is not None and "input_ids" in current_new_inputs: num_new_tokens += current_new_inputs["input_ids"].shape[1] # [update attention_mask] if not is_encoder_decoder: if "attention_mask" in model_kwargs: attention_mask = model_kwargs["attention_mask"] if current_new_inputs is not None and "attention_mask" in current_new_inputs: current_new_attention_mask = current_new_inputs["attention_mask"].to(inputs_device) # (previvous, generated_token's, current_new_inputs's) model_kwargs["attention_mask"] = torch.cat( [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1)), current_new_attention_mask], dim=-1 ) else: # (previvous, generated_token's) model_kwargs["attention_mask"] = torch.cat( [attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_tokens))], dim=-1 ) # NOTE: Currently, we do NOT have `decoder_attention_mask` in the outputs else: if "decoder_attention_mask" in model_kwargs: decoder_attention_mask = model_kwargs["decoder_attention_mask"] model_kwargs["decoder_attention_mask"] = torch.cat( [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1))], dim=-1, ) # [update position_ids] if "position_ids" in model_kwargs: # past_position_ids position_ids = model_kwargs["position_ids"] # shape: [batch_size] - find the maximum position_id in each batch last_positions = torch.max(position_ids, dim=1)[0] + 1 # position_ids for generated token (default to 0) generated_token_position_ids = torch.zeros((last_positions.shape[0], 1), dtype=position_ids.dtype, device=inputs_device) generated_token_position_ids = last_positions.unsqueeze(1) + generated_token_position_ids # position_ids for current_new_inputs if current_new_inputs is not None and "position_ids" in current_new_inputs: # get the position_ids of current_new_inputs from current_new_inputs current_new_position_ids = current_new_inputs["position_ids"].to(inputs_device) current_new_position_ids = current_new_position_ids + 1 else: # create base position_ids for current_new_inputs current_new_position_ids = torch.arange(1, num_new_tokens, dtype=position_ids.dtype, device=inputs_device).unsqueeze(0) current_new_position_ids = last_positions.unsqueeze(1) + current_new_position_ids # Apply attention_mask to set masked positions to pad_position_id if current_new_inputs is not None and "attention_mask" in current_new_inputs: attention_mask = current_new_inputs["attention_mask"].to(inputs_device) # Set masked positions (where attention_mask == 0) to pad position id current_new_position_ids = torch.where( attention_mask == 0, VIDEO_MLLAMA_PROCESSOR_PAD_POSITION_ID, current_new_position_ids ) # shape: [batch_size, num_new_tokens] current_new_position_ids = torch.cat([generated_token_position_ids, current_new_position_ids], dim=1) # final position_ids model_kwargs["position_ids"] = torch.cat([position_ids, current_new_position_ids], dim=1) # [update cache_position] # Note that: past_cache_position = model_kwargs.pop("cache_position") new_cache_position = torch.arange( past_cache_position[-1] + 1, past_cache_position[-1] + num_new_tokens + 1, dtype=past_cache_position.dtype ).to(past_cache_position.device) # If use_cache is True, we only need to update the last position if model_kwargs.get("use_cache", True): model_kwargs["cache_position"] = new_cache_position # Otherwise, we need to update the whole cache_position else: model_kwargs["cache_position"] = torch.cat((past_cache_position, new_cache_position)) # new image or video frames are added in this turn if frames_to_process: # [update pixel_values] if current_new_inputs is not None and "pixel_values" in current_new_inputs: model_kwargs["pixel_values"] = current_new_inputs["pixel_values"].to(inputs_device) # [update aspect_ratio_ids] if current_new_inputs is not None and "aspect_ratio_ids" in current_new_inputs: model_kwargs["aspect_ratio_ids"] = current_new_inputs["aspect_ratio_ids"].to(inputs_device) # [update aspect_ratio_mask] if current_new_inputs is not None and "aspect_ratio_mask" in current_new_inputs: model_kwargs["aspect_ratio_mask"] = current_new_inputs["aspect_ratio_mask"].to(inputs_device) # [update vision_position_ids] if "position_ids" in model_kwargs and processor.add_video_position_encoding: # current position_ids current_position_ids = model_kwargs["position_ids"] # calculate the first position of current_new_position_ids first_positions = current_position_ids[:, -num_new_tokens] + 1 # past vision_position_ids vision_position_ids = model_kwargs.pop("vision_position_ids", None) vision_position_ids = vision_position_ids if vision_position_ids is not None and vision_position_ids.numel() > 0 else None if current_new_inputs is not None and "vision_position_ids" in current_new_inputs: current_new_vision_position_ids = current_new_inputs["vision_position_ids"].to(inputs_device) else: current_new_vision_position_ids = torch.arange(0, len(frames_to_process), dtype=position_ids.dtype, device=inputs_device).unsqueeze(0) current_new_vision_position_ids = current_new_vision_position_ids + first_positions.unsqueeze(1) if vision_position_ids is not None: model_kwargs["vision_position_ids"] = torch.cat([vision_position_ids, current_new_vision_position_ids], dim=1) else: model_kwargs["vision_position_ids"] = current_new_vision_position_ids # [update vision_cache_position] # Note that: # The `vision_cache_position` actually retains the complete cache_position_ids for all images. # Therefore, when using it, we need to truncate it according to the actual num_images in `pixel_values`. past_vision_cache_position = model_kwargs.pop("vision_cache_position", None) # calculate the last position of vision_cache_position last_vision_cache_position = past_vision_cache_position[-1] + 1 if past_vision_cache_position is not None else 0 # calculate the new vision_cache_position new_vision_cache_position = torch.arange( last_vision_cache_position, last_vision_cache_position + len(frames_to_process), dtype=position_ids.dtype ).to(inputs_device) if past_vision_cache_position is None: model_kwargs["vision_cache_position"] = new_vision_cache_position else: model_kwargs["vision_cache_position"] = torch.cat((past_vision_cache_position, new_vision_cache_position)) else: model_kwargs["pixel_values"] = None model_kwargs["aspect_ratio_ids"] = None model_kwargs["aspect_ratio_mask"] = None model_kwargs["vision_position_ids"] = model_kwargs.pop("vision_position_ids", None) model_kwargs["vision_cache_position"] = model_kwargs.pop("vision_cache_position", None) # [update cross_attention_mask] # past_cross_attention_mask: (batch_size, past_seq_len, past_max_num_images, max_num_tiles) past_cross_attention_mask = model_kwargs.pop("cross_attention_mask", None) if current_new_inputs is not None and "cross_attention_mask" in current_new_inputs: # shape is (batch_size, cur_new_seq_len, cur_new_max_num_images, max_num_tiles) current_new_cross_attention_mask = current_new_inputs["cross_attention_mask"].to(inputs_device) batch_size, cur_new_seq_len, cur_new_max_num_images, max_num_tiles = current_new_cross_attention_mask.shape # the shape of cross_attention_mask for current_new_inputs is (batch_size, seq_len, max_num_images, max_num_tiles) # seq_len = 1 + cur_new_seq_len # max_num_images = past_max_num_images + cur_new_max_num_images # Note that: # (0) current_generated_token can be seen as the first token of current_new_inputs. # (1) all current new tokens can attend to all past images. # (2) the last token of past tokens also can attention to all past images. # add current_generated_token's cross_attention_mask to the beginning of current_new_cross_attention_mask # (batch_size, 1, cur_new_max_num_images, max_num_tiles) current_generated_token_cross_attention_mask = torch.zeros((batch_size, 1, cur_new_max_num_images, max_num_tiles), dtype=current_new_cross_attention_mask.dtype, device=inputs_device) # (batch_size, 1 + cur_new_seq_len, cur_new_max_num_images, max_num_tiles) current_new_cross_attention_mask = torch.cat([current_generated_token_cross_attention_mask, current_new_cross_attention_mask], dim=1) if past_cross_attention_mask is not None: # (batch_size, past_max_num_images, max_num_tiles) last_past_cross_attention_mask = past_cross_attention_mask[:, -1, :, :] # (batch_size, 1 + cur_new_seq_len, past_max_num_images, max_num_tiles) expanded_past_cross_attention_mask = last_past_cross_attention_mask.unsqueeze(1).repeat(1, 1 + cur_new_seq_len, 1, 1) # (batch_size, 1 + cur_new_seq_len, past_max_num_images + cur_new_max_num_images, max_num_tiles) current_new_cross_attention_mask = torch.cat([expanded_past_cross_attention_mask, current_new_cross_attention_mask], dim=2) else: if past_cross_attention_mask is not None: # (batch_size, 1, past_max_num_images, max_num_tiles) current_generated_token_cross_attention_mask = past_cross_attention_mask[:, -1, :, :].unsqueeze(1) # (batch_size, num_new_tokens, past_max_num_images, max_num_tiles) current_new_cross_attention_mask = current_generated_token_cross_attention_mask.repeat(1, num_new_tokens, 1, 1) else: # `past_cross_attention_mask` is None and `current_new_cross_attention_mask` is None, means that NO image or video frames # so cross_attention_mask is NOT needed current_new_cross_attention_mask = None # update cross_attention_mask model_kwargs["cross_attention_mask"] = current_new_cross_attention_mask return input_ids, model_kwargs @torch.no_grad() def _real_time_generate( self, new_video_frames: queue.Queue, new_prompts: queue.Queue, output_text_queue: queue.Queue, processor, inputs: Optional[torch.Tensor] = None, generation_config: Optional[GenerationConfig] = None, logits_processor: Optional[LogitsProcessorList] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None, synced_gpus: Optional[bool] = None, assistant_model: Optional["PreTrainedModel"] = None, streamer: Optional["BaseStreamer"] = None, negative_prompt_ids: Optional[torch.Tensor] = None, negative_prompt_attention_mask: Optional[torch.Tensor] = None, max_tokens_per_turn: int = 86400, is_in_offline_generate: bool = False, vision_chunked_length: int = 64, **kwargs, ) -> Union[GenerateOutput, torch.LongTensor]: # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call self._validate_model_class() # Pull this out first, we only use it for stopping criteria tokenizer = kwargs.pop("tokenizer", None) # only used for assisted generation assistant_tokenizer = kwargs.pop("assistant_tokenizer", None) # 2. Prepare generation config and model kwargs # Step 2.1: If generation_config is not provided, we use self.generation_config as default # Step 2.2: Use kwargs to update generation_config # Step 2.3: Set kwargs as model_kwargs generation_config, model_kwargs = self._prepare_generation_config(generation_config, **kwargs) # Verify model_kwargs based on the parameters of `self.prepare_inputs_for_generation` and `self.forward`. self._validate_model_kwargs(model_kwargs.copy()) # Only used for assisted generation # NOTE: Currently, we do NOT support assisted generation. self._validate_assistant(assistant_model, tokenizer, assistant_tokenizer) # 3. Set generation parameters if not already defined if synced_gpus is None: synced_gpus = (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and dist.get_world_size() > 1 assert not synced_gpus, "Currently, we do not support synced_gpus for real-time generation." # 4. Prepare logits processor and stopping criteria logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() # 5. Prepare model inputs # Step 5.1.1: Check if the model accepts attention mask accepts_attention_mask = "attention_mask" in set(inspect.signature(self.forward).parameters.keys()) # Step 5.1.2: Check if the model requires attention mask requires_attention_mask = "encoder_outputs" not in model_kwargs # Step 5.1.3: Check if the model has attention mask in model_kwargs kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None # Step 5.2: Prepare model inputs # Step 5.2.1: If input_ids is passed through inputs, return input_ids and the string "input_ids". # Step 5.2.2: If input_ids is passed through model_kwargs["input_ids"], return input_ids and the string "input_ids". # Step 5.2.3: If inputs_embeds is passed through model_kwargs["inputs_embeds"], return inputs_embeds, the string "inputs_embeds", and add an empty input_ids to model_kwargs (specifically, torch.ones((batch_size, 0))). # Step 5.2.4: If neither input_ids nor inputs_embeds is provided, return an input_ids that starts with the bos_token. # Notice: Only when `inputs_tensor` is "inputs_embeds" will "input_ids" be present in model_kwargs. Otherwise, after processing, there should be no "input_ids" in model_kwargs. # Notice: If "inputs_embeds" exist, they will still be retained in model_kwargs. inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs( inputs, generation_config.bos_token_id, model_kwargs ) batch_size = inputs_tensor.shape[0] # NOTE: Currently, we only support batch size 1 for real-time generation. assert batch_size == 1, "Currently, we only support batch size 1 for real-time generation." # Step 6: prepare special tokens # Step 6.1: Prepare generation_config.`special_tokens`: _bos_token_tensor, _eos_token_tensor, _pad_token_tensor, _decoder_start_token_tensor # Notice: `_eos_token_tensor` is 1D tensor, actually a list of eos_token_ids. device = inputs_tensor.device self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=device) # Step 7: padding side check # Step 7.1: decoder-only models must use left-padding for batched generation. if not self.config.is_encoder_decoder and not is_torchdynamo_compiling(): # If `input_ids` was given, check if the last id in any sequence is `pad_token_id` # Note: If using, `inputs_embeds` this check does not work, because we want to be more hands-off. if ( generation_config._pad_token_tensor is not None and batch_size > 1 and len(inputs_tensor.shape) == 2 and torch.sum(inputs_tensor[:, -1] == generation_config._pad_token_tensor) > 0 ): logger.warning( "A decoder-only architecture is being used, but right-padding was detected! For correct " "generation results, please set `padding_side='left'` when initializing the tokenizer." ) # Step 8: Define other model kwargs # Step 8.1: If the model is decoder-only and `inputs_embeds` is passed, set `use_cache` to True. if not self.config.is_encoder_decoder and model_input_name == "inputs_embeds": generation_config.use_cache = True # Step 8.2: Prepare attention mask # Step 8.2.1: If the model does not have attention mask in model_kwargs and requires attention mask, prepare attention mask. if not kwargs_has_attention_mask and requires_attention_mask and accepts_attention_mask: # Only when `pad_token_id` exists and is not equal to any of the `eos_token_ids` can the `attention_mask` be inferred from `input_ids`; # otherwise, the `attention_mask` should default to all ones. model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation( inputs_tensor, generation_config, model_kwargs ) # Step 8.2.2: If the model has attention mask in model_kwargs, check if it is 2D. elif kwargs_has_attention_mask: # TODO (joao): generalize this check with other types of inputs if model_input_name == "input_ids" and len(model_kwargs["attention_mask"].shape) > 2: raise ValueError("`attention_mask` passed to `generate` must be 2D.") # NOTE: Currently, we do not support encoder-decoder models. if self.config.is_encoder_decoder and "encoder_outputs" not in model_kwargs: model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation( inputs_tensor, model_kwargs, model_input_name, generation_config ) # NOTE: Currently, we do not support encoder-decoder models. if self.config.is_encoder_decoder: input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation( batch_size=batch_size, model_input_name=model_input_name, model_kwargs=model_kwargs, decoder_start_token_id=generation_config._decoder_start_token_tensor, device=inputs_tensor.device, ) else: input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids") if streamer is not None: streamer.put(input_ids.cpu()) # 9. Prepare `max_length` and `min_length` input_ids_length = input_ids.shape[-1] has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None generation_config = self._prepare_generated_length( generation_config=generation_config, has_default_max_length=has_default_max_length, has_default_min_length=has_default_min_length, model_input_name=model_input_name, inputs_tensor=inputs_tensor, input_ids_length=input_ids_length, ) self._validate_generated_length(generation_config, input_ids_length, has_default_max_length) # Step 10: prepare num_logits_to_keep # If the model supports `num_logits_to_keep` in forward(), set it to 1 to avoid computing the whole # logit matrix. This can save a lot of memory during the first forward pass. Note that assisted decoding # dynamically overrides this value as it can need more than the last token logits if self._supports_num_logits_to_keep() and "num_logits_to_keep" not in model_kwargs: model_kwargs["num_logits_to_keep"] = 1 # Step 11. Prepare the KV-Cache # Step 11.1: Calculate the maximum cache length max_cache_length = generation_config.max_length if ( inputs_tensor.shape[1] != input_ids_length and model_input_name == "inputs_embeds" and not self.config.is_encoder_decoder ): max_cache_length += inputs_tensor.shape[1] # Step 11.2: Prepare the cache for generation, KV-Cache will be ** added ** to `model_kwargs`. # NOTE: Currently, we only successfully support default cache implementation, which is `DynamicCache()`. self._prepare_cache_for_generation( generation_config, model_kwargs, assistant_model, batch_size, max_cache_length, device ) # Step 12. Determine generation mode generation_mode = generation_config.get_generation_mode(assistant_model=assistant_model) # NOTE: Currently, we only support `GenerationMode.SAMPLE` and `GenerationMode.GREEDY_SEARCH` for real-time generation. if streamer is not None and (generation_config.num_beams > 1): raise ValueError( "`streamer` cannot be used with beam search (yet!). Make sure that `num_beams` is set to 1." ) if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type: warnings.warn( "You are calling .generate() with the `input_ids` being on a device type different" f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model" f" is on {self.device.type}. You may experience unexpected behaviors or slower generation." " Please make sure that you have put `input_ids` to the" f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before" " running `.generate()`.", UserWarning, ) # Step 13. prepare logits processors and stopping criteria # Step 13.1: Prepare logits processor logits_processor = self._get_logits_processor( generation_config=generation_config, input_ids_seq_length=input_ids_length, encoder_input_ids=inputs_tensor, prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, logits_processor=logits_processor, device=inputs_tensor.device, model_kwargs=model_kwargs, negative_prompt_ids=negative_prompt_ids, negative_prompt_attention_mask=negative_prompt_attention_mask, ) # Step 13.2: Prepare stopping criteria stopping_criteria = self._get_stopping_criteria( generation_config=generation_config, stopping_criteria=stopping_criteria, tokenizer=tokenizer, **kwargs ) # Step 14. Set model_kwargs `use_cache` so we can use it later in forward runs model_kwargs["use_cache"] = generation_config.use_cache # Step 15. Go into different generation modes if generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH): # Step 15.1. expand input_ids with `num_return_sequences` additional sequences per batch input_ids, model_kwargs = self._expand_inputs_for_generation( input_ids=input_ids, expand_size=generation_config.num_return_sequences, is_encoder_decoder=self.config.is_encoder_decoder, **model_kwargs, ) # Step 15.2. run sample return self._real_time_sample( input_ids, new_video_frames=new_video_frames, new_prompts=new_prompts, output_text_queue=output_text_queue, processor=processor, logits_processor=logits_processor, stopping_criteria=stopping_criteria, generation_config=generation_config, synced_gpus=synced_gpus, streamer=streamer, max_tokens_per_turn=max_tokens_per_turn, is_in_offline_generate=is_in_offline_generate, vision_chunked_length =vision_chunked_length, **model_kwargs, ) else: raise NotImplementedError("This generation mode is not supported for real-time generation.") def _real_time_sample( self, input_ids: torch.LongTensor, new_video_frames: queue.Queue, new_prompts: queue.Queue, output_text_queue: queue.Queue, processor, logits_processor: LogitsProcessorList, stopping_criteria: StoppingCriteriaList, generation_config: GenerationConfig, synced_gpus: bool, streamer: Optional["BaseStreamer"], max_tokens_per_turn: int = 86400, is_in_offline_generate: bool = False, vision_chunked_length: int = 64, **model_kwargs, ) -> Union[GenerateOutput, torch.LongTensor]: # init values pad_token_id = generation_config._pad_token_tensor output_attentions = generation_config.output_attentions output_hidden_states = generation_config.output_hidden_states output_scores = generation_config.output_scores output_logits = generation_config.output_logits return_dict_in_generate = generation_config.return_dict_in_generate max_length = generation_config.max_length has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) do_sample = generation_config.do_sample # init attention / hidden states / scores tuples scores = () if (return_dict_in_generate and output_scores) else None raw_logits = () if (return_dict_in_generate and output_logits) else None decoder_attentions = () if (return_dict_in_generate and output_attentions) else None cross_attentions = () if (return_dict_in_generate and output_attentions) else None decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None # NOTE: Currently, we do not support encoder-decoder models. # if model is an encoder-decoder, retrieve encoder attention weights and hidden states if return_dict_in_generate and self.config.is_encoder_decoder: encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None encoder_hidden_states = ( model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None ) # keep track of which sequences are already finished batch_size, cur_len = input_ids.shape assert batch_size == 1, "Currently, we only support batch size 1 for real-time generation." this_peer_finished = False unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) # Add cache_position to model_kwargs model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) model_forward = self.__call__ # NOTE: Currently, we do not support static cache. if isinstance(model_kwargs.get("past_key_values"), StaticCache): raise NotImplementedError("StaticCache has not been supported for real-time generation yet.") if self.device.type == "cuda": logger.warning_once("Using `torch.compile`.") os.environ["TOKENIZERS_PARALLELISM"] = "0" model_forward = self.get_compiled_call(generation_config.compile_config) is_prefill = True current_generated_token_start_time = time.time() # Initialize token buffer for handling multi-byte characters token_buffer = deque() silence_token_id = processor.tokenizer.convert_tokens_to_ids(processor.tokenizer.tokenize("<|silence|>"))[0] ellipsis_token_id = processor.tokenizer.convert_tokens_to_ids(processor.tokenizer.tokenize("<|...|>"))[0] invalid_token_id = processor.tokenizer.convert_tokens_to_ids(processor.tokenizer.tokenize('�'))[0] while True: if not self.continue_generating: break if is_in_offline_generate and not self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device, cur_len=cur_len, max_length=max_length): break # prepare model inputs model_inputs = self.prepare_inputs_for_real_time_generation(input_ids, **model_kwargs) # prepare variable output controls (note: some models won't accept all output controls) model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) if is_prefill: if is_in_offline_generate: outputs = self.vision_chunked_forward(**model_inputs, vision_chunked_length=vision_chunked_length, return_dict=True) else: outputs = self(**model_inputs, return_dict=True) is_prefill = False else: outputs = model_forward(**model_inputs, return_dict=True) # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration # (the clone itself is always small) next_token_logits = outputs.logits[:, -1, :].clone().float() next_token_logits = next_token_logits.to(input_ids.device) # DEBUG: Log silence token logits silence_logit_value = next_token_logits[0, silence_token_id].item() logger.info(f"[DEBUG] <|silence|> token logit value: {silence_logit_value:.6f}") probs = nn.functional.softmax(next_token_logits, dim=-1) logger.info(f"[DEBUG] <|silence|> token probability: {probs[0, silence_token_id].item():.6f}") # pre-process distribution next_token_scores = logits_processor(input_ids, next_token_logits) # DEBUG: Log silence token scores after processing silence_score_value = next_token_scores[0, silence_token_id].item() logger.info(f"[DEBUG] <|silence|> token score after processing: {silence_score_value:.6f}") probs = nn.functional.softmax(next_token_scores, dim=-1) logger.info(f"[DEBUG] <|silence|> token probability: {probs[0, silence_token_id].item():.6f}") # Store scores, attentions and hidden_states when required if return_dict_in_generate: if output_scores: scores += (next_token_scores,) if output_logits: raw_logits += (next_token_logits,) if output_attentions: decoder_attentions += ( (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) ) # NOTE: Currently, we do not support encoder-decoder models. if self.config.is_encoder_decoder: cross_attentions += (outputs.cross_attentions,) if output_hidden_states: decoder_hidden_states += ( (outputs.decoder_hidden_states,) if self.config.is_encoder_decoder else (outputs.hidden_states,) ) # Adjust the probabilities of `<|...|>` and `<|silence|>` # -0- original method # probs = nn.functional.softmax(next_token_scores, dim=-1) # -1- Set ellipsis token probability to 0 and renormalize next_token_scores[0, ellipsis_token_id] = float('-inf') probs = nn.functional.softmax(next_token_scores, dim=-1) # -2- Apply silence token probability correction silence_prob = probs[0, silence_token_id].item() silence_threshold = 0.6 if silence_prob < silence_threshold: logger.info(f"[DEBUG] Silence prob {silence_prob:.6f} < {silence_threshold}, setting to 0") # Set silence token probability to 0 and renormalize probs[0, silence_token_id] = 0.0 # Renormalize probabilities to sum to 1 probs = probs / probs.sum(dim=-1, keepdim=True) else: logger.info(f"[DEBUG] Silence prob {silence_prob:.6f} >= {silence_threshold}, keeping original") # token selection if do_sample: # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) else: next_tokens = torch.argmax(probs, dim=-1) # next_tokens = torch.argmax(next_token_scores, dim=-1) if output_text_queue is not None: # Add new token to buffer current_token = next_tokens.item() token_buffer.append(current_token) # Try to decode all tokens in current buffer token_buffer_list = list(token_buffer) decoded_text = processor.tokenizer.decode(token_buffer_list, skip_special_tokens=False, clean_up_tokenization_spaces=False) # Check different conditions for outputting text is_silence_or_ellipsis = current_token in (silence_token_id, ellipsis_token_id) is_complete_text = decoded_text and decoded_text[-1] != '�' is_invalid_token_complete = (decoded_text and decoded_text[-1] == '�' and token_buffer_list[-1] == invalid_token_id) # Output text if any condition is met if is_silence_or_ellipsis or is_complete_text or is_invalid_token_complete: output_text_queue.put(decoded_text) token_buffer.clear() # Otherwise continue accumulating tokens, waiting for more tokens to complete multi-byte characters # Currently, we do use official stopping criteria. # finished sentences should have their next token be a padding token if is_in_offline_generate and has_eos_stopping_criteria: next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) # update generated ids, model inputs, and length for next step input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) if streamer is not None: streamer.put(next_tokens.cpu()) # Currently, we do use official stopping criteria for offline generate. if is_in_offline_generate: unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores) this_peer_finished = unfinished_sequences.max() == 0 cur_len += 1 current_generated_token_cost_time = time.time() - current_generated_token_start_time current_generated_token_wait_time = 1.0 / max_tokens_per_turn - current_generated_token_cost_time if current_generated_token_wait_time > 0: time.sleep(current_generated_token_wait_time) current_generated_token_start_time = time.time() is_in_silence_state = next_tokens.item() == processor.tokenizer.convert_tokens_to_ids("<|silence|>") should_wait_for_new_input = is_in_silence_state if is_in_offline_generate: should_wait_for_new_input = False # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping input_ids, model_kwargs = self._update_model_kwargs_for_real_time_generation( outputs, input_ids, model_kwargs, should_wait_for_new_input=should_wait_for_new_input, is_encoder_decoder=self.config.is_encoder_decoder, new_video_frames=new_video_frames, new_prompts=new_prompts, output_text_queue=output_text_queue, token_buffer=token_buffer, processor=processor, ) if synced_gpus and this_peer_finished: continue # This is needed to properly delete outputs.logits which may be very large for first iteration # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration del outputs if streamer is not None: streamer.end() if return_dict_in_generate: # NOTE: Currently, we do not support encoder-decoder models. if self.config.is_encoder_decoder: return GenerateEncoderDecoderOutput( sequences=input_ids, scores=scores, logits=raw_logits, encoder_attentions=encoder_attentions, encoder_hidden_states=encoder_hidden_states, decoder_attentions=decoder_attentions, cross_attentions=cross_attentions, decoder_hidden_states=decoder_hidden_states, past_key_values=model_kwargs.get("past_key_values"), ) else: return GenerateDecoderOnlyOutput( sequences=input_ids, scores=scores, logits=raw_logits, attentions=decoder_attentions, hidden_states=decoder_hidden_states, past_key_values=model_kwargs.get("past_key_values"), ) else: return input_ids def real_time_generate( self, new_video_frames: queue.Queue, new_prompts: queue.Queue, output_text_queue: queue.Queue, processor, max_tokens_per_turn: int = 86400, **generate_kwargs, ): """ Supports real-time video understanding. This function continuously processes incoming video frames and prompts. New video frames and prompts can be added to the `new_video_frames` and `new_prompts` lists from outside this function. Args: new_video_frames (queue.Queue): A queue to which new video frames (e.g., PIL Images or tensors) are added. new_prompts (queue.Queue): A queue to which new text prompts (strings) are added. output_text_queue (queue.Queue): A queue to which generated text is added. processor: The processor for preparing model inputs. **generate_kwargs: Additional keyword arguments for the model's `generate` method. """ system_prompt = ( "You are a helpful AI assistant. You perceive and understand the surrounding environment in real-time through the camera and interact with the user. Whether or not the user actively asks questions, you continuously observe and analyze visual information, maintaining awareness of the environment.\n\n" "Core Abilities:\n" "- Continuous Perception: Always observe and analyze the visual information captured by the camera in real-time. This perception does not stop even when there is no user interaction.\n" "- Observation-Based Answers: All answers must be based on visual observations, including both historical and real-time data. Do not make guesses without evidence from visual input.\n" "- Dynamic Adjustment: When you observe changes relevant to the user's question, promptly update and adjust your answers to ensure timeliness and accuracy.\n\n" "Interaction Rules:\n" "- Always base your answers on observed visual information.\n" "- If you notice significant changes related to the user's question, proactively and promptly update your answer.\n" "- When you are unable to answer or have finished answering, output `<|silence|>` directly.\n\n" "Your goal is to be the user's reliable \"eyes\", helping them understand and perceive the world around them." ) initial_messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": [ {"type": "text", "text": ""} ]} ] initial_input_text = processor.apply_chat_template(initial_messages, add_generation_prompt=True) initial_inputs = processor(text=initial_input_text, add_special_tokens=False, return_tensors="pt").to(self.device) self.start_real_time_generate() try: return self._real_time_generate( new_video_frames=new_video_frames, new_prompts=new_prompts, output_text_queue=output_text_queue, processor=processor, **initial_inputs, **generate_kwargs, max_tokens_per_turn=max_tokens_per_turn, ) finally: self.stop_real_time_generate()