modeling_hyperclovax.py

import ast
import contextlib
import gc
import json
import math
import os
from dataclasses import dataclass
from functools import partial
from itertools import chain
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.distributed as dist
import torch.nn as nn
from einops import rearrange
from timm.layers import LayerNorm, LayerNorm2d
from timm.models.regnet import RegStage
from torch.nn import CrossEntropyLoss
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    AutoTokenizer,
    PreTrainedModel,
)
from transformers.generation.utils import GenerationMixin
from transformers.modeling_utils import (
    is_fsdp_enabled,
    is_local_dist_rank_0,
    no_init_weights,
)
from transformers.models.auto import CONFIG_MAPPING
from transformers.utils import ModelOutput

from .configuration_hyperclovax import HCXVisionConfig
from .preprocessor import select_best_resolution

EOT = "<|endofturn|>"
IMG_LOC = "<|dummy3|>"


def get_rank():
    if dist.is_initialized():
        return dist.get_rank()
    return 0


def get_world_size():
    if torch.distributed.is_initialized():
        world_size = torch.distributed.get_world_size()
    else:
        world_size = 1
    return world_size


def unpad_image(tensor: torch.Tensor, original_size: Tuple[int, int]) -> torch.Tensor:
    """Unpads a PyTorch tensor of a padded and resized image.

    This function removes padding from a tensor image that was previously padded and resized.
    The padding is removed based on the aspect ratio difference between the original and current image dimensions.

    Args:
        tensor: The image tensor, assumed to be in CxHxW format.
        original_size: The original size of the image as (width, height).

    Returns:
        The unpadded image tensor.

    Examples:
        >>> import torch
        >>> # Example 1: Unpadding with height padding
        >>> padded_tensor = torch.randn(1, 64, 48)  # Padded tensor (C=1, H=64, W=48)
        >>> original_size = (32, 32)  # Original size (width=32, height=32)
        >>> unpadded_tensor = unpad_image(padded_tensor, original_size)
        >>> unpadded_tensor.shape
        torch.Size([1, 48, 48])
        >>> # Example 2: Unpadding with width padding
        >>> padded_tensor = torch.randn(1, 48, 64)  # Padded tensor (C=1, H=48, W=64)
        >>> original_size = (32, 32)  # Original size (width=32, height=32)
        >>> unpadded_tensor = unpad_image(padded_tensor, original_size)
        >>> unpadded_tensor.shape
        torch.Size([1, 48, 48])
    """
    original_width, original_height = original_size
    current_height, current_width = tensor.shape[1:]

    original_aspect_ratio = original_width / original_height
    current_aspect_ratio = current_width / current_height

    if original_aspect_ratio > current_aspect_ratio:
        scale_factor = current_width / original_width
        new_height = int(original_height * scale_factor)
        padding = (current_height - new_height) // 2
        unpadded_tensor = tensor[:, padding : current_height - padding, :]
    else:
        scale_factor = current_height / original_height
        new_width = int(original_width * scale_factor)
        padding = (current_width - new_width) // 2
        unpadded_tensor = tensor[:, :, padding : current_width - padding]

    return unpadded_tensor


def get_anyres_image_grid_shape(
    image_size: Tuple[int, int],
    grid_pinpoints: Union[str, List[Tuple[int, int]]],
    patch_size: int,
) -> Tuple[int, int]:
    """Calculates the image patch grid shape after any-resolution preprocessing.

    Selects the optimal resolution from predefined grid pinpoints based on input image
    dimensions using `select_best_resolution`, then computes the grid layout by
    dividing the selected resolution by the patch size using integer division.

    Args:
        image_size (Tuple[int, int]): Original image dimensions in (width, height) format.
        grid_pinpoints (Union[str, List[Tuple[int, int]]]): Accepts either:
            - List of (height, width) resolution tuples
            - String representation of list (e.g., "[(224, 224), (336, 336)]")
        patch_size (int): Spatial dimension of square patches for grid division.

    Returns:
        Tuple[int, int]: Grid dimensions as (num_patches_width, num_patches_height).

    Examples:
        >>> # Basic case with list input
        >>> get_anyres_image_grid_shape((1000, 800), [(224, 224), (448, 448)], 112)
        (4, 4)

        >>> # Basic case with string input
        >>> get_anyres_image_grid_shape((600, 400), "[(336, 336), (672, 672)]", 112)
        (6, 6)

        >>> # Case where resolution is not perfectly divisible by patch_size
        >>> # select_best_resolution picks (224, 224). 224 // 100 = 2
        >>> get_anyres_image_grid_shape((500, 500), [(224, 224)], 100)
        (2, 2)

        >>> # Different patch size
        >>> # select_best_resolution picks (448, 448). 448 // 224 = 2
        >>> get_anyres_image_grid_shape((1200, 900), [(448, 448), (224, 224)], 224)
        (2, 2)

    Note:
        String-formatted grid_pinpoints are converted via ast.literal_eval. Invalid formats
        may raise syntax exceptions. The actual resolution selection depends on the
        implementation of `select_best_resolution`. The doctests assume
        `select_best_resolution` picks the *first* resolution provided in `grid_pinpoints`.
    """
    possible_resolutions = grid_pinpoints if isinstance(grid_pinpoints, list) else ast.literal_eval(grid_pinpoints)

    original_width, original_height = image_size
    height, width = select_best_resolution((original_height, original_width), possible_resolutions)
    return width // patch_size, height // patch_size


def reshape_and_unpad_image_features(
    image_feature: torch.Tensor,
    height: int,
    width: int,
    image_size: Tuple[int, int],
    possible_resolutions: List[Tuple[int, int]],
    grid_size: int,
    unpad: bool,
    image_newline: torch.Tensor,
) -> torch.Tensor:
    """Reshapes and processes image features with optional unpadding operation.

    Processes input image features by:
    1. Separating base features from spatial features
    2. Reshaping spatial features into a 5D tensor (num_patch_height, num_patch_width, height, width, channels)
    3. Performing either unpadding operation or simple reshaping based on 'unpad' flag
    4. Concatenating processed features with base features

    Args:
        image_feature: Input tensor containing image features with shape
            [1 + num_patches, feature_dim] where the first element is the base feature
        height: Original image height in pixels
        width: Original image width in pixels
        image_size: Target image size as (width, height) tuple
        possible_resolutions: List of possible [height, width] resolutions for multi-scale processing
        grid_size: Grid dimension for patch arrangement
        unpad: Flag to enable unpadding operation
        image_newline: Special token tensor used as separator when unpadding

    Returns:
        torch.Tensor: Processed image features tensor with shape [1 + num_processed_patches, feature_dim]

    Raises:
        AssertionError: If base feature dimension doesn't match height*width
    """
    base_image_feature = image_feature[0]
    image_feature = image_feature[1:]

    assert (
        height * width == base_image_feature.shape[0]
    ), f"height: {height}, width: {width}, base_image_feature.shape[0]: {base_image_feature.shape[0]}"

    num_patch_width, num_patch_height = get_anyres_image_grid_shape(image_size, possible_resolutions, grid_size)
    image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1)

    if unpad:
        image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
        image_feature = image_feature.flatten(1, 2).flatten(2, 3)
        image_feature = unpad_image(image_feature, image_size)
        image_feature = torch.cat(
            (
                image_feature,
                image_newline[:, None, None].expand(*image_feature.shape[:-1], 1).to(image_feature.device),
            ),
            dim=-1,
        )
        image_feature = image_feature.flatten(1, 2).transpose(0, 1)
    else:
        image_feature = image_feature.permute(0, 2, 1, 3, 4).contiguous()
        image_feature = image_feature.flatten(0, 3)
    image_feature = torch.cat((base_image_feature, image_feature), dim=0)

    return image_feature


