optimum-rbln 0.8.4a8__py3-none-any.whl → 0.9.2__py3-none-any.whl

optimum/rbln/transformers/models/llava_next/modeling_llava_next.py

@@ -12,18 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import inspect
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
 from transformers import AutoModelForVision2Seq, LlavaNextForConditionalGeneration, PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.modeling_utils import no_init_weights
+from transformers.models.llava_next.modeling_llava_next import (
+    get_anyres_image_grid_shape,
+    image_size_to_num_patches,
+    unpad_image,
+)
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
 from ....utils.logging import get_logger
+from ...utils.rbln_runtime_wrapper import LoopProcessor
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 from ..decoderonly.modeling_decoderonly import RBLNDecoderOnlyOutput
 
 
@@ -33,33 +42,27 @@ if TYPE_CHECKING:
     from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig
 
 
-class LoopVisionTower:
-    def __init__(self, vision_tower: RBLNModel) -> None:
-        self.vision_tower = vision_tower
+class LoopVisionTower(LoopProcessor):
+    def __init__(self, vision_tower: "RBLNModel"):
+        super().__init__(model=vision_tower.model[0])
 
-    def forward(self, *args, **kwargs):
-        # Loop instead of batch
-        # shape of pixel_values : [batch, num_patches, num_channel, height, width]
-        pixel_values = args[0]
+    def _get_batch_size(self, pixel_values, **kwargs):
+        return pixel_values.shape[0]
 
-        batch_size = pixel_values.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.vision_tower.model[0](pixel_values[i : i + 1]))
+    def _prepare_inputs_for_iteration(self, index, common_inputs, pixel_values, **kwargs):
+        pixel_values_item = pixel_values[index : index + 1]
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
+        return ([pixel_values_item], {"out": out_buffer})
 
-        last_hidden_states = [output[0] for output in outputs]
-        pooler_output = [output[1] for output in outputs]
+    def _process_outputs(self, outputs: list, **kwargs) -> "BaseModelOutputWithPooling":
+        output = kwargs["out"]
+        last_hidden_states = output[0]
+        pooler_output = output[1]
 
-        # FIXME:: This can be optimized using out= API of rbln runtime.
-        last_hidden_states = torch.cat(last_hidden_states, dim=0)
-        pooler_output = torch.cat(pooler_output, dim=0)
-
-        hidden_states = [output[2:] for output in outputs]  # batch x (hidden x 1)
-
-        hidden_states = tuple(
-            torch.cat(tuple((hidden_states[n][i] for n in range(batch_size))), dim=0)
-            for i in range(len(hidden_states[0]))
-        )  # hidden x (batch,)
+        if not output[2:]:
+            hidden_states = None
+        else:
+            hidden_states = tuple(output[2:])
 
         return BaseModelOutputWithPooling(
             last_hidden_state=last_hidden_states,
@@ -67,38 +70,25 @@ class LoopVisionTower:
             hidden_states=hidden_states,
         )
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
-
-    def __repr__(self) -> str:
-        return repr(self.vision_tower)
-
 
-class LoopProjector:
-    def __init__(self, multi_modal_projector) -> None:
-        self.multi_modal_projector = multi_modal_projector
+class LoopProjector(LoopProcessor):
+    def __init__(self, multi_modal_projector: "RBLNModel"):
+        super().__init__(model=multi_modal_projector)
 
-    def forward(self, *args, **kwargs):
-        # Loop instead of batch
-        image_feature = args[0]
+    def _get_batch_size(self, image_feature, **kwargs):
+        return image_feature.shape[0]
 
-        batch_size = image_feature.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.multi_modal_projector(image_feature[i : i + 1]))
+    def _prepare_inputs_for_iteration(self, index, common_inputs, image_feature, **kwargs):
+        image_feature_item = image_feature[index : index + 1]
+        out_buffer = [tensor[index : index + 1] for tensor in kwargs["out"]]
+        return ([image_feature_item], {"out": out_buffer})
 
