modeling_hyperclovax.py

import ast
import contextlib
import gc
import json
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 .image_processing_hyperclovax import select_best_resolution

EOT = "<|endofturn|>"
IMAGE_LOC = "<|dummy3|>"
VIDEO_LOC = "<|_unuse_missing_100270|>"


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: List[torch.FloatTensor],
    image_sizes: List[List[int]],
    possible_resolutions: List[Tuple[int, int]],
    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.
        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)

    # post-processing (unpad, add newline)
    new_image_features = []
    for image_idx, image_feature in enumerate(image_forward_outs):
        if image_feature.shape[0] > 1:
            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[0]
            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


@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 = ["SiglipEncoderLayer", "LlamaDecoderLayer", "HyperCLOVAXDecoderLayer"]
    supports_gradient_checkpointing = True
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn_2 = True
    _supports_sdpa = True

    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 text_config is not defined.
        """
        super().__init__(config)  # self.config = config

        # init configs
        text_config = self._init_text_config(config)
        vision_config = self._init_vision_config(config)
        
        ## possible_resolution should be matched with preprocessor_config.json
        config.possible_resolutions = self._init_possible_resolutions(config, vision_config)

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

        self.mm_projector = self._init_mm_projector(config, text_config, vision_config)

        self.language_model = AutoModelForCausalLM.from_config(text_config)
        self.lm_head_vocab_size = getattr(text_config, "padded_vocab_size", text_config.vocab_size)
        self.language_model.lm_head = nn.Linear(text_config.hidden_size, self.lm_head_vocab_size, bias=False)

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

        # modify configs or model settings
        if text_config.model_type in ["llama", "hyperclovax", "gpt2"]:
            self.language_model.gradient_checkpointing_enable()
        if text_config.model_type == "hyperclovax" and self.use_liger:
            self.language_model._get_apply_liger_kernel_converter()(model=self.language_model)

        # update configs
        self.vision_config = vision_config = self.vision_model.config
        self.text_config = text_config = self.language_model.config
        config.update({"vision_config": vision_config})
        config.update({"text_config": text_config})

        # etc
        self.use_liger = kwargs.pop("use_liger", False)
        self.use_fused_ce = kwargs.pop("use_fused_ce", False)
        self.use_meansum_loss = kwargs.pop("use_meansum_loss", False)
        self.freeze_before_sampler = kwargs.pop("freeze_before_sampler", False)
        self.use_turnmeansum_loss = kwargs.pop("use_turnmeansum_loss", False)
        self.vision_input_chunk_size = kwargs.pop("vision_input_chunk_size", None)
        self.is_safetensor_save = kwargs.get("is_safetensor_save", True)

        use_sum_loss = True if kwargs.pop("use_sum_loss", False) else False
        self.reduction = self._init_reduction_type(use_sum_loss)

        self.vision_model_use_no_grad = None  # forward 시 체크 및 할당

        self._backward_compatibility_gradient_checkpointing()  # self.post_init() 에 포함되어 있는 gc 가능한지 확인하고 켜주는 함수

    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 _init_reduction_type(self, use_sum_loss):
        assert not (
            self.use_meansum_loss and self.use_turnmeansum_loss
        ), "use_meansum_loss and use_turnmeansum_loss cannot both be True; only one or neither may be True."
        if self.use_meansum_loss or self.use_turnmeansum_loss:
            reduction = "none"
        elif use_sum_loss:
            reduction = "sum"
        else:
            reduction = "mean"
        return reduction

    def _init_vision_config(self, config):
        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.to_dict())
            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")

        vision_config.anyres = config.anyres
        vision_config.max_num_grids = config.max_num_grids
        return vision_config

    def _init_text_config(self, config):
        if hasattr(config, "text_config") and config.text_config is not None:
            model_type = config.text_config.model_type
            text_config = CONFIG_MAPPING[model_type](**config.text_config.to_dict())
        else:
            raise ValueError("text_config is not defined")
        text_config._attn_implementation = config._attn_implementation
        if text_config.model_type != "hyperclovax":
            text_config.logits_scaling = 1.0
        return text_config

    def _init_possible_resolutions(self, config, vision_config):
        """possible_resolution should be matched with preprocessor_config.json"""
        if not getattr(config, "possible_resolutions", []):
            possible_resolutions = []
            if config.anyres:
                assert config.max_num_grids > 0
                for i in range(1, config.max_num_grids + 1):
                    for j in range(1, config.max_num_grids + 1):
                        if i == 1 and j == 1 and not config.use_1x1_grid:
                            continue
                        if i * j <= config.max_num_grids:
                            possible_resolutions.append([i, j])

                possible_resolutions = [
                    [ys * vision_config.image_size, xs * vision_config.image_size] for ys, xs in possible_resolutions
                ]
            return possible_resolutions
        else:
            return config.possible_resolutions

    def _init_mm_projector(self, config, text_config, vision_config):
        input_hidden_size = vision_config.hidden_size
        if config.mm_projector_type == "linear":
            mm_projector = nn.Linear(input_hidden_size, text_config.hidden_size)
            mm_projector.dtype = next(mm_projector.parameters()).dtype
        elif config.mm_projector_type == "cabstractor":
            mm_projector = HCXVisionCAbstractor(
                num_queries=config.num_queries_vis_abstractor_image,
                num_input_tokens=(vision_config.image_size // vision_config.patch_size) ** 2,
                encoder_hidden_size=input_hidden_size,
                hidden_size=input_hidden_size,
                output_hidden_size=text_config.hidden_size,
                pos_emb=config.proj_pos_emb,
                prenorm=config.proj_prenorm,
            )
        else:
            mm_projector = HCXVisionMlp(
                config.mm_projector_type,
                input_hidden_size,
                hidden_features=input_hidden_size,  # TODO: llava 처럼 hidden_size 를 input_hidden_size 가 아니라 LLM embedding size 로 바꿔주기
                out_features=self.text_config.hidden_size,
            )
        return mm_projector

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values_images: Optional[List[List[torch.FloatTensor]]] = None,
        image_sizes_images: Optional[List[List[Tuple[int, int]]]] = None,
        pixel_values_videos: Optional[List[List[torch.FloatTensor]]] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        position_ids: Optional[torch.LongTensor] = 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,
        **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:
            if pixel_values_images is not None or pixel_values_videos is not None:
                inputs_embeds = self.extract_inputs_embeds(
                    input_ids=input_ids,
                    pixel_values_images=pixel_values_images,
                    image_sizes_images=image_sizes_images,
                    pixel_values_videos=pixel_values_videos,
                )
            else:
                inputs_embeds = self.get_input_embeddings()(input_ids)

        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,
            position_ids=position_ids,
            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.text_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,
        )

    # 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_images: Optional[List[List[torch.FloatTensor]]] = None,
        image_sizes_images: Optional[List[List[Tuple[int, int]]]] = None,
        pixel_values_videos: Optional[List[List[torch.FloatTensor]]] = 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.
            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.
        """
        # for convert back to List of List format
        len_pixel_values_images = [len(pixel_value) for pixel_value in pixel_values_images] if pixel_values_images else []
        len_pixel_values_videos = [len(pixel_value) for pixel_value in pixel_values_videos] if pixel_values_videos else []

        if sum(len_pixel_values_images) + sum(len_pixel_values_videos) == 0:
            return None

        inputs_embeds = self.get_input_embeddings()(input_ids)

        if sum(len_pixel_values_images) > 0:
            image_features_batch = self.forward_images(
                pixel_values_images, image_sizes_images, len_pixel_values_images
            )
            for i, image_features in enumerate(image_features_batch):
                if len(image_features) > 0:
                    image_token_indices = (input_ids[i] == self.config.image_token_id).nonzero().squeeze()
                    inputs_embeds[i][image_token_indices] = torch.cat(image_features).to(inputs_embeds.dtype)

        if sum(len_pixel_values_videos) > 0:
            video_features_batch = self.forward_videos(pixel_values_videos, len_pixel_values_videos)
            for i, video_features in enumerate(video_features_batch):
                if len(video_features) > 0:
                    video_token_indices = (input_ids[i] == self.config.video_token_id).nonzero().squeeze()
                    inputs_embeds[i][video_token_indices] = torch.cat(video_features).to(inputs_embeds.dtype)

        return inputs_embeds

    def forward_images(
        self,
        pixel_values_images: List[List[torch.FloatTensor]],
        image_sizes_images: List[List[Tuple[int, int]]],
        len_pixel_values_images: List[int],
    ) -> List[List[torch.Tensor]]:
        if sum(len_pixel_values_images) == 0:
            return None

        concat_pixel_values_images = torch.cat(list(chain(*pixel_values_images)), dim=0)

        visual_token_idx = 0 if "siglip" in self.vision_config.model_type else 1
        context_vision_model = torch.no_grad() if self.vision_model_use_no_grad else contextlib.nullcontext()
        with context_vision_model:
            if self.config.use_nth_layer == -1:
                # Replace post_layernorm of the last layer with Identity
                self.vision_model.vision_model.post_layernorm = nn.Identity()
                image_forward_outs = self.vision_model(concat_pixel_values_images)
                image_forward_outs = image_forward_outs.last_hidden_state[:, visual_token_idx:]
            else:
                image_forward_outs = self.vision_model(concat_pixel_values_images, output_hidden_states=True)
                image_forward_outs = image_forward_outs.hidden_states[self.config.use_nth_layer][:, visual_token_idx:]

        image_forward_outs = image_forward_outs.to(dtype=self.mm_projector.dtype)
        image_forward_outs = self.mm_projector(image_forward_outs)  # b (h w) d

        # feature 를 분할. e.g. torch.Size([18, 81, 3072]) -> [torch.Size([9, 81, 3072]), torch.Size([9, 81, 3072])]
        split_sizes = [pixel_value.shape[0] for pixel_value in chain(*pixel_values_images)]
        image_forward_outs = torch.split(image_forward_outs, split_sizes, dim=0)

        # newline 붙여주기 (anyres postprocessing)
        image_features = anyres_postprocessing(
            image_forward_outs=image_forward_outs,
            image_sizes=[image_size for image_sizes in image_sizes_images for image_size in image_sizes],
            num_queries_vis_abstractor=self.config.num_queries_vis_abstractor_image,
            unpad=self.config.unpad,
            patch_size=self.vision_config.patch_size,
            grid_size=self.vision_config.image_size,
            image_newline=self.image_newline,
            possible_resolutions=self.config.possible_resolutions,
        )

        # 원래 pixel_values_images 형태로 복원
        image_features = [
            image_features[sum(len_pixel_values_images[:i]) : sum(len_pixel_values_images[: i + 1])]
            for i in range(len(len_pixel_values_images))
        ]

        return image_features

    def forward_videos(
        self,
        pixel_values_videos: List[List[torch.FloatTensor]],
        len_pixel_values_videos: List[int],
    ) -> List[torch.Tensor]:

        len_video_grids = sum(len_pixel_values_videos)
        if len_video_grids == 0:
            return None

        # Run Vision Model
        concat_pixel_values_videos = torch.cat(list(chain(*pixel_values_videos)), dim=0)

        visual_token_idx = 0 if "siglip" in self.vision_config.model_type else 1
        context_vision_model = torch.no_grad() if self.vision_model_use_no_grad else contextlib.nullcontext()
        with context_vision_model:
            if self.config.use_nth_layer == -1:
                # Replace post_layernorm of the last layer with Identity
                self.vision_model.vision_model.post_layernorm = nn.Identity()
                video_forward_outs = self.vision_model(concat_pixel_values_videos)
                video_forward_outs = video_forward_outs.last_hidden_state[:, visual_token_idx:]
            else:
                video_forward_outs = self.vision_model(concat_pixel_values_videos, output_hidden_states=True)
                video_forward_outs = video_forward_outs.hidden_states[self.config.use_nth_layer][:, visual_token_idx:]

        video_forward_outs = video_forward_outs.to(dtype=self.mm_projector.dtype)

        # Run MM-Projector
        # len(num_grids) == len(num_queries_vis_abstractors) + 1
        grid_idx = 0
        num_grids = [grid_idx]  # e.g. [0, 9, 18, 19, 27, 28, 36, 37, 45, 46, 54, 55, 56]
        num_queries_vis_abstractors = []  # e.g. [81, 81, 81, 9, 81, 9, 81, 9, 81, 9, 81, 9]
        len_total_frames = video_forward_outs.shape[0]

        if self.config.first_last_frames_slow:
            # TODO: 동작 확인 안 했음. 해야 함.
            # slowfast (first_last_frames_slow)
            assert len_total_frames != 0
            if len_total_frames <= 2:
                num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_slow)
                grid_idx += len_total_frames
                num_grids.append(grid_idx)
            else:
                num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_slow)
                grid_idx += 1
                num_grids.append(grid_idx)

                num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_fast)
                grid_idx += len_total_frames - 2
                num_grids.append(grid_idx)

                num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_slow)
                grid_idx += 1
                num_grids.append(grid_idx)
        else:
            # slowfast
            for pixel_values_frames in pixel_values_videos:
                for pixel_values_frame in pixel_values_frames:
                    if len(pixel_values_frame) > 0:
                        num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_slow)
                        grid_idx += 1
                        num_grids.append(grid_idx)
                        num_queries_vis_abstractors.append(self.config.num_queries_vis_abstractor_video_fast)
                        grid_idx = grid_idx + len(pixel_values_frame) - 1
                        num_grids.append(grid_idx)

        video_forward_outs = self.mm_projector(video_forward_outs, num_queries_vis_abstractors, num_grids)

        # video_group 별로 concat 처리.
        # 예를 들어, 3x3 grid 를 사용했을 경우, 총 9개의 feature 가 모일 때까지, grouped_features 에 리스트를 모아주고, concat 처리.
        video_features = []  # what we want to return
        target_features = []
        target_group_size = 0
        group_counter = 0
        video_groups = [
            len(frame) for frames in pixel_values_videos for frame in frames
        ]  # for concat video features after projector

        for forward_out in video_forward_outs:
            target_group_size += len(forward_out)
            target_features.append(forward_out.flatten(0, 1))

            video_group_size = video_groups[group_counter]
            if video_group_size == target_group_size:
                video_features.append(torch.cat(target_features, dim=0))
                target_features = []
                group_counter += 1
                target_group_size = 0

            elif video_group_size < target_group_size:
                raise RuntimeError(f"video_group_size < target_group_size!! [{video_group_size} < {target_group_size}]")

        assert len(target_features) == 0, f"target_features is not empty!! {target_features}"
        assert len(video_groups) == len(video_features)

        # 원래 pixel_values_videos 형태로 복원
        video_features = [
            video_features[sum(len_pixel_values_videos[:i]) : sum(len_pixel_values_videos[: i + 1])]
            for i in range(len(len_pixel_values_videos))
        ]

        return video_features

    @torch.no_grad()
    def generate(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values_images: Optional[List[List[torch.FloatTensor]]] = None,
        image_sizes_images: Optional[List[List[Tuple[int, int]]]] = None,
        pixel_values_videos: Optional[List[List[torch.FloatTensor]]] = 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,
        verbose: bool = False,
        **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.text_config.bos_token_id,
                ],
                [
                    self.config.text_config.eos_token_id,
                ],
            ]

        if (pixel_values_images is None or all(len(pixel_values) == 0 for pixel_values in pixel_values_images)) and (
            pixel_values_videos is None or all(len(pixel_values) == 0 for pixel_values in pixel_values_videos)
        ):
            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_images=pixel_values_images,
            image_sizes_images=image_sizes_images,
            pixel_values_videos=pixel_values_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_new_tokens=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,
        )

        if verbose:
            llm_query = self.tokenizer.batch_decode(
                [
                    [token_id for token_id in input_ids_row if token_id != self.tokenizer.pad_token_id]
                    for input_ids_row in input_ids.detach().cpu().tolist()
                ],
                skip_special_tokens=False,
            )[0]
            llm_pred = self.tokenizer.batch_decode(
                [
                    [token_id for token_id in pred_row if token_id != self.tokenizer.pad_token_id]
                    for pred_row in pred.detach().cpu().tolist()
                ],
                skip_special_tokens=False,
            )[0]
            print(f"# [info] llm_query: {llm_query}")
            print(f"# [info] llm_pred: {llm_pred}")

        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)

        image_token_id = model.tokenizer.encode(IMAGE_LOC, add_special_tokens=False)
        assert (
            len(image_token_id) == 1
        ), f'"<|dummy3|>" was not encoded into a single special token. Encoding result: {image_token_id}'
        model.config.image_token_id = image_token_id[0]
        
        video_token_id = model.tokenizer.encode(VIDEO_LOC, add_special_tokens=False)
        assert (
            len(video_token_id) == 1
        ), f'"<|_unuse_missing_100270|>" was not encoded into a single special token. Encoding result: {video_token_id}'
        model.config.video_token_id = video_token_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



class HCXVisionMlp(nn.Module):
    def __init__(
        self,
        mm_projector_type,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.mm_projector_type = mm_projector_type
        if self.mm_projector_type == "mlp":
            self.fc1 = nn.Linear(in_features, hidden_features)
            self.act = act_layer()
            self.fc2 = nn.Linear(hidden_features, out_features)
        elif self.mm_projector_type == "inverted_mlp":
            self.fc1 = nn.Linear(in_features, 2 * hidden_features)
            self.act = act_layer()
            self.fc2 = nn.Linear(2 * hidden_features, out_features)
        else:
            raise NotImplementedError("{} is not implemented".format(self.mm_projector_type))

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x


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)