modeling.py

from typing import List, Optional, Tuple, Union
from dataclasses import dataclass

import torch
from torch import nn
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    LlavaNextForConditionalGeneration,
    LlavaNextModel,
)
from transformers.models.llava_next.modeling_llava_next import (
    LlavaNextCausalLMOutputWithPast,
    LlavaNextPreTrainedModel,
    LlavaNextMultiModalProjector,
    get_anyres_image_grid_shape,
    image_size_to_num_patches,
    unpad_image,
    LlavaNextModelOutputWithPast
)
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, can_return_tuple, logging
from accelerate import init_empty_weights
from transformers import Blip2QFormerConfig, Blip2QFormerModel
from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig
from .configuration import Granite4VisionConfig, Granite4VisionConfigNaflex
from .downsampling import BilinearDownsampler, QFormerDownsampler, WindowQFormerDownsampler
import math
import numpy as np
from fractions import Fraction
from transformers.modeling_utils import flash_attention_forward
from transformers.models.granitemoehybrid.modeling_granitemoehybrid import HybridMambaAttentionDynamicCache
IGNORE_INDEX = -100
logger = logging.get_logger(__name__)

@dataclass
class Granite4VisionModelOutputWithPast(LlavaNextModelOutputWithPast):
    r"""
    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    image_hidden_states (`torch.FloatTensor`, *optional*):
        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
    """

    balancing_loss: Optional[torch.FloatTensor] = None

@dataclass
class Granite4VisionCausalLMOutputWithPast(LlavaNextCausalLMOutputWithPast):
    r"""
    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    image_hidden_states (`torch.FloatTensor`, *optional*):
        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
    """

    balancing_loss: Optional[torch.FloatTensor] = None


class ParamWrapper(nn.Module):
    def __init__(self, param):
        super().__init__()
        self.param = param