-        # FIXME:: This can be optimized using out= API of rbln runtime.
-        outputs = torch.cat(outputs, dim=0)
-        return outputs
+    def _process_outputs(self, outputs: list, **kwargs):
+        output = kwargs["out"]
+        return output[0]
 
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
 
-    def __repr__(self) -> str:
-        return repr(self.multi_modal_projector)
-
-
-class RBLNLlavaNextForConditionalGeneration(RBLNModel):
+class RBLNLlavaNextForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     RBLNLlavaNextForConditionalGeneration is a multi-modal model that combines vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
@@ -148,6 +138,23 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
     def can_generate(self):
         return True
 
+    @classmethod
+    def get_pytorch_model(cls, *args, **kwargs):
+        model = super().get_pytorch_model(*args, **kwargs)
+
+        with no_init_weights():
+            model_cls_name = model.model.language_model.__class__.__name__
+            causal_model_cls_name = model_cls_name.replace("Model", "ForCausalLM")
+            causal_model_cls = getattr(importlib.import_module("transformers"), causal_model_cls_name)
+            new_language_model = causal_model_cls(model.model.language_model.config)
+
+        new_language_model.lm_head = model.lm_head
+        new_language_model.model = model.model.language_model
+        model.model.language_model = new_language_model
+        model.lm_head = None
+        del model.lm_head
+        return model
+
     @classmethod
     def save_torch_artifacts(
         cls,
@@ -159,7 +166,7 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
         # If you are unavoidably running on a CPU rather than an RBLN device,
         # store the torch tensor, weight, etc. in this function.
         save_dict = {}
-        save_dict["image_newline"] = model.image_newline
+        save_dict["image_newline"] = model.model.image_newline
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     def __post_init__(self, **kwargs):
@@ -206,7 +213,11 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
             selected_image_feature_dim = num_positions
 
         input_info = [
-            ("image_features", [rbln_config.batch_size, selected_image_feature_dim, feature_size], "float32")
+            (
+                "image_features",
+                [rbln_config.vision_tower.batch_size, selected_image_feature_dim, feature_size],
+                "float32",
+            )
         ]
         rbln_compile_config = RBLNCompileConfig(input_info=input_info)
         rbln_config.set_compile_cfgs([rbln_compile_config])
@@ -217,86 +228,62 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
         input_ids,
         inputs_embeds=None,
         pixel_values=None,
-        image_sizes=None,
         attention_mask=None,
+        cache_position=None,
+        image_sizes=None,
         generate_idx=None,
         **kwargs,
     ):
-        # Prepare HF generation
         is_prefill_phase = generate_idx is None
-        batch_size = input_ids.shape[0]
-
-        model_inputs = self.language_model.prepare_inputs_for_generation(
-            input_ids=input_ids,
-            inputs_embeds=inputs_embeds,
-            generate_idx=generate_idx,  # Not affect
-            attention_mask=attention_mask,
-            **kwargs,
-        )
+        model_inputs = {}
 
         if is_prefill_phase:
-            model_inputs["generate_idx"] = torch.zeros((batch_size, 1), dtype=torch.int32)
-            model_inputs.update(
-                {
-                    "pixel_values": pixel_values,
-                    "image_sizes": image_sizes,
-                }
-            )
+            generate_idx = attention_mask.sum(dim=-1, keepdim=True).int()
+            cache_position = None
+            pixel_values = pixel_values
+            model_inputs.update({"image_sizes": image_sizes})
+        else:
+            if inputs_embeds is not None:
+                raise NotImplementedError("Specifying inputs_embeds in decoder phase is not supported.")
 
-        model_inputs["attention_mask"] = attention_mask
+            pixel_values = None
+            input_ids = input_ids[:, -1:]
+            cache_position = generate_idx
+            generate_idx = generate_idx + 1
+            model_inputs.update({"input_ids": input_ids})
+
+        if inputs_embeds is not None:
+            if self.rbln_config.use_inputs_embeds:
+                model_inputs.update({"inputs_embeds": inputs_embeds})
+            else:
+                raise ValueError(
+                    "The specifying inputs_embeds is only supported when using a compiled RBLN model with 'rbln_use_inputs_embeds' set to True."
+                )
+        else:
+            model_inputs.update({"input_ids": input_ids})
+
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "pixel_values": pixel_values,
+                "cache_position": cache_position,
+                "generate_idx": generate_idx,
+            }
+        )
         return model_inputs
 