def anyres_postprocessing(
    image_forward_outs: torch.FloatTensor,
    split_sizes: List[int],
    image_sizes: List[List[int]],
    possible_resolutions: List[Tuple[int, int]],
    is_videos: List[bool],
    patch_size: int,
    grid_size: int,
    image_newline: torch.FloatTensor,
    num_queries_vis_abstractor: int = -1,
    unpad: bool = False,
) -> List[torch.FloatTensor]:
    """Processes 2D visual features into 1D sequences with post-processing steps.

    Performs AnyRes postprocessing by flattening 2D visual features from grid partitions into 1D sequences, adding
    newline embeddings at row boundaries for images, and optionally removing padding regions based on original image
    sizes. For video data, processes each frame's features separately into a single sequence per video and disables
    unpadding and newline insertion.

    Args:
        image_forward_outs (List[torch.FloatTensor]): List of input tensors with shape
            (number_of_images_in_grid, total_patches, feature_dim) containing visual features.
        split_sizes (List[int]): A list containing the number of patches for each sample in the batch. The sum of
            `split_sizes` should equal `image_forward_outs.shape[0]`.
        image_sizes (List[List[int]]): A list where each element is a list `[width, height]` representing the original
            dimensions of the corresponding image sample. Used for unpadding.
        possible_resolutions (List[Tuple[int, int]]): A list of supported resolution tuples `(height, width)` used by
            `reshape_and_unpad_image_features` for spatial reconstruction, especially during unpadding.
        is_videos (List[bool]): A list of boolean flags indicating whether each corresponding sample in the batch is a
            video [`True`] or an image [`False`].
        patch_size (int): The spatial dimension (height and width) of the square patches the image was divided into.
        grid_size (int): The spatial dimension (height and width) of the square grid onto which patches are mapped.
            `grid_size` should be divisible by `patch_size`.
        image_newline (torch.FloatTensor): A learnable tensor representing the newline embedding, typically with shape
            (1, feature_dim). Added after each row of image patches when not unpadding.
        num_queries_vis_abstractor (int, optional): If a visual abstractor with a fixed number of output queries is used
            instead of grid patching, this specifies the number of queries. Must be a perfect square if > 0.
            Defaults to -1 (indicating standard grid patching is used).
        unpad (bool, optional): If `True`, removes padding tokens from image features based on `image_sizes` and
            `possible_resolutions`. Does not apply to video features. Defaults to False.

    Returns:
        List[torch.FloatTensor]: A list of tensors, where each tensor represents the processed 1D sequence of visual
            features for a single sample from the input batch. The length of the sequence varies depending on processing
            (unpadding, newlines, video flattening).

    Raises:
        AssertionError: If `num_queries_vis_abstractor` is greater than 0 but not a perfect square.
    """
    height = width = grid_size // patch_size

    if num_queries_vis_abstractor > 0:
        assert (num_queries_vis_abstractor**0.5).is_integer(), "n_queries must be square number"
        height = width = int(num_queries_vis_abstractor**0.5)

    image_features = torch.split(image_forward_outs, split_sizes, dim=0)

    # post-processing (unpad, add newline)
    new_image_features = []
    for image_idx, (image_feature, is_video) in enumerate(zip(image_features, is_videos)):
        if image_feature.shape[0] > 1:
            if not is_video:
                image_feature = reshape_and_unpad_image_features(
                    image_feature=image_feature,
                    height=height,
                    width=width,
                    image_size=image_sizes[image_idx],
                    possible_resolutions=possible_resolutions,
                    grid_size=grid_size,  # Pass grid info if needed by helper
                    unpad=unpad,
                    image_newline=image_newline,
                )
            else:
                image_feature = image_feature.flatten(0, 1)
        else:
            image_feature = image_feature[0]
            if unpad and not is_video:
                image_feature = torch.cat((image_feature, image_newline[None].to(image_feature.device)), dim=0)
        new_image_features.append(image_feature)
    image_features = new_image_features
    return image_features


def adaptive_anyres_postprocessing(
    image_forward_outs: torch.FloatTensor,
    image_sizes: List[List[int]],
    possible_resolutions: List[Tuple[int, int]],
    is_videos: List[bool],
    group_ids: List[List[int]],
    num_queries_vis_abstractors: List[List[int]],
    grid_size: int,
    image_newline: torch.FloatTensor,
    unpad: bool = False,
) -> List[torch.FloatTensor]:
    """Adaptive AnyRes postprocessing for multi-group feature aggregation.

    Processes 2D visual features into 1D sequences with group-wise adaptive processing. Each image can belong to
    multiple processing groups with different query configurations. Features are processed per group and aggregated
    according to group_ids.

    Args:
        image_forward_outs (List[torch.FloatTensor]): List of input tensors with shape
            (number_of_images_in_grid, total_patches, feature_dim) containing visual features.
        image_sizes (List[List[int]]): Original image dimensions for each sample. [[width, height], ... ]
        possible_resolutions (List[Tuple[int, int]]): Supported resolutions. [[height, width], ... ]
        is_videos (List[bool]): Flags indicating video inputs
        group_ids (List[List[int]]): Group indices for feature aggregation. Each group means a single grid.
        num_queries_vis_abstractors (List[List[int]]): Query numbers per group
        grid_size (int): Total grid size for spatial processing
        image_newline (torch.FloatTensor): Sample-wise config. Newline embedding tensor
        unpad (bool, optional): Sample-wise config. Enable padding removal. Defaults to False.

    Returns:
        List[torch.FloatTensor]: Aggregated features per group

    Raises:
        AssertionError: If num_queries is not square number in any group
    """
    # post-processing (unpad, add newline)
    new_image_features = []
    for image_idx, (image_feature, is_video) in enumerate(zip(image_forward_outs, is_videos)):
        num_queries_vis_abstractor = num_queries_vis_abstractors[image_idx]
        assert (num_queries_vis_abstractor**0.5).is_integer(), "n_queries must be square number"
        height = width = int(num_queries_vis_abstractor**0.5)

        if image_feature.shape[0] > 1:
            if not is_video:
                image_feature = reshape_and_unpad_image_features(
                    image_feature=image_feature,
                    height=height,
                    width=width,
                    image_size=image_sizes[image_idx],
                    possible_resolutions=possible_resolutions,
                    grid_size=grid_size,
                    unpad=unpad,
                    image_newline=image_newline,
                )
            else:
                image_feature = image_feature.flatten(0, 1)
        else:
            image_feature = image_feature[0]
            if unpad and not is_video:
                image_feature = torch.cat((image_feature, image_newline[None].to(image_feature.device)), dim=0)
        new_image_features.append(image_feature)

    image_features = [
        torch.cat([new_image_features[group_id] for group_id in group_ids_list], dim=0) for group_ids_list in group_ids
    ]
    return image_features


@dataclass
class HCXVisionOutput(ModelOutput):
    """Output class for vision models, containing various computation results.

    Args:
        loss (Optional[torch.FloatTensor], optional): Total cross-entropy loss calculated from logits and labels.
        loss_per_sample (Optional[torch.FloatTensor], optional): Per-sample loss values for advanced loss processing.
        logits (torch.FloatTensor): Classification scores (before SoftMax) of shape (batch_size, num_classes).
        past_key_values (Optional[Tuple[Tuple[torch.FloatTensor]]], optional): Contains precomputed hidden-states
            that can be used (see `past_key_values` input) to speed up sequential decoding.
        hidden_states (Optional[Tuple[torch.FloatTensor]], optional):
            Tuple of torch.FloatTensor (one for the output of the embeddings + one for the output of each layer) of
            shape (batch_size, sequence_length, hidden_size).
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (Optional[Tuple[torch.FloatTensor]], optional): Tuple of torch.FloatTensor (one for each layer)
            of shape (batch_size, num_heads, sequence_length, sequence_length). Attentions weights after the attention
            softmax, used to compute the weighted average in the self-attention heads.
    """

    loss: Optional[torch.FloatTensor] = None
    loss_per_sample: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