class Granite4VisionForConditionalGeneration(LlavaNextForConditionalGeneration):
    config_class = Granite4VisionConfig

    def __init__(self, config: Granite4VisionConfig):
        # Update config with pretrained models if specified
        if config.pretrained_vision_tower:
            config.vision_config = AutoConfig.from_pretrained(
                config.pretrained_vision_tower, **config.vision_config.to_dict()
            )
            config.vision_config = (
                config.vision_config.vision_config
                if hasattr(config.vision_config, "vision_config")
                else config.vision_config
            )
        if config.pretrained_language_model:
            config.text_config = AutoConfig.from_pretrained(
                config.pretrained_language_model, **config.text_config.to_dict()
            )

        # Initialize parent
        LlavaNextPreTrainedModel.__init__(self, config)

        # Create custom model instance
        self.model = Granite4VisionModel(config)

        # Create lm_head
        self.lm_head = nn.Linear(
            config.text_config.hidden_size, config.text_config.vocab_size, bias=False
        )

        # Load pretrained components if specified
        if config.pretrained_vision_tower:
            self._load_pretrained_vision_tower(config)
            config.pretrained_vision_tower = ""

        if config.pretrained_language_model:
            self._load_pretrained_language_model(config)
            config.pretrained_language_model = ""

        self.post_init()

    def _load_pretrained_vision_tower(self, config):
        """Load pretrained vision tower weights"""
        print(f"Loading vision tower from: {config.pretrained_vision_tower}")
        vision_tower = AutoModel.from_pretrained(
            config.pretrained_vision_tower,
            attn_implementation="flash_attention_2",
            device_map="cpu",
            dtype=torch.bfloat16,

        )
        self.model.vision_tower = self.model.vision_tower.to(torch.bfloat16)
        print(self.model.vision_tower.load_state_dict(vision_tower.state_dict(), strict=False).missing_keys)
        self.model.vision_tower.config._attn_implementation = "flash_attention_2"
        # todo: (Avihu) would have done this but afraid - maybe something I'm missing
        # self.model.vision_tower = vision_tower
        self.config.vision_config = (
            self.model.vision_tower.config.vision_config
            if hasattr(self.model.vision_tower.config, "vision_config")
            else self.model.vision_tower.config
        )

    def _load_pretrained_language_model(self, config):
        """Load pretrained language model weights"""
        print(f"Loading language model from: {config.pretrained_language_model}")
        language_model = AutoModelForCausalLM.from_pretrained(
            config.pretrained_language_model,
            device_map="cpu",
            attn_implementation="flash_attention_2",
            dtype=torch.bfloat16,
            # use_kernels=True,
        )
        if self.config.image_token_index >= language_model.config.vocab_size:
            language_model.resize_token_embeddings(self.config.image_token_index + 1)
        # load weights in quantized mode with kernels
        self.model.language_model = language_model.model
        self.lm_head = language_model.lm_head
        # Load weights into the language model inside self.model
        self.config.text_config = self.model.language_model.config


    @can_return_tuple
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        image_sizes: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        vision_feature_layer: Optional[Union[int, list[int]]] = None,
        vision_feature_select_strategy: Optional[str] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        spatial_shapes: Optional[torch.LongTensor] = None,
        pixel_attention_mask: Optional[torch.Tensor] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Union[tuple, Granite4VisionCausalLMOutputWithPast]:

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        vision_feature_layer = (
            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
        )
        vision_feature_select_strategy = (
            vision_feature_select_strategy
            if vision_feature_select_strategy is not None
            else self.config.vision_feature_select_strategy
        )

        outputs = self.model(
            input_ids,
            pixel_values=pixel_values,
            image_sizes=image_sizes,
            vision_feature_layer=vision_feature_layer,
            vision_feature_select_strategy=vision_feature_select_strategy,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            cache_position=cache_position,
            spatial_shapes=spatial_shapes,
            pixel_attention_mask=pixel_attention_mask,
            **kwargs,
        )

        hidden_states = outputs.last_hidden_state
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])
        logits = logits / self.config.text_config.logits_scaling

        loss = None
        if labels is not None:
            # Upcast to float if we need to compute the loss to avoid potential precision issues
            # Avihu: removed the .float(), didn't make a huge difference and requires more memory
            # logits = logits.float()
            # Flatten the tokens
            loss = self.loss_function(
                logits,
                labels,
                vocab_size=self.config.text_config.vocab_size,
                **kwargs,
            )

        return Granite4VisionCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            image_hidden_states=outputs.image_hidden_states,
            balancing_loss=outputs.balancing_loss
        )
        
    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        inputs_embeds=None,
        pixel_values=None,
        image_sizes=None,
        attention_mask=None,
        cache_position=None,
        logits_to_keep=None,
        **kwargs,
    ):
        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model

        model_inputs = super().prepare_inputs_for_generation(
            input_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_position=cache_position,
            logits_to_keep=logits_to_keep,
            **kwargs,
        )

        # Check if the model or its langauge model are moe type - requires special cache handling
        if any(class_name in self.__class__.__name__.lower() or class_name in self.language_model.__class__.__name__.lower() for class_name in ["moe"]):
            model_inputs = self.prepare_inputs_for_generation_granite_moe(**model_inputs)
            
        # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore
        # Otherwise we need pixel values to be passed to model
        if cache_position[0] == 0:
            model_inputs["pixel_values"] = pixel_values
            model_inputs["image_sizes"] = image_sizes

        return model_inputs

    # Avihu: would have used the GraniteMoeSharedForCausalLM method, but we don't store this object anymore (split the model / lm head)
    def prepare_inputs_for_generation_granite_moe(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        cache_position=None,
        position_ids=None,
        use_cache=True,
        **kwargs,
    ):
        # Overwritten -- has a unique cache type, `HybridMambaAttentionDynamicCache`
        # Note: (Avihu) in transformers v4, the past_key_values is already an empty DynamicCache object. Testing that too
        empty_past_kv = past_key_values is None or (isinstance(past_key_values, DynamicCache) and past_key_values[0][0] is None)

        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
        # Exception 1: when passing input_embeds, input_ids may be missing entries
        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
        # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
        #              (we can't check exception 3 while compiling)
        if not empty_past_kv:
            if (
                inputs_embeds is not None  # Exception 1
                or cache_position[-1] >= input_ids.shape[1]  # Exception 3
            ):
                input_ids = input_ids[:, -cache_position.shape[0] :]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]
        elif use_cache:
            past_key_values = HybridMambaAttentionDynamicCache(
                self.model.language_model.config, input_ids.shape[0], self.dtype, device=self.device
            )

        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if not empty_past_kv:
                position_ids = position_ids[:, -input_ids.shape[1] :]

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and empty_past_kv:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids.contiguous()}  # `contiguous()` needed for compilation use cases

        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_cache": use_cache,
                "attention_mask": attention_mask,
                "cache_position": cache_position,
            }
        )

        # Forward ALL kwargs that are uninitialized (e.g. `use_cache`).
        for key, value in kwargs.items():
            if key not in model_inputs:
                model_inputs[key] = value

        return model_inputs