-    def _update_model_kwargs_for_generation(
-        self,
-        outputs: RBLNDecoderOnlyOutput,
-        model_kwargs: Dict[str, Any],
-        **kwargs,
-    ) -> Dict[str, Any]:
+    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder, **kwargs):
         # update generate_idx
         model_kwargs["generate_idx"] = outputs.generate_idx
-
         return model_kwargs
 
-    def text_embedding(
+    def get_image_features(
         self,
-        input_ids: torch.LongTensor,
-    ) -> torch.Tensor:
-        for_inputs_embeds_ids = input_ids.clone()
-        for_inputs_embeds_ids[(input_ids == self.config.image_token_index)] = 0
-        inputs_embeds = self.get_input_embeddings()(for_inputs_embeds_ids)
-
-        return inputs_embeds
-
-    def image_embedding(
-        self,
-        image_sizes: torch.Tensor,
         pixel_values: torch.FloatTensor,
-        vision_feature_layer: int,
+        image_sizes: torch.Tensor,
+        vision_feature_layer: Union[int, List[int]],
         vision_feature_select_strategy: str,
     ):
-        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
-        )
-
-        """
-        Obtains image last hidden states from the vision tower and apply multimodal projection.
-
-        Args:
-            pixel_values (torch.FloatTensor): The tensors corresponding to the input images
-                whose shape is `(batch_size, num_patches, channels, height, width)`.
-            image_sizes (torch.Tensor): Actual image size of each images (H, W).
-            vision_feature_layer (int): The index of the layer to select the vision feature.
-            vision_feature_select_strategy (str): 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)`).
-        """
         # ! infer image_num_patches from image_sizes
         image_num_patches = [
             image_size_to_num_patches(
@@ -306,6 +293,26 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
             )
             for imsize in image_sizes
         ]
+
+        # prepare out buffer for pre-allocation
+        vision_out_size = [
+            pixel_values.shape[0] * pixel_values.shape[1],
+            (self.config.vision_config.image_size // self.config.vision_config.patch_size) ** 2 + 1,
+            self.config.vision_config.hidden_size,
+        ]
+        pooler_out_size = [pixel_values.shape[0] * pixel_values.shape[1], self.config.vision_config.hidden_size]
+        vision_out_buffer = []
+        for i in range(self.config.vision_config.num_hidden_layers + 2):
+            vision_out_buffer.append(torch.empty(size=vision_out_size, dtype=torch.float32, device="cpu"))
+        vision_out_buffer.insert(1, torch.empty(size=pooler_out_size, dtype=torch.float32, device="cpu"))
+
+        projector_out_size = [
+            pixel_values.shape[0] * pixel_values.shape[1],
+            (self.config.vision_config.image_size // self.config.vision_config.patch_size) ** 2,
+            self.config.text_config.hidden_size,
+        ]
+        projector_out_buffer = [torch.empty(size=projector_out_size, dtype=torch.float32, device="cpu")]
+
         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)]
@@ -314,114 +321,25 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
             # 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)
-        selected_image_feature = image_features.hidden_states[vision_feature_layer]
+        image_features = self.vision_tower(pixel_values, output_hidden_states=True, out=vision_out_buffer)
+        # 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:]
         elif vision_feature_select_strategy == "full":
             selected_image_feature = selected_image_feature