class HCXVisionForCausalLM(PreTrainedModel, GenerationMixin):
    """HCX Vision model for causal language modeling with vision-language capabilities.

    This class combines a vision model with a language model to create a multimodal model
    capable of processing images or videos and generating text based on the visual inputs.

    Attributes:
        config_class: Configuration class for the model.
        vision_model_name: Name of the vision model component.
        _no_split_modules: List of modules that should not be split during parallel processing.
        supports_gradient_checkpointing: Whether the model supports gradient checkpointing.
        _skip_keys_device_placement: Keys to skip during device placement.
    """

    config_class = HCXVisionConfig
    vision_model_name = "vision_model"
    _no_split_modules = ["CLIPAttention", "SiglipVisionModel"]
    supports_gradient_checkpointing = True
    _skip_keys_device_placement = "past_key_values"

    def __init__(
        self,
        config: HCXVisionConfig,
        **kwargs: Optional[Any],
    ) -> None:
        """Initialize the HCXVisionForCausalLM model.

        Args:
            config: Configuration object for the model containing parameters for both
                vision and language components.
            **kwargs: Additional keyword arguments:
                - use_liger: Whether to use liger kernel for hyperclovax models.
                - use_fused_ce: Whether to use fused cross-entropy loss.
                - use_sum_loss: Whether to use sum reduction for loss instead of mean.
                - is_safetensor_save: Whether to save model using safetensors format.

        Raises:
            ValueError: If vision_config is not defined or if language_config is not defined.
        """
        super().__init__(config)

        self.flag_changed_max_position_embeddings = False

        vision_model_type = config.vision_config["model_type"]
        if vision_model_type in CONFIG_MAPPING:
            vision_config = CONFIG_MAPPING[vision_model_type](**config.vision_config)
            vision_config.auto_map = {}
        else:
            if config.vision_model_name_or_path is not None:
                vision_config = AutoConfig.from_pretrained(config.vision_model_name_or_path, trust_remote_code=True)
            elif config.vision_config["_name_or_path"] is not None:
                vision_config = AutoConfig.from_pretrained(
                    config.vision_config["_name_or_path"], trust_remote_code=True
                )
            else:
                raise ValueError("vision_config is not defined")

        self.use_liger = kwargs.pop("use_liger", False)
        self.use_fused_ce = kwargs.pop("use_fused_ce", False)
        self.reduction = "sum" if kwargs.pop("use_sum_loss", False) else "mean"

        self.vision_config = vision_config
        vision_config.anyres = config.anyres
        vision_config.max_num_grids = config.max_num_grids

        possible_resolutions = []
        if config.anyres:
            assert config.max_num_grids > 0
            possible_resolutions = [
                [i, j]
                for i in range(1, config.max_num_grids + 1)
                for j in range(1, (config.max_num_grids // i) + 1)
                if not (i == 1 and j == 1 and not config.use_1x1_grid)
            ]

            possible_resolutions = [
                [ys * vision_config.image_size, xs * vision_config.image_size] for ys, xs in possible_resolutions
            ]

        self.possible_resolutions = possible_resolutions

        with no_init_weights():
            self.vision_model = AutoModel.from_config(
                vision_config, trust_remote_code=True
            )  # weight will be loaded in from_pretrained

        assert config.language_config["model_type"] == "llama"
        language_config = CONFIG_MAPPING["llama"](**config.language_config)
        language_config._attn_implementation = kwargs.get("attn_implementation", "sdpa")  # activate flash attention
        language_config.logits_scaling = 1.0

        self.language_config = language_config
        self.language_model = AutoModelForCausalLM.from_config(language_config)

        self.language_model.gradient_checkpointing_enable()
        self.num_queries_vis_abstractor = config.num_queries_vis_abstractor

        # mm_projctor(==connector); vision_model_hidden_size -> LLM embedding size
        input_hidden_size = vision_config.hidden_size
        self.mm_projector = HCXVisionCAbstractor(
            num_queries=self.num_queries_vis_abstractor,
            num_input_tokens=(self.vision_config.image_size // self.vision_config.patch_size) ** 2,
            encoder_hidden_size=input_hidden_size,
            hidden_size=input_hidden_size,
            output_hidden_size=language_config.hidden_size,
            pos_emb=config.proj_pos_emb,
            prenorm=config.proj_prenorm,
        )
        self.use_nth_layer = config.use_nth_layer
        self.config.update({"vision_config": self.vision_model.config.to_dict()})
        self.config.update({"language_config": self.language_model.config.to_dict()})
        self.lm_head_vocab_size = (
            language_config.padded_vocab_size
            if hasattr(language_config, "padded_vocab_size")
            else language_config.vocab_size
        )
        self.language_model.lm_head = nn.Linear(language_config.hidden_size, self.lm_head_vocab_size, bias=False)
        self.model_parallel = False
        self.device_map = None
        self.use_no_grad = None
        self.decoder_max_length = config.decoder_max_length

        self.anyres = config.anyres
        self.unpad = config.unpad
        if self.anyres:
            self.image_newline = nn.Parameter(torch.empty(language_config.hidden_size, dtype=self.dtype))

        self.is_safetensor_save = kwargs.get("is_safetensor_save", True)
        self._backward_compatibility_gradient_checkpointing()

    def _init_weights(self, module):
        # copies from https://github.com/kakaobrain/honeybee/blob/main/honeybee/common_layers.py#L55
        if (
            isinstance(module, nn.Conv2d)  # noqa: SIM101
            or isinstance(module, nn.Embedding)
            or isinstance(module, nn.Linear)
        ):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if hasattr(module, "bias") and module.bias is not None:
                module.bias.data.zero_()

        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        elif isinstance(module, nn.Parameter):
            embed_std = 1 / torch.sqrt(torch.tensor(module.size(0), dtype=torch.float)).to(module.dtype)
            module.data.normal_(mean=0.0, std=embed_std)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[List[List[torch.FloatTensor]]] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        attention_mask: Optional[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,
        image_sizes: Optional[List[List[List[int]]]] = None,
        vision_query_lengths: Optional[List[List[int]]] = None,
        non_vision_query_lengths: Optional[List[int]] = None,
        img_start_ids_list: Optional[List[List[int]]] = None,
        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
        num_queries_vis_abstractors_slow: Optional[List[List[int]]] = None,
        first_last_frames_slows: Optional[List[bool]] = None,
        is_video_list: Optional[List[bool]] = None,
        **kwargs,
    ) -> Union[Tuple, HCXVisionOutput]:
        """Forward pass of the model.

        This method processes the input tokens and images, combines them into a unified
        representation, and generates text output based on the inputs.

        Args:
            input_ids: Input token IDs. In positions where images are inputted, the value is replaced by "<|dummy3|>"
            pixel_values: List of lists of 4D tensors for images. Each outer list corresponds to a batch and contains
                inner lists of image tensors.
            past_key_values: Pre-computed key and value states of the attention layers for faster inference.
            attention_mask: Mask to avoid performing attention on padding token indices.
            inputs_embeds: Input embeddings. If provided, input_ids will not be used.
            labels: Labels for computing the language modeling loss.
            use_cache: Whether to use past key/values for faster inference.
            output_attentions: Whether to return attention weights of each layer.
            output_hidden_states: Whether to return hidden states of each layer.
            return_dict: Whether to return a ModelOutput instead of a tuple.
            image_sizes: List of lists representing image dimensions (width, height).
            vision_query_lengths: List of lists containing lengths when each image is converted into visual tokens.
            non_vision_query_lengths: List of lengths of text tokens (excluding visual tokens) for each sample.
            img_start_ids_list: List of lists containing indices of img_start_id tokens for each sample.
            num_queries_vis_abstractors: List of lists containing number of visual tokens for each image grid.\
                For video frames, this is the number of visual tokens for the fast part.
            num_queries_vis_abstractors_slow: List of lists containing number of visual tokens for
                the slow part when applying the slowfast algorithm to video frames.
            first_last_frames_slows: List of booleans indicating whether the slowfast algorithm is
                applied to the first or last frames of the video.
            is_video_list: List of booleans indicating which inputs are videos.
            **kwargs: Additional keyword arguments.

        Returns:
            If return_dict=True, returns an HCXVisionOutput object containing:
                - loss: Language modeling loss if labels are provided, otherwise None.
                - loss_per_sample: Per-sample loss if labels are provided, otherwise None.
                - logits: Prediction scores of the language modeling head.
                - past_key_values: Past key/values for faster inference if use_cache=True.
                - hidden_states: Hidden states of all layers if output_hidden_states=True.
                - attentions: Attention weights of all layers if output_attentions=True.
            If return_dict=False, returns a tuple containing the above items except loss_per_sample.
        """
        output_attentions = (
            output_attentions if output_attentions is not None else self.config.vision_config["output_attentions"]
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.vision_config["output_hidden_states"]
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if inputs_embeds is None and past_key_values is None:
            inputs_embeds = self.extract_inputs_embeds(
                input_ids=input_ids,
                pixel_values=pixel_values,
                past_key_values=past_key_values,
                image_sizes=image_sizes,
                vision_query_lengths=vision_query_lengths,
                non_vision_query_lengths=non_vision_query_lengths,
                img_start_ids_list=img_start_ids_list,
                num_queries_vis_abstractors=num_queries_vis_abstractors,
                num_queries_vis_abstractors_slow=num_queries_vis_abstractors_slow,
                first_last_frames_slows=first_last_frames_slows,
                is_videos=is_video_list,
            )

        if inputs_embeds is not None:
            input_ids = None

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.language_model.base_model(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        hidden_states = hidden_states * self.language_config.logits_scaling

        loss = None
        loss_per_sample = None
        logits = self.language_model.lm_head(hidden_states)
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss(reduction="none")  # ignore IGNORE_INDEX(-100)
            shift_logits = shift_logits.view(-1, self.lm_head_vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model/pipeline parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)
            if get_rank() == 0:
                loss_per_sample = loss.view(logits.shape[0], -1).sum(axis=1) / (
                    shift_labels.view(logits.shape[0], -1) != self.config.ignore_index
                ).sum(axis=1)
            loss = loss[shift_labels != self.config.ignore_index].mean()
        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return HCXVisionOutput(
            loss=loss,
            loss_per_sample=loss_per_sample,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def determine_non_vision_query_lengths(
        self, input_ids: torch.LongTensor, pad_id: int, img_start_id: int
    ) -> List[int]:
        """Calculate the lengths of non-vision query parts in the input.

        This method calculates the length of text tokens (excluding visual tokens) for each sample.
        When input_ids are collated, they are padded with pad_id on the right, so this method finds
        these values by identifying pad tokens and img_start_id tokens.

        Args:
            input_ids: Input token IDs with img_start_id markers for image positions.
            pad_id: Token ID used for padding.
            img_start_id: Token ID marking the start of image data.

        Returns:
            List of lengths of non-vision query parts for each sample in the batch.
        """
        non_vision_query_lengths = []
        batch_size, len_seq = input_ids.size(0), input_ids.size(1)

        for i in range(batch_size):
            temp_idx = (input_ids[i] == pad_id).nonzero()
            eos_idx = temp_idx[0, 0].item() if len(temp_idx) > 0 else len_seq
            num_imgs = (input_ids[i] == img_start_id).sum().item()
            non_vision_query_lengths.append(eos_idx - num_imgs)

        if all([pad_id in input_id for input_id in input_ids.tolist()]):
            non_vision_query_lengths = [
                non_vision_query_length + 1 for non_vision_query_length in non_vision_query_lengths
            ]

        return non_vision_query_lengths

    def determine_vision_query_lengths(
        self, image_features: List[List[torch.Tensor]], image_cnts: List[int]
    ) -> List[List[int]]:
        """Calculate the lengths of vision query parts in the input.

        This method calculates the lengths of visual tokens for each image in each sample based on
        the shapes of image feature tensors. For samples without any images, a dummy image is included
        but then converted to an empty list.

        Args:
            image_features: List of lists of image features tensors.
            image_cnts: List of counts of images for each sample in the batch.

        Returns:
            List of lists of lengths of visual tokens for each image in each sample.
        """
        vision_query_lengths = [
            [image_feature.size(0) for image_feature in image_feature_list] for image_feature_list in image_features
        ]

        for i, image_cnt in enumerate(image_cnts):
            if image_cnt == 0:
                assert len(vision_query_lengths[i]) == 1  # 현재 검정 이미지 1개 들어가있음
                vision_query_lengths[i] = []  # 빈 list 로 변환

        return vision_query_lengths

    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_input_embeddings
    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_input_embeddings
    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_output_embeddings
    def get_output_embeddings(self):
        return self.language_model.get_output_embeddings()

    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_output_embeddings
    def set_output_embeddings(self, new_embeddings):
        self.language_model.set_output_embeddings(new_embeddings)

    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.set_decoder
    def set_decoder(self, decoder):
        self.language_model.set_decoder(decoder)

    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.get_decoder
    def get_decoder(self):
        return self.language_model.get_decoder()

    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.tie_weights
    def tie_weights(self):
        return self.language_model.tie_weights()

    # Copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration.resize_token_embeddings
    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, pad_to_multiple_of=None) -> nn.Embedding:
        model_embeds = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
        self.config.text_config.vocab_size = model_embeds.num_embeddings
        self.vocab_size = model_embeds.num_embeddings
        return model_embeds

    def extract_inputs_embeds(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[List[List[torch.FloatTensor]]] = None,  # list of list of 4D tensors
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        image_sizes: Optional[List[List[List[int]]]] = None,
        vision_query_lengths: Optional[List[List[int]]] = None,
        non_vision_query_lengths: Optional[List[int]] = None,
        img_start_ids_list: Optional[List[List[int]]] = None,
        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
        num_queries_vis_abstractors_slow: Optional[List[List[int]]] = None,
        first_last_frames_slows: Optional[List[bool]] = None,
        is_videos: Optional[List[str]] = None,
    ):
        """Extract input embeddings by processing text tokens and visual features.

        This method processes the input tokens and image features, extracts the visual features
        using the vision model, and combines them with the text token embeddings to create
        a unified input representation for the language model.

        Args:
            input_ids: Input token IDs with img_start_id markers for image positions.
            pixel_values: List of lists of image tensors.
            past_key_values: Pre-computed key and value states for faster inference.
            image_sizes: List of lists of image dimensions (width, height).
            vision_query_lengths: List of lists of lengths when each image is converted to visual tokens.
            non_vision_query_lengths: List of lengths of text tokens (excluding visual tokens) for each sample.
            img_start_ids_list: List of lists containing indices of img_start_id tokens for each sample.
            num_queries_vis_abstractors: List of lists containing number of visual tokens for each image grid.
            num_queries_vis_abstractors_slow: List of lists containing number of visual tokens for
                the slow part when applying the slowfast algorithm to video frames.
            first_last_frames_slows: List of booleans indicating whether the slowfast algorithm is
                applied to the first or last frames of the video.
            is_videos: List of booleans indicating which inputs are videos.

        Returns:
            Combined embeddings of text tokens and visual features.
        """
        inputs_embeds = None
        if past_key_values:
            pass
        else:
            # Flatten CLIP and connector for feature encoding, then convert back to List of List format
            len_pixel_values = [len(pixel_value) for pixel_value in pixel_values]
            concat_pixel_values = torch.cat(list(chain(*pixel_values)), dim=0)  # list of list of 4D Tensor
            visual_token_idx = 0 if "siglip" in self.vision_config.model_type else 1
            # Check if all parameters of the model require_grad=False
            if self.use_no_grad is None:
                self.use_no_grad = all(not p.requires_grad for p in self.vision_model.vision_model.encoder.parameters())
            context = torch.no_grad() if self.use_no_grad else contextlib.nullcontext()
            with context:
                if self.use_no_grad:
                    # Fixed number of for-loop iterations to 10.
                    # Currently no memory effect observed, so proceeding without chunking.
                    n_chunks = 1
                else:
                    n_chunks = 1
                total_len = concat_pixel_values.size(0)
                # Calculate the size of each chunk based on total data length (divided into 10 chunks)
                chunk_size = math.ceil(total_len / n_chunks) if total_len > 0 else 1
                image_forward_outs_chunks = []

                for i in range(n_chunks):
                    start = i * chunk_size
                    end = (i + 1) * chunk_size
                    # Current chunk slice (could be an empty tensor if there's no data)
                    chunk = concat_pixel_values[start:end].to(self.vision_model.dtype)
                    # If the current chunk size is smaller than chunk_size, pad with dummy data
                    if chunk.size(0) < chunk_size:
                        # print(f"chunk.size(0): {chunk.size(0)}, chunk_size: {chunk_size}")
                        pad_size = chunk_size - chunk.size(0)
                        # Create dummy tensor based on concat_pixel_values shape
                        dummy_shape = (pad_size,) + tuple(concat_pixel_values.shape[1:])
                        dummy = torch.zeros(
                            dummy_shape,
                            dtype=concat_pixel_values.dtype,
                            device=concat_pixel_values.device,
                        )
                        chunk = torch.cat([chunk, dummy], dim=0)

                    # Pass the chunk through the vision model (processed according to use_nth_layer)
                    if self.use_nth_layer == -1:
                        # Replace post_layernorm of the last layer with Identity
                        self.vision_model.vision_model.post_layernorm = nn.Identity()
                        outs = self.vision_model(chunk)
                        outs = outs.last_hidden_state[:, visual_token_idx:]
                    else:
                        outs = self.vision_model(chunk, output_hidden_states=True)
                        outs = outs.hidden_states[self.use_nth_layer][:, visual_token_idx:]
                    image_forward_outs_chunks.append(outs)

                # Concatenate results from all chunks
                image_forward_outs = torch.cat(image_forward_outs_chunks, dim=0).to(image_forward_outs_chunks[0].dtype)

            if num_queries_vis_abstractors is None:
                assert num_queries_vis_abstractors_slow is None
                image_sizes = list(chain(*image_sizes))
                if is_videos is not None:
                    is_videos = list(chain(*is_videos))
                group_ids = None
                image_forward_outs = image_forward_outs.to(dtype=self.mm_projector.dtype)
                image_forward_outs = self.mm_projector(image_forward_outs)
            else:
                # adaptive anyres is only implemented in HCXVisionCAbstractor
                assert isinstance(self.mm_projector, HCXVisionCAbstractor)

                (
                    num_queries_vis_abstractors,
                    num_grids,
                    image_sizes,
                    is_videos,
                    group_ids,
                ) = self.compute_adaptive_params(
                    pixel_values,
                    num_queries_vis_abstractors,
                    num_queries_vis_abstractors_slow,
                    image_sizes,
                    is_videos,
                    first_last_frames_slows,
                )

                image_forward_outs = image_forward_outs.to(dtype=self.mm_projector.dtype)
                image_forward_outs = self.mm_projector(
                    image_forward_outs,
                    num_queries_vis_abstractors=num_queries_vis_abstractors,
                    num_grids=num_grids,
                )

            if self.anyres:
                split_sizes = [pixel_value.shape[0] for pixel_value in chain(*pixel_values)]

                if num_queries_vis_abstractors is None:
                    image_features = anyres_postprocessing(
                        image_forward_outs=image_forward_outs,
                        split_sizes=split_sizes,
                        image_sizes=image_sizes,
                        num_queries_vis_abstractor=self.num_queries_vis_abstractor,
                        unpad=self.unpad,
                        is_videos=is_videos,
                        patch_size=self.vision_model.config.patch_size,
                        grid_size=self.vision_model.config.image_size,
                        image_newline=self.image_newline,
                        possible_resolutions=self.possible_resolutions,
                    )
                else:
                    image_features = adaptive_anyres_postprocessing(
                        image_forward_outs=image_forward_outs,
                        image_sizes=image_sizes,
                        num_queries_vis_abstractors=num_queries_vis_abstractors,
                        unpad=self.unpad,
                        is_videos=is_videos,
                        grid_size=self.vision_model.config.image_size,
                        image_newline=self.image_newline,
                        possible_resolutions=self.possible_resolutions,
                        group_ids=group_ids,
                    )
            else:
                if num_queries_vis_abstractors is None:
                    image_features = [image_forward_out for image_forward_out in image_forward_outs]
                else:
                    image_features = [image_forward_out.unsqueeze(0) for image_forward_out in image_forward_outs]

            # print(f"BEFORE GROUPING: len(image_features): {len(image_features)}")
            image_features = [
                image_features[sum(len_pixel_values[:i]) : sum(len_pixel_values[: i + 1])]
                for i in range(len(len_pixel_values))
            ]

            batch_size = input_ids.size(0)
            image_feature_dim = image_features[0][0].size(1)
            image_feature_dtype = image_features[0][0].dtype

            if img_start_ids_list is None:
                image_cnts = (input_ids == self.config.img_start_id).sum(dim=1).tolist()
            else:
                image_cnts = [len(img_start_ids) for img_start_ids in img_start_ids_list]

            if non_vision_query_lengths is None:
                non_vision_query_lengths = self.determine_non_vision_query_lengths(
                    input_ids, self.tokenizer.pad_token_id, self.config.img_start_id
                )

            if vision_query_lengths is None:
                vision_query_lengths = self.determine_vision_query_lengths(image_features, image_cnts)

            # Slicing is faster than concatenation
            len_inputs_embeds = max(
                [
                    sum(vision_query_length) + non_vision_query_length
                    for non_vision_query_length, vision_query_length in zip(
                        non_vision_query_lengths, vision_query_lengths
                    )
                ]
            )
            len_inputs_embeds = min(self.decoder_max_length, len_inputs_embeds)

            inputs_embeds = torch.zeros(
                [batch_size, len_inputs_embeds, image_feature_dim],
                dtype=image_feature_dtype,
                device=self.device,
                requires_grad=True,
            ).clone()
            # temp_embeds : torch.bfloat16 : [batchsize, 174, 3072]
            temp_embeds = self.get_input_embeddings()(input_ids)

            # The complete format is <PROMPT><USER_PREFIX><VISION_QUERIES>Sentence
            for batch_idx, sample in enumerate(input_ids):
                # Concatenate with visual tokens and then slice
                non_vision_query_length = non_vision_query_lengths[batch_idx]
                # Safely concatenate with visual tokens and then slice
                sample = sample[: non_vision_query_length + image_cnts[batch_idx]]

                if image_cnts[batch_idx] == 0:  # Text instruction data doesn't insert image features
                    temp_idx = 0
                    # Reference: https://github.com/haotian-liu/LLaVA/commit/44e0562f9497fb79f042427307472a87d266d90a#diff-4477387d506ccb1897a13972cba26c9da3fad4d3e1c32ec4b8bd8ff7acd3f292
                    # https://github.com/intel/intel-extension-for-transformers/issues/1201#issuecomment-1915875119
                    inputs_embeds[batch_idx, :non_vision_query_length] = temp_embeds[batch_idx][
                        :non_vision_query_length
                    ]
                    inputs_embeds[batch_idx, temp_idx:temp_idx] = image_features[batch_idx][0][
                        0:0
                    ]  # First image of batch_idx sample (dummy image)
                else:
                    if img_start_ids_list is None:
                        img_start_ids = (sample == self.config.img_start_id).nonzero()
                    else:
                        img_start_ids = img_start_ids_list[batch_idx]
                    assert len(img_start_ids) == image_cnts[batch_idx] == len(image_features[batch_idx])
                    # Initialize starting points for input embeddings and temporary embeddings
                    input_start, temp_start = 0, 0

                    # Iterate through each image starting point in the batch
                    for multi_img_idx, img_start_idx in enumerate(img_start_ids):
                        # Calculate token length up to the current image starting point
                        token_len = img_start_idx - temp_start

                        # Copy tokens to inputs_embeds
                        inputs_embeds[batch_idx, input_start : input_start + token_len] = temp_embeds[
                            batch_idx, temp_start : temp_start + token_len
                        ]

                        inputs_embeds[
                            batch_idx,
                            input_start
                            + token_len : input_start
                            + token_len
                            + vision_query_lengths[batch_idx][multi_img_idx],
                        ] = image_features[batch_idx][multi_img_idx]

                        # Update starting points for next token processing
                        input_start += token_len + vision_query_lengths[batch_idx][multi_img_idx]
                        temp_start += token_len + 1  # Increase by 1 to skip the image start token

                    # Process tokens after the last image end token
                    token_len = min(sample[temp_start:].size(0), inputs_embeds.size(1) - input_start)
                    inputs_embeds[batch_idx, input_start : input_start + token_len] = temp_embeds[
                        batch_idx, temp_start : temp_start + token_len
                    ]
        return inputs_embeds

    @torch.no_grad()
    def generate(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[List[List[torch.FloatTensor]]] = None,
        image_sizes: Optional[List[List[List[int]]]] = None,
        vision_query_lengths: Optional[List[List[int]]] = None,
        non_vision_query_lengths: Optional[List[int]] = None,
        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
        num_queries_vis_abstractors_slow: Optional[List[List[int]]] = None,
        first_last_frames_slows: Optional[List[bool]] = None,
        is_videos: Optional[List[bool]] = None,
        img_start_ids_list: Optional[List[List[int]]] = None,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[int] = None,
        bad_words_ids: Optional[List[List[int]]] = None,
        max_length: int = 196,
        min_length: int = 2,
        do_sample: bool = True,
        num_beams: int = 1,
        top_p: float = 0.6,
        top_k: int = 0,
        temperature: float = 0.5,
        repetition_penalty: float = 1.0,
        length_penalty: int = 1,
        use_cache: bool = True,
        **kwargs,
    ) -> torch.LongTensor:
        """Generate text based on input tokens and images.

        This method generates text based on the provided input tokens and images using
        beam search and/or sampling strategies.

        Args:
            input_ids: Input token IDs with img_start_id markers for image positions.
            pixel_values: List of lists of image tensors.
            image_sizes: List of lists of image dimensions (width, height).
            vision_query_lengths: List of lists of lengths when each image is converted to visual tokens.
            non_vision_query_lengths: List of lengths of text tokens (excluding visual tokens) for each sample.
            num_queries_vis_abstractors: List of lists containing number of visual tokens for each image grid.
            num_queries_vis_abstractors_slow: List of lists containing number of visual tokens for the slow part when
                applying the slowfast algorithm to video frames.
            first_last_frames_slows: List of booleans indicating whether the slowfast algorithm is applied to the first
                or last frames of the video.
            is_videos: List of booleans indicating which inputs are videos.
            img_start_ids_list: List of lists containing indices of img_start_id tokens for each sample.
            pad_token_id: Token ID used for padding.
            eos_token_id: Token ID used to signal the end of a sequence.
            bad_words_ids: List of token ID sequences that should not be generated.
            max_length: Maximum length of the sequence to be generated (input length + max_new_tokens).
            min_length: Minimum length of the sequence to be generated (input length + min_new_tokens).
            do_sample: Whether to use sampling for generation (otherwise uses greedy decoding).
            num_beams: Number of beams for beam search. 1 means no beam search.
            top_p: Nucleus sampling parameter. Tokens with cumulative probability > top_p are kept.
            top_k: Number of highest probability tokens to keep for top-k-filtering.
            temperature: Value used to modulate the next token probabilities.
            repetition_penalty: Penalty applied to tokens that have already appeared in the sequence.
            length_penalty: Exponential penalty applied to sequence length.
            use_cache: Whether to use past key/values for faster inference.
            **kwargs: Additional keyword arguments.

        Returns:
            Generated token IDs.
        """
        # inputs_embeds: torch.bfloat16 : [batchsize, variable(visual token, text token, system prompt 모두 포함)]
        if pad_token_id is None:
            pad_token_id = self.tokenizer.pad_token_id
        if eos_token_id is None:
            eos_token_id = self.tokenizer.encode("<|endofturn|>")[0]
        if bad_words_ids is None:
            bad_words_ids = [
                [
                    self.config.language_config["bos_token_id"],
                ],
                [
                    self.config.language_config["eos_token_id"],
                ],
            ]

        if pixel_values is None:
            return self.language_model.generate(
                input_ids, pad_token_id=pad_token_id, eos_token_id=eos_token_id, bad_words_ids=bad_words_ids, **kwargs
            )
        inputs_embeds = self.extract_inputs_embeds(
            input_ids=input_ids,
            pixel_values=self.to_vision_model_device(pixel_values),
            image_sizes=image_sizes,
            vision_query_lengths=vision_query_lengths,
            non_vision_query_lengths=non_vision_query_lengths,
            img_start_ids_list=img_start_ids_list,
            num_queries_vis_abstractors=num_queries_vis_abstractors,
            num_queries_vis_abstractors_slow=num_queries_vis_abstractors_slow,
            first_last_frames_slows=first_last_frames_slows,
            is_videos=is_videos,
        )
        inputs_embeds = inputs_embeds.to(device=self.language_model.device, dtype=self.language_model.dtype)

        # pred : torch.int64 : [batchsize, generated token_length]
        pred = self.language_model.generate(
            inputs_embeds=inputs_embeds,
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
            bad_words_ids=bad_words_ids,
            max_length=max_length,
            min_length=min_length,
            num_beams=num_beams,
            do_sample=(False if temperature == 0.0 else do_sample),  # set do_sample=False if invalid temperature
            top_k=top_k,
            top_p=top_p,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            early_stopping=(False if num_beams <= 1 else True),  # set early_stopping=False when not beam_search
            use_cache=use_cache,
            **kwargs,
        )

        return pred

    def to_vision_model_device(self, input_tensor: Union[torch.Tensor, List]) -> Union[torch.Tensor, List]:
        """Move input tensors to the vision model's device.
        This method recursively moves input tensors or lists of tensors to the vision model's device.

        Args:
            input_tensor: Input tensor or list of tensors to be moved to the vision model's device.

        Returns:
            The input tensor or list of tensors moved to the vision model's device.

        Raises:
            TypeError: If the input is neither a tensor nor a list.
        """
        if isinstance(input_tensor, list):
            return [self.to_vision_model_device(item) for item in input_tensor]
        elif isinstance(input_tensor, torch.Tensor):
            return input_tensor.to(self.vision_model.device)
        else:
            raise TypeError("Unsupported data type. Only tensors and lists are allowed.")

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        **kwargs,
    ) -> Dict[str, Any]:
        """Prepare inputs for the generation algorithm.

        This method prepares the input for each generation step based on the model's needs.

        Args:
            input_ids: Input token IDs.
            past_key_values: Pre-computed key and value states for faster inference.
            attention_mask: Mask to avoid performing attention on padding token indices.
            inputs_embeds: Input embeddings. If provided, input_ids will not be used.
            **kwargs: Additional keyword arguments.

        Returns:
            Dictionary containing the prepared inputs for the model.
        """
        input_ids = kwargs.get("decoder_input_ids", input_ids)

        if past_key_values:
            input_ids = input_ids[:, -1:]

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "pixel_values": kwargs.get("pixel_values", None),
            }
        )
        return model_inputs

    @classmethod
    def from_config(cls, config, vision_model_name_or_path):
        return cls(config, vision_model_name_or_path)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None,
        *model_args,
        **kwargs,
    ) -> "HCXVisionForCausalLM":
        assert pretrained_model_name_or_path is not None

        save_only_vision = kwargs.pop("save_only_vision") if "save_only_vision" in kwargs else False
        save_only_qformer = kwargs.pop("save_only_qformer") if "save_only_qformer" in kwargs else False
        save_shard_size = kwargs.pop("save_shard_size") if "save_shard_size" in kwargs else "5GB"

        if pretrained_model_name_or_path is not None:  # when evaluate or load instruction tunned model
            model: HCXVisionForCausalLM = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
            model.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)

        img_start_id = model.tokenizer.encode(IMG_LOC, add_special_tokens=False)
        assert (
            len(img_start_id) == 1
        ), f'"<|dummy3|>" was not encoded into a single special token. Encoding result: {img_start_id}'
        model.config.img_start_id = img_start_id[0]

        model.save_only_vision = save_only_vision
        model.save_only_qformer = save_only_qformer
        model.save_shard_size = save_shard_size

        return model

    def get_language_model(self):
        return self.language_model.base_model

    def get_vision_model(self):
        return self.vision_model

    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        *args,
        **kwargs,
    ):
        state_dict = kwargs["state_dict"] if "state_dict" in kwargs else self.state_dict()
        partial_state_dict = self.get_pretrained_state_dict(
            state_dict,
            save_directory,
        )
        kwargs["state_dict"] = partial_state_dict
        kwargs["safe_serialization"] = self.is_safetensor_save
        kwargs.setdefault("max_shard_size", self.save_shard_size)
        super().save_pretrained(save_directory, *args, **kwargs)

    def get_pretrained_state_dict(self, state_dict, save_dir):
        vision_key = "vision_model."
        llm_keys = ["language_model."]
        head_key = "lm_head."

        for key in list(state_dict.keys()):
            if self.save_only_vision:
                for llm_key in llm_keys:
                    if llm_key in key:
                        state_dict.pop(key)
                if key.startswith(head_key):
                    state_dict.pop(key)

                elif self.save_only_qformer:
                    if f"{vision_key}" in key:
                        state_dict.pop(key)

        return state_dict

    def compute_adaptive_params(
        self,
        pixel_values: Optional[List[List[torch.FloatTensor]]] = None,
        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
        num_queries_vis_abstractors_slow: Optional[List[List[int]]] = None,
        image_sizes: Optional[List[List[List[int]]]] = None,
        is_videos: Optional[List[bool]] = None,
        first_last_frames_slows: Optional[List[bool]] = None,
    ) -> Tuple[List[int], List[int], List[List[int]], List[bool], List[List[int]]]:
        """Compute adaptive parameters for processing different image and video inputs.

        This method calculates parameters needed for adaptive processing, especially when handling
        variable resolutions or applying the slowfast algorithm to video frames. It flattens
        batch-level inputs (lists of lists) into single lists representing all images/frames
        in the batch. Based on slowfast configuration, it may split video frames into 'slow'
        and 'fast' components, adjusting query counts and grid indices accordingly.

        Args:
            pixel_values: List of lists of image tensors (per sample). Used to determine the initial number of grids per
                image/frame.
            num_queries_vis_abstractors: List of lists (per sample) containing the base number of visual tokens
                generated by the visual abstractor for each image grid
                (e.g., 81 for a full grid, 9 for a subsampled/fast grid).
            num_queries_vis_abstractors_slow: List of lists (per sample) containing the number of visual tokens for the
                'slow' path when applying slowfast. Non-zero values here trigger the slowfast processing logic.
            image_sizes: List of lists (per sample) of original image dimensions ([width, height]).
            is_videos: List of lists (per sample) of booleans indicating if each input item is part of a video sequence.
            first_last_frames_slows: List (per sample) of booleans. If True, slowfast logic
                (if active based on `num_queries_vis_abstractors_slow`) is applied only to the first or last frame(s)
                within each video sequence.

        Returns:
            Tuple containing:
                - num_queries_vis_abstractors: Flattened list of final query counts per processed grid.
                  Values might be adjusted based on slow/fast splitting
                  (e.g., using values from `num_queries_vis_abstractors_slow` for slow frames).
                  Example: [81, 81, 81, 9, 81, 9, ...] (Image, Image, Vid_Slow, Vid_Fast, Vid_Slow, Vid_Fast...)
                - num_grids: Flattened list representing cumulative grid counts, acting as end indices for slicing the
                  flattened `image_forward_outs`. Adjusted for slow/fast splits.
                  Example: [0, 1, 9, 10, 18, 19, 27, ...] (Indices after Grid0_Slow(1),
                  Grid1_Fast(8), Grid2_Slow(1), Grid3_Fast(8)...).
                - image_sizes: Flattened list of image dimensions ([width, height]), potentially duplicated if slow/fast
                  splitting occurred.
                - is_videos: Flattened list of booleans indicating video status, potentially duplicated for
                  slow/fast splits. Example: [False, False, True, True, True, True, ...]
                  (Image1, Image2, Vid_grid1_slow, Vid_grid1_fast, Vid_grid2_slow, Vid_grid2_fast...)
                - group_ids: List of lists, grouping indices that correspond to the same original image or frame.
                  If a frame is split into slow/fast, its group will contain multiple indices.
                  Example: [[0], [1], [2, 3], [4, 5], ...]
                  (Group for Image1, Group for Image2, Group for Vid1_Slow+Fast, Group for Vid2_Slow+Fast...).

        Raises:
            AssertionError: If input validation fails (e.g., negative query counts).
            Exception: If an unexpected case is encountered during slowfast processing.
        """

        # Check if all elements are integers greater than or equal to 0
        assert all(
            all(isinstance(value, int) and value >= 0 for value in sublist) for sublist in num_queries_vis_abstractors
        ), "All values in num_queries_vis_abstractors must be integers >= 0."

        assert all(
            all(isinstance(value, int) and value >= 0 for value in sublist)
            for sublist in num_queries_vis_abstractors_slow
        ), "All values in num_queries_vis_abstractors_slow must be integers >= 0."

        assert is_videos is not None

        # Is it the first or last image? (for applying slowfast to video processing)
        is_first_images = []
        is_last_images = []
        for is_video in is_videos:
            for idx, is_video_item in enumerate(is_video):
                if idx == 0:
                    is_first_images.append(True)
                else:
                    is_first_images.append(False)
                if idx == len(is_video) - 1:
                    is_last_images.append(True)
                else:
                    is_last_images.append(False)

        num_queries_vis_abstractors = list(chain(*num_queries_vis_abstractors))
        num_queries_vis_abstractors_slow = list(chain(*num_queries_vis_abstractors_slow))
        image_sizes = list(chain(*image_sizes))
        is_videos = list(chain(*is_videos))
        first_last_frames_slows = list(chain(*first_last_frames_slows))

        # Use slowfast mode if there's at least one visual token count greater than 0 in num_queries_vis_abstractors_slow
        use_slowfast = any([num_query > 0 for num_query in num_queries_vis_abstractors_slow])
        num_grids = [pixel_value.shape[0] for pixel_value in chain(*pixel_values)]
        num_grids = [0] + num_grids
        group_ids = []

        if use_slowfast:
            new_num_grids = [num_grids[0]]
            new_num_queries = []
            new_image_sizes = []
            new_is_videos = []

            # When using slowfast, split more finely
            # 0th local grid is slow frame, remaining local grids are fast frames
            for (
                num_query,
                num_query_slow,
                num_grid,
                image_size,
                is_video,
                first_last_frames_slow,
                is_first_image,
                is_last_image,
            ) in zip(
                num_queries_vis_abstractors,
                num_queries_vis_abstractors_slow,
                num_grids[1:],
                image_sizes,
                is_videos,
                first_last_frames_slows,
                is_first_images,
                is_last_images,
            ):

                if not first_last_frames_slow and num_query_slow > 0:  # Process all image in slowfast mode
                    assert is_video  # slowfast mode is only applied to videos

                    this_group_ids = [group_ids[-1][-1] + 1 if group_ids else 0]

                    # slow frame (first grid)
                    new_num_grids.append(new_num_grids[-1] + 1)
                    new_num_queries.append(num_query_slow)
                    new_image_sizes.append(image_size)
                    new_is_videos.append(is_video)

                    if num_grid >= 2:
                        # fast frames
                        new_num_grids.append(new_num_grids[-1] + num_grid - 1)
                        new_num_queries.append(num_query)
                        new_image_sizes.append(image_size)
                        new_is_videos.append(is_video)
                        this_group_ids.append(this_group_ids[-1] + 1)

                    group_ids.append(this_group_ids)
                elif (
                    first_last_frames_slow and num_query_slow > 0 and (is_first_image or is_last_image)
                ):  # Process only first/last image in slowfast mode
                    # Case for special treatment of first/last frames in slow mode
                    assert is_video  # slowfast mode is only applied to videos

                    this_group_ids = [group_ids[-1][-1] + 1 if group_ids else 0]

                    if num_grid == 1:
                        # Simply process with slow since there's only one grid
                        new_num_grids.append(new_num_grids[-1] + 1)
                        new_num_queries.append(num_query_slow)
                        new_image_sizes.append(image_size)
                        new_is_videos.append(is_video)

                    if num_grid >= 2:
                        # Special treatment for first or last grid depending on is_first_image or is_last_image

                        if is_first_image:  # includes both first and last
                            # slow frame (first grid)
                            new_num_grids.append(new_num_grids[-1] + 1)
                            new_num_queries.append(num_query_slow)
                            new_image_sizes.append(image_size)
                            new_is_videos.append(is_video)
                            # fast frames
                            new_num_grids.append(new_num_grids[-1] + num_grid - 1)
                            new_num_queries.append(num_query)
                            new_image_sizes.append(image_size)
                            new_is_videos.append(is_video)
                            this_group_ids.append(this_group_ids[-1] + 1)
                        elif is_last_image:
                            # fast frames
                            new_num_grids.append(new_num_grids[-1] + num_grid - 1)
                            new_num_queries.append(num_query)
                            new_image_sizes.append(image_size)
                            new_is_videos.append(is_video)
                            # slow frame (last grid)
                            new_num_grids.append(new_num_grids[-1] + 1)
                            new_num_queries.append(num_query_slow)
                            new_image_sizes.append(image_size)
                            new_is_videos.append(is_video)
                            this_group_ids.append(this_group_ids[-1] + 1)
                        else:
                            raise Exception("This case should not be reached.")
                    group_ids.append(this_group_ids)
                else:
                    # Not in slowfast mode, so reduce all by num_query (fast)
                    new_num_grids.append(new_num_grids[-1] + num_grid)
                    new_num_queries.append(num_query)
                    new_image_sizes.append(image_size)
                    new_is_videos.append(is_video)

                    start_group_id = group_ids[-1][-1] + 1 if group_ids else 0
                    group_ids.append([start_group_id])

            num_grids = new_num_grids
            num_queries_vis_abstractors = new_num_queries
            image_sizes = new_image_sizes
            is_videos = new_is_videos
        else:
            num_grids = [sum(num_grids[:i]) for i in range(1, len(num_grids) + 1)]
            group_ids = [[group_id] for group_id in range(len(is_videos))]

        return num_queries_vis_abstractors, num_grids, image_sizes, is_videos, group_ids


