modeling_video_mllama.py

# 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 <image> 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 <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>.",
                # 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 <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>.",
            }
        }
        
        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()