-        image_features = self.multi_modal_projector(selected_image_feature)
-        image_features = torch.split(image_features, image_num_patches, dim=0)
-
-        # NOTE we only support multimodal_patch_merge_type == "spatial_unpad"
-        image_features, feature_lens = self.pack_image_features(
-            image_features,
-            image_sizes,
-            vision_feature_select_strategy=vision_feature_select_strategy,
-            image_newline=self.image_newline,
-        )
-
-        return image_features, feature_lens
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: torch.LongTensor = None,
-        pixel_values: torch.FloatTensor = None,
-        image_sizes: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        vision_feature_layer: Optional[int] = None,
-        vision_feature_select_strategy: Optional[str] = None,
-        cache_position: torch.Tensor = None,
-        generate_idx: Optional[torch.Tensor] = None,
-        batch_idx: Optional[int] = None,
-        **kwargs,
-    ) -> Union[Tuple, RBLNDecoderOnlyOutput]:
-        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 inputs_embeds is not None:
-            raise NotImplementedError("Specifying inputs_embeds is not supported.")
-        inputs_embeds = self.get_input_embeddings()(input_ids)
-
-        if pixel_values is not None and pixel_values.size(0) > 0:
-            image_features, _ = self.image_embedding(
-                pixel_values=pixel_values,
-                image_sizes=image_sizes,
-                vision_feature_layer=vision_feature_layer,
-                vision_feature_select_strategy=vision_feature_select_strategy,
-            )
-
-            n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
-            n_image_features = image_features.shape[0]
-            if n_image_tokens != n_image_features:
-                raise ValueError(
-                    f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
-                )
-            special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
-            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
-            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
-            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
-
-        is_prefill_phase = not generate_idx.bool().all()
-
-        if is_prefill_phase:
-            logits = []
-            batch_size = input_ids.shape[0]
-            inputs_embeds = [inputs_embeds[i : i + 1, attention_mask[i].bool()] for i in range(batch_size)]
-            for batch_idx in range(batch_size):
-                generate_idx[batch_idx] = inputs_embeds[batch_idx].shape[-2]
-                output = self.language_model.prefill_decoder(
-                    inputs_embeds=inputs_embeds[batch_idx],
-                    batch_idx=batch_idx,
-                    cache_position=torch.arange(
-                        0,
-                        generate_idx[batch_idx].item(),
-                        dtype=torch.int32,
-                    ).unsqueeze(0),
-                )
 
-                logits.append(output.logits)
-            logits = torch.cat(logits, dim=0)
-        else:
-            output = self.language_model.decoder(
-                inputs_embeds=inputs_embeds,
-                cache_position=cache_position,
-            )
-            logits = output.logits
-        return RBLNDecoderOnlyOutput(logits=logits, generate_idx=generate_idx)
+        image_features = self.multi_modal_projector(selected_image_feature, out=projector_out_buffer)
+        image_features = torch.split(image_features, image_num_patches, dim=0)
+        return image_features
 
-    # Almost copied from : https://github.com/huggingface/transformers/blob/6b550462139655d488d4c663086a63e98713c6b9/src/transformers/models/llava_next/modeling_llava_next.py
     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]): List of image feature tensor, each contains all the visual feature of all patches.
-        #         Its length is num_images, and each of shape is `(num_patches, image_length, embed_dim)`
-        #     image_sizes (torch.Tensor): 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): New line embedding vector whose shape is `embed_dim`.
-
-        # Returns:
-        #     image_features (torch.Tensor): A torch.Tensor of shape `(all_feat_len, embed_dim)`)
-        #     feature_lens (List[int]): A token length of each image in image_features
-
         new_image_features = []
         feature_lens = []
         for image_idx, image_feature in enumerate(image_features):
@@ -430,18 +348,22 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
                 image_feature = image_feature[1:]
                 height = width = self.config.vision_config.image_size // self.config.vision_config.patch_size
 
-                if vision_feature_select_strategy == "default":
-                    expected_num_patches = height * width
-                elif vision_feature_select_strategy == "full":
-                    expected_num_patches = height * width + 1
-                if expected_num_patches != base_image_feature.shape[0]:
-                    raise ValueError("The number of patches is not consistent with the image 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 (
+                    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)
@@ -468,137 +390,106 @@ class RBLNLlavaNextForConditionalGeneration(RBLNModel):
         feature_lens = torch.tensor(feature_lens, dtype=torch.long, device=image_features.device)
         return image_features, feature_lens
 