class HCXVisionCAbstractor(nn.Module):
    """
    This module is based on C-Abstractor, whose license is under apache-2.0.
    You can check the original code at https://github.com/khanrc/honeybee/blob/main/honeybee/projectors/projectors.py
    and we made necessary modifications.
    """

    def __init__(
        self,
        num_queries: int,
        num_input_tokens: int,
        encoder_hidden_size: int,
        hidden_size: int,
        output_hidden_size: int,
        pos_emb: bool = True,
        prenorm: bool = False,
    ):
        super().__init__()
        self.num_input_tokens = num_input_tokens
        self.output_hidden_size = output_hidden_size

        # Positional embedding
        if pos_emb:
            self.pos_emb = torch.nn.Parameter(torch.zeros(1, num_input_tokens, encoder_hidden_size))
            self.pos_emb.data.normal_(mean=0.0, std=0.02)
        else:
            self.pos_emb = None

        # (Optional) Pre-normalization layer
        if prenorm:
            self.prenorm = LayerNorm(encoder_hidden_size)
        else:
            self.prenorm = None

        self.build_net(num_queries, encoder_hidden_size, hidden_size, output_hidden_size)
        self.dtype = next(self.parameters()).dtype

    def forward(
        self,
        x: torch.Tensor,
        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
        num_grids: Optional[List[int]] = None,
    ) -> torch.Tensor:
        """
        Args:
            x: (B, L, encoder_hidden_size) tensor from the visual backbone (e.g. CLIP visual encoder), including cls token.
        """
        if self.prenorm is not None:
            x = self.prenorm(x)

        if self.pos_emb is not None:
            x = x + self.pos_emb

        x = self._forward(
            x,
            num_queries_vis_abstractors=num_queries_vis_abstractors,
            num_grids=num_grids,
        )  # (B, L, output_hidden_size)

        return x

    def _forward(
        self,
        x: torch.Tensor,
        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
        num_grids: Optional[List[int]] = None,
    ) -> torch.Tensor:
        # x: [B, L, dim]
        B, L, dim = x.shape
        hw = int(L ** 0.5)
        x = rearrange(x, "b (h w) d -> b d h w", h=hw, w=hw)

        if num_queries_vis_abstractors is not None:
            assert num_grids is not None
            return self._forward_adaptive_num_query(x, num_queries_vis_abstractors, num_grids)

        x = self.net(x)
        x = rearrange(x, "b d h w -> b (h w) d")
        x = self.readout(x)
        return x

    def _forward_adaptive_num_query(
        self,
        x: torch.Tensor,
        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
        num_grids: Optional[List[int]] = None,
    ) -> List[torch.Tensor]:
        # self.net is consisted by 3 layers (s1, sampler, s2)
        assert len(self.net) == 3

        x = self.net[0](x)  # s1
        new_x = []
        for i, num_queries in enumerate(num_queries_vis_abstractors):
            hw = int(num_queries**0.5)
            sampler = nn.AdaptiveAvgPool2d((hw, hw))
            out = sampler(x[num_grids[i]:num_grids[i + 1], :])
            out = self.net[2](out)  # s2

            out = rearrange(out, "b d h w -> b (h w) d")
            out = self.readout(out)

            new_x.append(out)
        return new_x

    def build_net(
        self,
        n_queries: int,
        encoder_hidden_size: int,
        hidden_size: int,
        output_hidden_size: int,
        depth: int = 3,
        mlp_depth: int = 2,
    ):
        assert (n_queries ** 0.5).is_integer(), f"n_queries must be square number. n_queries: {n_queries}"
        hw = int(n_queries ** 0.5)

        # RegBlock = ResBlock + SE
        RegBlock = partial(
            RegStage,
            stride=1,
            dilation=1,
            act_layer=nn.SiLU,
            norm_layer=LayerNorm2d,
        )

        s1 = RegBlock(
            depth,
            encoder_hidden_size,
            hidden_size,
        )
        sampler = nn.AdaptiveAvgPool2d((hw, hw))
        s2 = RegBlock(
            depth,
            hidden_size,
            hidden_size,
        )

        self.net = nn.Sequential(s1, sampler, s2)
        self.readout = self.build_mlp(mlp_depth, hidden_size, output_hidden_size)

    def build_mlp(
        self,
        depth: int,
        hidden_size: int,
        output_hidden_size: int,
    ):
        layers = [nn.Linear(hidden_size, output_hidden_size)]
        for _ in range(1, depth):
            layers.append(nn.SiLU())
            layers.append(nn.Linear(output_hidden_size, output_hidden_size))
        return nn.Sequential(*layers)