class Granite4VisionModel(LlavaNextPreTrainedModel):
    config_class = Granite4VisionConfig
    
    def __init__(self, config: Granite4VisionConfig):
        super().__init__(config)
        self.vision_tower = AutoModel.from_config(config.vision_config)
        self.multi_modal_projector = LlavaNextMultiModalProjector(config)
        self.downsampler = None
        self.downsample_rate = config.downsample_rate
        if config.downsample_rate is not None:
            if config.downsample_method in ["interpolate", "bilinear"]:
                self.downsampler = BilinearDownsampler(config)
            elif config.downsample_method == "qformer":
                self.downsampler = QFormerDownsampler(config)
            elif config.downsample_method == "window_qformer":
                self.downsampler = WindowQFormerDownsampler(config)
        self.image_newline = None
        if config.use_image_newline_parameter:
            embed_std = 1 / math.sqrt(config.text_config.hidden_size)
            image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size, dtype=self.dtype) * embed_std)
            self.model_type = config.model_type
            if self.model_type in ["gpt_vision", "granite4_vision"]:
                # this hack allows to do lora training from scratch, so image_newline would be in modules_to_keep
                self.image_newline = ParamWrapper(image_newline)
            else:
                self.image_newline = image_newline
        self.vocab_size = config.text_config.vocab_size
        
        # with init_empty_weights(): # Avihu: hack to load the model faster
        self.language_model = AutoModel.from_config(config.text_config)
        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
        self.post_init()
        
    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    def set_decoder(self, decoder):
        self.language_model = decoder

    def get_decoder(self):
        return self.language_model


    def pack_image_features(self, image_features, image_sizes, vision_feature_select_strategy, image_newline=None):
        """
        Reshape, unpad and then pack each image_feature into a single image_features tensor containing all visual vectors.

        Args:
            image_features (`list[torch.Tensor]` of length num_images, each of shape `(num_patches, image_length, embed_dim)`)
                List of image feature tensor, each contains all the visual feature of all patches.
            image_sizes (`torch.Tensor` of shape `(num_images, 2)`)
                Actual image size of each images (H, W).
            vision_feature_select_strategy (`str`)
                The feature selection strategy used to select the vision feature from the vision backbone.
            image_newline (`torch.Tensor` of shape `(embed_dim)`)
                New line embedding vector.
        Returns:
            image_features (`torch.Tensor` of shape `(all_feat_len, embed_dim)`)
            feature_lens (`list[int]`)
                token length of each image in image_features
        """
        new_image_features = []
        feature_lens = []
        for image_idx, image_feature in enumerate(image_features):
            if image_feature.shape[0] > 1:
                base_image_feature = image_feature[0]
                image_feature = image_feature[1:]
                height = width = self.config.vision_config.image_size // self.config.vision_config.patch_size

                num_patch_height, num_patch_width = get_anyres_image_grid_shape(
                    image_sizes[image_idx],
                    self.config.image_grid_pinpoints,
                    self.config.vision_config.image_size,
                )
                if self.downsampler is not None:
                    ds_rate = Fraction(self.downsample_rate)
                    height = int(height  * ds_rate)
                    width = int(width  * ds_rate)

                if (
                    np.prod(image_feature.shape) % (num_patch_height * num_patch_width * height * width) != 0
                    and vision_feature_select_strategy == "default"
                ):
                    logger.warning_once(
                        "Image feature shape does not line up with the provided patch size. "
                        "You may be using the `default` vision_feature_select_strategy with a"
                        " visual encoder that does not have CLS."
                    )

                image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1)
                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_sizes[image_idx])
                if image_newline is not None:
                    image_feature = torch.cat(
                        (
                            image_feature,
                            image_newline[:, None, None]
                            .expand(*image_feature.shape[:-1], 1)
                            .to(image_feature.device, image_feature.dtype),
                        ),
                        dim=-1,
                    )
                image_feature = image_feature.flatten(1, 2).transpose(0, 1)
                image_feature = torch.cat((base_image_feature, image_feature), dim=0)
            else:
                image_feature = image_feature[0]
                if image_newline is not None:
                    image_feature = torch.cat((image_feature, image_newline[None].to(image_feature)), dim=0)
            new_image_features.append(image_feature)
            feature_lens.append(image_feature.size(0))
        feature_lens = torch.tensor(feature_lens, dtype=torch.long, device=image_features[0].device)
        return new_image_features, feature_lens

    def get_image_features(
        self,
        pixel_values: torch.FloatTensor,
        image_sizes: torch.Tensor,
        vision_feature_layer: Optional[Union[int, list[int]]] = None,
        vision_feature_select_strategy: Optional[str] = None,
    ):
        """
        Obtains image last hidden states from the vision tower and apply multimodal projection.

        Args:
            pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`)
               The tensors corresponding to the input images.
            image_sizes (`torch.Tensor` of shape `(num_images, 2)`)
                Actual image size of each images (H, W).
            vision_feature_layer (`Union[int, list[int]]`, *optional*):
                The index of the layer to select the vision feature. If multiple indices are provided,
                the vision feature of the corresponding indices will be concatenated to form the
                vision features.
            vision_feature_select_strategy (`str`, *optional*):
                The feature selection strategy used to select the vision feature from the vision backbone.
                Can be one of `"default"` or `"full"`
        Returns:
            image_features (list[`torch.Tensor`]): List of image feature tensor, each contains all the visual feature of all patches
            and are of shape `(num_patches, image_length, embed_dim)`).
        """
        vision_feature_layer = (
            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
        )
        vision_feature_select_strategy = (
            vision_feature_select_strategy
            if vision_feature_select_strategy is not None
            else self.config.vision_feature_select_strategy
        )

        # ! infer image_num_patches from image_sizes
        image_num_patches = [
            image_size_to_num_patches(
                image_size=imsize,
                grid_pinpoints=self.config.image_grid_pinpoints,
                patch_size=self.config.vision_config.image_size,
            )
            for imsize in image_sizes
        ]
        if pixel_values.dim() == 5:
            # stacked if input is (batch_size, num_patches, num_channels, height, width)
            _pixel_values_list = [pix_val[:num_patch] for pix_val, num_patch in zip(pixel_values, image_num_patches)]
            pixel_values = torch.cat(_pixel_values_list, dim=0)
        elif pixel_values.dim() != 4:
            # otherwise has to be stacked from list of (num_patches, num_channels, height, width)
            raise ValueError(f"pixel_values of shape {pixel_values.shape}, expect to be of 4 or 5 dimensions")

        image_features = self.vision_tower(pixel_values, output_hidden_states=True)
        # If we have one vision feature layer, return the corresponding hidden states,
        # otherwise, select the hidden states of each feature layer and concatenate them
        if isinstance(vision_feature_layer, int):
            selected_image_feature = image_features.hidden_states[vision_feature_layer]
        else:
            hs_pool = [image_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer]
            selected_image_feature = torch.cat(hs_pool, dim=-1)

        if vision_feature_select_strategy == "default":
            selected_image_feature = selected_image_feature[:, 1:]

        image_features = self.multi_modal_projector(selected_image_feature)
        if self.downsampler is not None:
            # training this with peft+deepspeed had this issue+fix: 
            # https://github.com/deepspeedai/DeepSpeed/issues/7203#issuecomment-3007490737
            image_features = self.downsampler(image_features)
        if image_features.shape[0] != sum(image_num_patches):
            print("about to crash on split", pixel_values.shape, image_sizes, image_num_patches)
        image_features = torch.split(image_features, image_num_patches, dim=0)

        # NOTE we only support multimodal_patch_merge_type == "spatial_unpad"
        image_newline = self.image_newline.param if self.model_type in ["gpt_vision", "granite4_vision"] else self.image_newline
        image_features, feature_lens = self.pack_image_features(
            image_features,
            image_sizes,
            vision_feature_select_strategy=vision_feature_select_strategy,
            image_newline=image_newline,
        )
        return image_features

    def get_image_features_naflex(
        self,
        pixel_values: torch.FloatTensor,
        spatial_shapes, 
        pixel_attention_mask,
        vision_feature_layer: Optional[Union[int, list[int]]] = None,
    ):
        vision_feature_layer = (
            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
        )
        # todo: (Avihu): Hack! siglip2 naflex now supports pad-free
        # todo: This was done by manually editing the siglip modeling code
        # todo: Until we have a better solution, to run this, consult with me
        # Note! siglip gets a stacked tensor
        image_features = self.vision_tower(pixel_values, spatial_shapes=spatial_shapes,
                                           pixel_attention_mask=pixel_attention_mask, output_hidden_states=True)
        # If we have one vision feature layer, return the corresponding hidden states,
        # otherwise, select the hidden states of each feature layer and concatenate them
        if isinstance(vision_feature_layer, int):
            selected_image_feature = image_features.hidden_states[vision_feature_layer]
        else:
            hs_pool = [image_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer]
            selected_image_feature = torch.cat(hs_pool, dim=-1)

        image_features = self.multi_modal_projector(selected_image_feature)
        # Note (Avihu): downsampling would and newline is more complex at the moment
        assert self.downsampler is None, "downsampler not supported for naflex yet"
        assert self.image_newline is None, "newline not supported for naflex yet"
        return image_features


    
    def get_placeholder_mask(
        self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
    ):
        """
        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
        equal to the length of multimodal features. If the lengths are different, an error is raised.
        """
        if input_ids is None:
            special_image_mask = inputs_embeds == self.get_input_embeddings()(
                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
            )
            special_image_mask = special_image_mask.all(-1)
        else:
            special_image_mask = input_ids == self.config.image_token_id
            
        n_image_tokens = special_image_mask.sum()
        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
        if inputs_embeds[special_image_mask].numel() != image_features.numel():
            raise ValueError(
                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}"
            )
        return special_image_mask

    @can_return_tuple
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        image_sizes: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        vision_feature_layer: Optional[Union[int, list[int]]] = None,
        vision_feature_select_strategy: Optional[str] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        spatial_shapes: Optional[torch.LongTensor] = None,
        pixel_attention_mask: Optional[torch.Tensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> Union[tuple, Granite4VisionModelOutputWithPast]:
        r"""
        vision_feature_select_strategy (`str`, *optional*, defaults to `"default"`):
            The feature selection strategy used to select the vision feature from the vision backbone.
            Can be one of `"default"` or `"full"`. If `"default"`, the CLS token is removed from the vision features.
            If `"full"`, the full vision features are used.
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        vision_feature_layer = (
            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
        )
        vision_feature_select_strategy = (
            vision_feature_select_strategy
            if vision_feature_select_strategy is not None
            else self.config.vision_feature_select_strategy
        )

        if (input_ids is None) ^ (inputs_embeds is not None):
            print(input_ids, inputs_embeds, position_ids, pixel_values, image_sizes, kwargs, )
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

        if pixel_values is not None and pixel_values.size(0) > 0:
            
            if spatial_shapes is not None and pixel_attention_mask is not None:
                # naflex setup
                image_features = self.get_image_features_naflex(
                    pixel_values,
                    spatial_shapes, 
                    pixel_attention_mask,
                    vision_feature_layer=vision_feature_layer
                )
            else:
                image_features = self.get_image_features(
                    pixel_values,
                    image_sizes,
                    vision_feature_layer=vision_feature_layer,
                    vision_feature_select_strategy=vision_feature_select_strategy,
                )
                image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
            special_image_mask = self.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, image_features=image_features
            )
            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
        elif torch.is_grad_enabled():
            self.run_dummy_encoder_forward(inputs_embeds, vision_feature_layer, vision_feature_select_strategy)
        try:
            outputs = self.language_model(
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=True,
                cache_position=cache_position,
                **kwargs,
            )
        except Exception as e:
            print(e)
            print(attention_mask)
            print(position_ids)
            print(inputs_embeds)
            print(input_ids)
            print(kwargs)
            raise e
        return Granite4VisionModelOutputWithPast(
            last_hidden_state=outputs.last_hidden_state,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            image_hidden_states=image_features if pixel_values is not None else None,
        )

    def run_dummy_encoder_forward(self, inputs_embeds, vision_feature_layer, vision_feature_select_strategy):
        if isinstance(self.config.vision_config, Siglip2VisionConfig):
            print("no pixel values, using dummy data to get grads - naflex mode")
            dummy_pixel_values = torch.zeros((1, 256, 768), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
            dummy_spatial_shapes = torch.tensor([[16, 16]], device=inputs_embeds.device)
            dummy_pixel_attention_mask = torch.ones((1,256), device=inputs_embeds.device)
            other_embeds = self.get_image_features_naflex(
                    dummy_pixel_values,
                    dummy_spatial_shapes, 
                    dummy_pixel_attention_mask,
                    vision_feature_layer=vision_feature_layer
                )
            other_embeds = other_embeds[0][:1] * 0 # adding zeros tensor
            inputs_embeds[0, :1] = inputs_embeds[0, :1] + other_embeds

        else:
            print("no pixel values, using dummy data to get grads")
            dummy_data = torch.zeros(
                (3, 3, 384, 384), dtype=inputs_embeds.dtype, device=inputs_embeds.device
            )
            dummy_sizes = torch.tensor([[768, 384]], device=inputs_embeds.device)
            other_embeds = self.get_image_features(dummy_data, dummy_sizes, 
                                                    vision_feature_layer=vision_feature_layer,
                                                    vision_feature_select_strategy=vision_feature_select_strategy)
            other_embeds = other_embeds[0][:1] * 0 # adding zeros tensor
            inputs_embeds[0, :1] = inputs_embeds[0, :1] + other_embeds

class Granite4VisionForConditionalGenerationNaflex(Granite4VisionForConditionalGeneration):
    config_class = Granite4VisionConfigNaflex