+    def _preprocess_prefill(
+        self,
+        input_ids: torch.LongTensor = None,
+        pixel_values: torch.FloatTensor = None,
+        image_sizes: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        vision_feature_layer: Optional[int] = None,
+        vision_feature_select_strategy: Optional[str] = None,
+        **kwargs,
+    ):
+        vision_feature_layer = (
+            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
+        )
 
-# Almost copied from : https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/llava_next/modeling_llava_next.py
-def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
-    # Calculate the shape of the image patch grid after the preprocessing for images of any resolution.
+        vision_feature_select_strategy = (
+            vision_feature_select_strategy
+            if vision_feature_select_strategy is not None
+            else self.config.vision_feature_select_strategy
+        )
 
-    # Args:
-    #     image_size (tuple): The size of the input image in the format (width, height).
-    #     grid_pinpoints (list): A list containing possible resolutions.
-    #         Each item in the list should be a tuple or list of the form `(height, width)`.
-    #     patch_size (int): The size of each image patch.
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
-    # Returns:
-    #     tuple: The shape of the image patch grid in the format (width, height).
+        if pixel_values is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
+            )
 
-    if not isinstance(grid_pinpoints, list):
-        raise TypeError("grid_pinpoints should be a list of tuples or lists")
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
 
-    # ! VERY IMPORTANT if image_size is tensor, must convert to into tuple, otherwise it will cause wrong calculate
-    if not isinstance(image_size, (list, tuple)):
-        if not isinstance(image_size, (torch.Tensor, np.ndarray)):
-            raise TypeError(
-                f"image_size invalid type: {type(image_size)} not valid, should be either list, tuple, np.ndarray or tensor"
+        if pixel_values is not None and pixel_values.size(0) > 0:
+            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_size = image_size.tolist()
 
-    height, width = select_best_resolution(image_size, grid_pinpoints)
-    return height // patch_size, width // patch_size
+            # NOTE we only support multimodal_patch_merge_type == "spatial_unpad"
+            image_features, feature_lens = self.pack_image_features(
+                image_features,
+                image_sizes,
+                vision_feature_select_strategy=vision_feature_select_strategy,
+                image_newline=self.image_newline,
+            )
 
+            special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
+            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)
+            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
 
-# Almost copied from : https://github.com/huggingface/transformers/blob/1feebb5b4150882deabddd190a541f336f3be817/src/transformers/models/llava_next/modeling_llava_next.py#L115C1-L152C1
-def unpad_image(tensor, original_size):
-    # Unpads a PyTorch tensor of a padded and resized image.
+        return inputs_embeds
 
-    # Args:
-    #     tensor (torch.Tensor): The image tensor, assumed to be of shape (num_channels, height, width).
-    #     original_size (tuple): The original size of the image (height, width).
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.LongTensor = None,
+        pixel_values: torch.FloatTensor = None,
+        image_sizes: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: torch.Tensor = None,
+        generate_idx: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, RBLNDecoderOnlyOutput]:
+        # Prefill
+        if cache_position is None:
+            inputs_embeds = self._preprocess_prefill(
+                input_ids=input_ids, inputs_embeds=inputs_embeds, pixel_values=pixel_values, image_sizes=image_sizes
+            )
+            logits = []
+            inputs = inputs_embeds if inputs_embeds is not None else input_ids
+            batch_size = inputs.shape[0]
 
-    # Returns:
-    #     (torch.Tensor): The unpadded image tensor.
+            for b_idx in range(batch_size):
+                cache_position = torch.arange(0, generate_idx[b_idx].item(), dtype=torch.int32).unsqueeze(0)
+                output = self.language_model.prefill_decoder(
+                    input_ids=inputs[b_idx : b_idx + 1] if inputs_embeds is None else None,
+                    inputs_embeds=inputs[b_idx : b_idx + 1] if inputs_embeds is not None else None,
+                    attention_mask=attention_mask[b_idx] if attention_mask is not None else None,
+                    cache_position=cache_position,
+                    batch_idx=b_idx,
+                )
+                logits.append(output.logits)
 