def load_sharded_checkpoint(
    model, folder, pick_prefix="", replace_prefix_list=[], replace_prefix_dict={}, print_info=True
):
    if folder is None:
        return {}

    files = os.listdir(folder)

    # find relevant files
    pytorch_bin_files = [file for file in files if file.startswith("pytorch_model") and file.endswith(".bin")]
    safetensor_files = [file for file in files if file.endswith(".safetensors")]
    shard_index_file = [file for file in files if file.endswith(".index.json")]

    # check if sharded
    index_present = len(shard_index_file) > 0
    index_file = os.path.join(folder, shard_index_file[0]) if index_present else []

    # check if safetensor
    is_safetensor = len(safetensor_files) > 0

    model_keys = model.state_dict().keys()

    if is_safetensor:
        from safetensors.torch import load_file

        load_function = load_file
        shard_files = safetensor_files
    else:
        load_function = partial(torch.load, map_location="cpu")
        shard_files = pytorch_bin_files

    # sharded case
    if index_present:
        with open(index_file, "r", encoding="utf-8") as f:
            index = json.load(f)
        loaded_keys = index["weight_map"].keys()
        if pick_prefix:
            loaded_keys = [k[len(pick_prefix) :] for k in loaded_keys if k.startswith(pick_prefix)]
        if replace_prefix_list:
            for rep_prefix in replace_prefix_list:
                loaded_keys = [k[len(rep_prefix) :] if k.startswith(rep_prefix) else k for k in loaded_keys]
        if replace_prefix_dict:
            for rep_prefix in replace_prefix_dict:
                loaded_keys = [
                    k.replace(rep_prefix, replace_prefix_dict[rep_prefix]) if k.startswith(rep_prefix) else k
                    for k in loaded_keys
                ]

    for i, shard_file in enumerate(shard_files):
        state_dict = load_function(os.path.join(folder, shard_file))

        # if pick_prefix, use only pick
        if pick_prefix:
            state_dict = {k[len(pick_prefix) :]: v for k, v in state_dict.items() if k.startswith(pick_prefix)}

        for rep_prefix in replace_prefix_list:
            state_dict = {k[len(rep_prefix) :] if k.startswith(rep_prefix) else k: v for k, v in state_dict.items()}

        for rep_prefix in replace_prefix_dict:
            state_dict = {
                k.replace(rep_prefix, replace_prefix_dict[rep_prefix]) if k.startswith(rep_prefix) else k: v
                for k, v in state_dict.items()
            }

        if is_fsdp_enabled():
            if is_local_dist_rank_0():
                model.load_state_dict(state_dict, strict=False)
        else:
            model.load_state_dict(state_dict, strict=False)
        # Make sure memory is freed before we load the next state dict.

        if not index_present:
            loaded_keys = state_dict.keys()

        del state_dict
        gc.collect()

    # missing keys
    missing_keys = [key for key in model_keys if key not in loaded_keys]
    unexpected_keys = [key for key in loaded_keys if key not in model_keys]

    if get_rank() == 0 and print_info:
        print(f"[info] missing_keys: {missing_keys}")
        print(f"[info] unexpected_keys: {unexpected_keys}")

    return {"missing_keys": missing_keys, "unexpected_keys": unexpected_keys}