-    if not isinstance(original_size, (list, tuple)):
-        if not isinstance(original_size, (torch.Tensor, np.ndarray)):
-            raise TypeError(
-                f"image_size invalid type: {type(original_size)} not valid, should be either list, tuple, np.ndarray or tensor"
+            logits = torch.cat(logits, dim=0)
+
+        # Decoder
+        else:
+            logits = self.language_model.decoder(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                cache_position=cache_position,
+            ).logits
+
+        if not return_dict:
+            return logits, generate_idx
+        else:
+            return RBLNDecoderOnlyOutput(
+                logits=logits,
+                generate_idx=generate_idx,
             )
-        original_size = original_size.tolist()
-    original_height, original_width = 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(round(original_height * scale_factor, 7))
-        padding = (current_height - new_height) // 2
-        unpadded_tensor = tensor[:, padding : current_height - padding, :]
-    else:
-        scale_factor = current_height / original_height
-        new_width = int(round(original_width * scale_factor, 7))
-        padding = (current_width - new_width) // 2
-        unpadded_tensor = tensor[:, :, padding : current_width - padding]
-
-    return unpadded_tensor
-
-
-# Almost copied from : https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/llava_next/modeling_llava_next.py
-def select_best_resolution(original_size: tuple, possible_resolutions: list) -> tuple:
-    # Selects the best resolution from a list of possible resolutions based on the original size.
-
-    # This is done by calculating the effective and wasted resolution for each possible resolution.
-
-    # The best fit resolution is the one that maximizes the effective resolution and minimizes the wasted resolution.
-
-    # Args:
-    #     original_size (tuple): The original size of the image in the format (height, width).
-    #     possible_resolutions (List(tuple)): A list of possible resolutions in the format [(height1, width1), (height2, width2), ...].
-
-    # Returns:
-    #     (tuple): The best fit resolution in the format (height, width).
-
-    original_height, original_width = original_size
-    best_fit = None
-    max_effective_resolution = 0
-    min_wasted_resolution = float("inf")
-
-    for height, width in possible_resolutions:
-        scale = min(width / original_width, height / original_height)
-        downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)
-        effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)
-        wasted_resolution = (width * height) - effective_resolution
-
-        if effective_resolution > max_effective_resolution or (
-            effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution
-        ):
-            max_effective_resolution = effective_resolution
-            min_wasted_resolution = wasted_resolution
-            best_fit = (height, width)
-
-    return best_fit
-
-
-# Almost copied from : https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/llava_next/modeling_llava_next.py
-def image_size_to_num_patches(image_size, grid_pinpoints, patch_size: int):
-    # Calculate the number of patches after the preprocessing for images of any resolution.
-
-    # Args:
-    #     image_size (Union[torch.LongTensor, np.ndarray, Tuple[int, int]): The size of the input image in the format (height, width).
-    #     grid_pinpoints (list): A list containing possible resolutions.
-    #         Each item in the list should be a tuple or list of the form `(height, width)`.
-    #     patch_size (int): The size of each image patch.
-
-    # Returns:
-    #     (int): the number of patches.
-
-    if not isinstance(grid_pinpoints, list):
-        raise TypeError("grid_pinpoints should be a list of tuples or lists")
-
-    # ! VERY IMPORTANT if image_size is tensor, must convert to into tuple, otherwise it will cause wrong calculate
-    if not isinstance(image_size, (list, tuple)):
-        if not isinstance(image_size, (torch.Tensor, np.ndarray)):
-            raise TypeError(f"image_size invalid type {type(image_size)} with value {image_size}")
-        image_size = image_size.tolist()
-
-    best_resolution = select_best_resolution(image_size, grid_pinpoints)
-    height, width = best_resolution
-    num_patches = 0
-    # consider change to ceil(height/patch_size)*ceil(width/patch_size) + 1
-    for i in range(0, height, patch_size):
-        for j in range(0, width, patch_size):
-            num_patches += 1
-    # add the base patch
-    num_patches += 1
-    return num_patches