optimum-rbln 0.8.2a4__py3-none-any.whl → 0.9.3rc0__py3-none-any.whl

optimum/rbln/transformers/models/auto/auto_factory.py

@@ -14,13 +14,14 @@
 import importlib
 import inspect
 import warnings
-from typing import Type
+from pathlib import Path
+from typing import Any, Dict, Optional, Type, Union
 
-from transformers import AutoConfig, PretrainedConfig
+from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
 from transformers.models.auto.auto_factory import _get_model_class
 
-from optimum.rbln.configuration_utils import RBLNAutoConfig
+from optimum.rbln.configuration_utils import RBLNAutoConfig, RBLNModelConfig
 from optimum.rbln.modeling_base import RBLNBaseModel
 from optimum.rbln.utils.model_utils import (
     MODEL_MAPPING,
@@ -43,10 +44,10 @@ class _BaseAutoModelClass:
     @classmethod
     def get_rbln_cls(
         cls,
-        pretrained_model_name_or_path,
-        *args,
-        export=True,
-        **kwargs,
+        pretrained_model_name_or_path: Union[str, Path],
+        *args: Any,
+        export: bool = None,
+        **kwargs: Any,
     ):
         """
         Determine the appropriate RBLN model class based on the given model ID and configuration.
@@ -59,6 +60,20 @@ class _BaseAutoModelClass:
         Returns:
             RBLNBaseModel: The corresponding RBLN model class.
         """
+        if isinstance(pretrained_model_name_or_path, Path):
+            pretrained_model_name_or_path = pretrained_model_name_or_path.as_posix()
+
+        if export is None:
+            export = not RBLNBaseModel._is_compiled(
+                model_id=pretrained_model_name_or_path,
+                token=kwargs.get("token"),
+                revision=kwargs.get("revision"),
+                force_download=kwargs.get("force_download", False),
+                cache_dir=kwargs.get("cache_dir"),
+                subfolder=kwargs.get("subfolder", ""),
+                local_files_only=kwargs.get("local_files_only", False),
+            )
+
         if export:
             hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
             rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
@@ -85,9 +100,9 @@ class _BaseAutoModelClass:
     @classmethod
     def infer_hf_model_class(
         cls,
-        pretrained_model_name_or_path,
-        *args,
-        **kwargs,
+        pretrained_model_name_or_path: Union[str, Path],
+        *args: Any,
+        **kwargs: Any,
     ):
         """
         Infer the HuggingFace model class based on the configuration or model name.
@@ -140,7 +155,7 @@ class _BaseAutoModelClass:
         return model_class
 
     @classmethod
-    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path, **kwargs):
+    def get_rbln_model_cls_name(cls, pretrained_model_name_or_path: Union[str, Path], **kwargs):
         """
         Retrieve the path to the compiled model directory for a given RBLN model.
 
@@ -163,17 +178,77 @@ class _BaseAutoModelClass:
         return rbln_config.rbln_model_cls_name
 
     @classmethod
-    def from_pretrained(cls, model_id, *args, **kwargs):
-        rbln_cls = cls.get_rbln_cls(model_id, *args, **kwargs)
-        return rbln_cls.from_pretrained(model_id, *args, **kwargs)
+    def from_pretrained(
+        cls,
+        model_id: Union[str, Path],
+        export: bool = None,
+        rbln_config: Optional[Union[Dict, RBLNModelConfig]] = None,
+        **kwargs,
+    ):
+        """
+        Load an RBLN-accelerated model from a pretrained checkpoint or a compiled RBLN artifact.
+
+        This convenience method determines the concrete `RBLN*` model class that matches the
+        underlying HuggingFace architecture and dispatches to that class's
+        `from_pretrained()` implementation. Depending on whether a compiled RBLN folder is
+        detected (or if `export=True` is passed), it will either:
+
+        - Compile from a HuggingFace checkpoint to an RBLN model
+        - Or load an already-compiled RBLN model directory/repository
+
+        Args:
+            model_id:
+                HF repo id or local path. For compiled models, this should point to a directory
+                (optionally under `subfolder`) that contains `*.rbln` files and `rbln_config.json`.
+            export:
+                Force compilation from a HuggingFace checkpoint. When `None`, this is inferred by
+                checking whether compiled artifacts exist at `model_id`.
+            rbln_config:
+                RBLN compilation/runtime configuration. May be provided as a dictionary or as an
+                instance of the specific model's config class (e.g., `RBLNLlamaForCausalLMConfig`).
+            kwargs: Additional keyword arguments.
+                - Arguments prefixed with `rbln_` are forwarded to the RBLN config.
+                - Remaining arguments are forwarded to the HuggingFace loader (e.g., `revision`,
+                  `token`, `trust_remote_code`, `cache_dir`, `subfolder`, `local_files_only`).
+
+        Returns:
+            An instantiated RBLN model ready for inference on RBLN NPUs.
+        """
+        rbln_cls = cls.get_rbln_cls(model_id, export=export, **kwargs)
+        return rbln_cls.from_pretrained(model_id, export=export, rbln_config=rbln_config, **kwargs)
 
     @classmethod
-    def from_model(cls, model, *args, **kwargs):
+    def from_model(
+        cls,
+        model: PreTrainedModel,
+        config: Optional[PretrainedConfig] = None,
+        rbln_config: Optional[Union[RBLNModelConfig, Dict]] = None,
+        **kwargs: Any,
+    ) -> RBLNBaseModel:
+        """
+        Convert and compile an in-memory HuggingFace model into an RBLN model.
+
+        This method resolves the appropriate concrete `RBLN*` class from the input model's class
+        name (e.g., `LlamaForCausalLM` -> `RBLNLlamaForCausalLM`) and then delegates to that
+        class's `from_model()` implementation.
+
+        Args:
+            model: A HuggingFace model instance to convert.
+            config: The configuration object associated with the model.
+            rbln_config:
+                RBLN compilation/runtime configuration. May be provided as a dictionary or as an
+                instance of the specific model's config class.
+            kwargs: Additional keyword arguments.
+                - Arguments prefixed with `rbln_` are forwarded to the RBLN config.
+
+        Returns:
+            An instantiated RBLN model ready for inference on RBLN NPUs.
+        """
         rbln_cls = get_rbln_model_cls(f"RBLN{model.__class__.__name__}")
-        return rbln_cls.from_model(model, *args, **kwargs)
+        return rbln_cls.from_model(model, config=config, rbln_config=rbln_config, **kwargs)
 
     @staticmethod
-    def register(rbln_cls: Type[RBLNBaseModel], exist_ok=False):
+    def register(rbln_cls: Type[RBLNBaseModel], exist_ok: bool = False):
         """
         Register a new RBLN model class.
 

optimum/rbln/transformers/models/auto/modeling_auto.py

@@ -35,8 +35,12 @@ from transformers.models.auto.modeling_auto import (
     MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES,
     MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING,
     MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES,
+    MODEL_FOR_TEXT_ENCODING_MAPPING,
+    MODEL_FOR_TEXT_ENCODING_MAPPING_NAMES,
     MODEL_FOR_VISION_2_SEQ_MAPPING,
     MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES,
+    MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING,
+    MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES,
     MODEL_MAPPING,
     MODEL_MAPPING_NAMES,
 )
@@ -53,65 +57,106 @@ MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.update(
 
 
 class RBLNAutoModel(_BaseAutoModelClass):
+    """Automatically detect all supported transformers models."""
+
     _model_mapping = MODEL_MAPPING
     _model_mapping_names = MODEL_MAPPING_NAMES
 
 
 class RBLNAutoModelForCTC(_BaseAutoModelClass):
+    """Automatically detect Connectionist Temporal Classification (CTC) head Models."""
+
     _model_mapping = MODEL_FOR_CTC_MAPPING
     _model_mapping_names = MODEL_FOR_CTC_MAPPING_NAMES
 
 
 class RBLNAutoModelForCausalLM(_BaseAutoModelClass):
+    """Automatically detect Casual Language Models."""
+
+    """"""
     _model_mapping = MODEL_FOR_CAUSAL_LM_MAPPING
     _model_mapping_names = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForSeq2SeqLM(_BaseAutoModelClass):
+    """Automatically detect Sequence to Sequence Language Models."""
+
     _model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
     _model_mapping_names = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForSpeechSeq2Seq(_BaseAutoModelClass):
+    """Automatically detect Sequence to Sequence Generation Models."""
+
     _model_mapping = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING
     _model_mapping_names = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
 
 
 class RBLNAutoModelForDepthEstimation(_BaseAutoModelClass):
+    """Automatically detect Speech Sequence to Sequence Language Models."""
+
     _model_mapping = MODEL_FOR_DEPTH_ESTIMATION_MAPPING
     _model_mapping_names = MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForSequenceClassification(_BaseAutoModelClass):
+    """Automatically detect Sequence Classification Models."""
+
     _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
     _model_mapping_names = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForVision2Seq(_BaseAutoModelClass):
+    """Automatically detect Vision to Sequence Generation Models."""
+
     _model_mapping = MODEL_FOR_VISION_2_SEQ_MAPPING
     _model_mapping_names = MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES
 
 
 class RBLNAutoModelForImageTextToText(_BaseAutoModelClass):
+    """Automatically detect Image and Text to Text Generation Models."""
+
     _model_mapping = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING
     _model_mapping_names = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
 
 
 class RBLNAutoModelForMaskedLM(_BaseAutoModelClass):
+    """Automatically detect Masked Lanuage Models."""
+
     _model_mapping = MODEL_FOR_MASKED_LM_MAPPING
     _model_mapping_names = MODEL_FOR_MASKED_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForAudioClassification(_BaseAutoModelClass):
+    """Automatically detect Audio Classification Models."""
+
     _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
     _model_mapping_names = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForImageClassification(_BaseAutoModelClass):
+    """Automatically detect Image Classification Models."""
+
     _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
     _model_mapping_names = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForQuestionAnswering(_BaseAutoModelClass):
+    """Automatically detect Question Answering Models."""
+
     _model_mapping = MODEL_FOR_QUESTION_ANSWERING_MAPPING
     _model_mapping_names = MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
+
+
+class RBLNAutoModelForTextEncoding(_BaseAutoModelClass):
+    """Automatically detect Text Encoding Models."""
+
+    _model_mapping = MODEL_FOR_TEXT_ENCODING_MAPPING
+    _model_mapping_names = MODEL_FOR_TEXT_ENCODING_MAPPING_NAMES
+
+
+class RBLNAutoModelForZeroShotObjectDetection(_BaseAutoModelClass):
+    """Automatically detect Zero Shot Object Detection Models."""
+
+    _model_mapping = MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING
+    _model_mapping_names = MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES

optimum/rbln/transformers/models/bart/bart_architecture.py

@@ -16,9 +16,7 @@ from typing import Tuple
 
 import torch
 from torch import nn
-from transformers.modeling_attn_mask_utils import (
-    _prepare_4d_attention_mask,
-)
+from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
 from transformers.utils import logging
 
 from ..seq2seq.seq2seq_architecture import (

optimum/rbln/transformers/models/bart/configuration_bart.py

@@ -32,3 +32,5 @@ class RBLNBartForConditionalGenerationConfig(RBLNModelForSeq2SeqLMConfig):
     This configuration class stores the configuration parameters specific to
     RBLN-optimized BART models for conditional text generation tasks.
     """
+
+    support_paged_attention = True

optimum/rbln/transformers/models/bert/bert_architecture.py

@@ -0,0 +1,16 @@
+import torch
+
+
+class BertModelWrapper(torch.nn.Module):
+    def __init__(self, model, rbln_config):
+        super().__init__()
+        self.model = model
+        self.rbln_config = rbln_config
+
+    def forward(self, *args, **kwargs):
+        output = self.model(*args, **kwargs)
+        if isinstance(output, torch.Tensor):
+            return output
+        elif isinstance(output, tuple):
+            return tuple(x for x in output if x is not None)
+        return output

optimum/rbln/transformers/models/bert/modeling_bert.py

@@ -12,15 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ....utils.logging import get_logger
+import torch
+
 from ...modeling_generic import (
     RBLNModelForMaskedLM,
     RBLNModelForQuestionAnswering,
     RBLNTransformerEncoderForFeatureExtraction,
 )
-
-
-logger = get_logger(__name__)
+from .bert_architecture import BertModelWrapper
+from .configuration_bert import RBLNBertModelConfig
 
 
 class RBLNBertModel(RBLNTransformerEncoderForFeatureExtraction):
@@ -34,6 +34,10 @@ class RBLNBertModel(RBLNTransformerEncoderForFeatureExtraction):
 
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
+    @classmethod
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNBertModelConfig) -> torch.nn.Module:
+        return BertModelWrapper(model, rbln_config)
+
 
 class RBLNBertForMaskedLM(RBLNModelForMaskedLM):
     """

optimum/rbln/transformers/models/blip_2/configuration_blip_2.py

@@ -12,9 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
+from ....utils.logging import get_logger
+
+
+logger = get_logger(__name__)
 
 
 class RBLNBlip2VisionModelConfig(RBLNModelConfig):
@@ -25,6 +29,16 @@ class RBLNBlip2VisionModelConfig(RBLNModelConfig):
     RBLN-optimized BLIP-2 vision encoder models for multimodal tasks.
     """
 
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
 
 class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
     """
@@ -36,24 +50,34 @@ class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
 
     def __init__(
         self,
+        batch_size: Optional[int] = None,
         num_query_tokens: Optional[int] = None,
         image_text_hidden_size: Optional[int] = None,
         **kwargs,
     ):
         """
         Args:
-            batch_size (Optional[int]): The batch size for inference. Defaults to 1.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
-
-        Raises:
-            ValueError: If batch_size is not a positive integer.
+            num_query_tokens (Optional[int]): The number of query tokens passed through the Transformer.
+            image_text_hidden_size (Optional[int]): Dimensionality of the hidden state of the image-text fusion layer.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
         """
         super().__init__(**kwargs)
+        self.batch_size = batch_size or 1
+        if not isinstance(self.batch_size, int) or self.batch_size < 0:
+            raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
+
         self.num_query_tokens = num_query_tokens
         self.image_text_hidden_size = image_text_hidden_size
 
 
 class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLNBlip2ForConditionalGeneration.
+
+    This configuration class stores the configuration parameters specific to
+    RBLN-optimized BLIP-2 models for conditional generation tasks that involve both image and text inputs.
+    """
+
     submodules = ["vision_model", "qformer", "language_model"]
 
     def __init__(
@@ -62,14 +86,15 @@ class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
         vision_model: Optional[RBLNModelConfig] = None,
         qformer: Optional[RBLNModelConfig] = None,
         language_model: Optional[RBLNModelConfig] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
             vision_model (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            qformer (Optional[RBLNModelConfig]): Configuration for the RBLN-optimized BLIP-2 Q-Former model.
             language_model (Optional[RBLNModelConfig]): Configuration for the language model component.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If batch_size is not a positive integer.
@@ -79,6 +104,12 @@ class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
         if not isinstance(self.batch_size, int) or self.batch_size < 0:
             raise ValueError(f"batch_size must be a positive integer, got {self.batch_size}")
 
-        self.vision_model = self.init_submodule_config(RBLNBlip2VisionModelConfig, vision_model)
-        self.language_model = language_model
-        self.qformer = self.init_submodule_config(RBLNBlip2QFormerModelConfig, qformer)
+        if self.batch_size != 1:
+            logger.warning("Ignore batch_size for Blip2 vision model. It will be set to 1.")
+            logger.warning("Ignore batch_size for Blip2 qformer. It will be set to 1.")
+
+        self.vision_model = self.initialize_submodule_config(
+            submodule_config=vision_model, batch_size=1, force_kwargs=True
+        )
+        self.qformer = self.initialize_submodule_config(submodule_config=qformer, batch_size=1, force_kwargs=True)
+        self.language_model = self.initialize_submodule_config(submodule_config=language_model)

optimum/rbln/transformers/models/blip_2/modeling_blip_2.py

@@ -30,38 +30,31 @@ from transformers.utils import logging
 
 from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
+from ...utils.rbln_runtime_wrapper import LoopProcessor
+from ..decoderonly.generation_decoderonly import RBLNDecoderOnlyGenerationMixin
 
 
 logger = logging.get_logger(__name__)
 
 if TYPE_CHECKING:
-    from transformers import (
-        AutoFeatureExtractor,
-        AutoProcessor,
-        AutoTokenizer,
-    )
+    import rebel
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
 
 
-class LoopProjector:
-    def __init__(self, language_projection) -> None:
-        self.language_projection = language_projection
+class LoopProjector(LoopProcessor):
+    def __init__(self, language_projection: Union[RBLNModel, "rebel.Runtime"]):
+        super().__init__(model=language_projection)
 
-    def forward(self, *args, **kwargs):
-        query_output = args[0]
+    def _get_batch_size(self, query_output, **kwargs):
+        return query_output.shape[0]
 
-        batch_size = query_output.shape[0]
-        outputs = []
-        for i in range(batch_size):
-            outputs.append(self.language_projection(query_output[i : i + 1]))
-
-        outputs = torch.cat(outputs, dim=0)
-        return outputs
-
-    def __call__(self, *args: Any, **kwds: Any) -> Any:
-        return self.forward(*args, **kwds)
+    def _prepare_inputs_for_iteration(self, index, common_inputs, query_output, **kwargs):
+        query_output_item = query_output[index : index + 1]
+        return ([query_output_item], {})
 
-    def __repr__(self) -> str:
-        return repr(self.language_projection)
+    def _process_outputs(self, outputs: list, **kwargs):
+        output = torch.cat(outputs, dim=0)
+        return output
 
 
 class RBLNBlip2VisionModel(RBLNModel):
@@ -72,6 +65,8 @@ class RBLNBlip2VisionModel(RBLNModel):
     on RBLN devices, supporting image encoding for multimodal vision-language tasks.
     """
 
+    _tp_support = False
+
     def get_input_embeddings(self):
         return self.embeddings
 
@@ -100,8 +95,7 @@ class RBLNBlip2VisionModel(RBLNModel):
             (
                 "pixel_values",
                 [
-                    # support for vllm CB (prefill)
-                    1,
+                    rbln_config.batch_size,
                     model_config.num_channels,
                     model_config.image_size,
                     model_config.image_size,
@@ -116,7 +110,7 @@ class RBLNBlip2VisionModel(RBLNModel):
 
     def forward(
         self,
-        pixel_values,
+        pixel_values: torch.FloatTensor,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
@@ -151,6 +145,8 @@ class RBLNBlip2QFormerModel(RBLNModel):
     mechanisms for multimodal understanding tasks.
     """
 
+    _tp_support = False
+
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
 
@@ -178,7 +174,12 @@ class RBLNBlip2QFormerModel(RBLNModel):
         return Blip2QFormerModelWrapper(model).eval()
 
     @classmethod
-    def _update_submodule_config(cls, model: "PreTrainedModel", rbln_config: "RBLNModelConfig") -> "RBLNModelConfig":
+    def _update_submodule_config(
+        cls,
+        model: "PreTrainedModel",
+        rbln_config: RBLNModelConfig,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+    ):
         if rbln_config.num_query_tokens is None:
             rbln_config.num_query_tokens = model.config.num_query_tokens
 
@@ -199,7 +200,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "query_embeds",
                 [
-                    1,
+                    rbln_config.batch_size,
                     rbln_config.num_query_tokens,
                     model_config.hidden_size,
                 ],
@@ -208,7 +209,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "encoder_hidden_states",
                 [
-                    1,
+                    rbln_config.batch_size,
                     # image_text_hidden_size + cls token
                     rbln_config.image_text_hidden_size + 1,
                     model_config.encoder_hidden_size,
@@ -218,7 +219,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
             (
                 "encoder_attention_mask",
                 # image_text_hidden_size + cls token
-                [1, rbln_config.image_text_hidden_size + 1],
+                [rbln_config.batch_size, rbln_config.image_text_hidden_size + 1],
                 "int64",
             ),
         ]
@@ -265,7 +266,7 @@ class RBLNBlip2QFormerModel(RBLNModel):
         )
 
 
-class RBLNBlip2ForConditionalGeneration(RBLNModel):
+class RBLNBlip2ForConditionalGeneration(RBLNModel, RBLNDecoderOnlyGenerationMixin):
     """
     RBLNBlip2ForConditionalGeneration is a multi-modal model that integrates vision and language processing capabilities,
     optimized for RBLN NPUs. It is designed for conditional generation tasks that involve both image and text inputs.
@@ -433,3 +434,66 @@ class RBLNBlip2ForConditionalGeneration(RBLNModel):
             )
 
         return inputs_embeds
+
+    @torch.no_grad()
+    def generate(
+        self,
+        pixel_values: torch.FloatTensor,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        interpolate_pos_encoding: bool = False,
+        **generate_kwargs,
+    ) -> torch.LongTensor:
+        batch_size = pixel_values.shape[0]
+        image_embeds = self.vision_model(
+            pixel_values,
+            return_dict=True,
+            interpolate_pos_encoding=interpolate_pos_encoding,
+        ).last_hidden_state
+        image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
+
+        query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
+        query_outputs = self.qformer(
+            query_embeds=query_tokens,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_attention_mask,
+            return_dict=True,
+        )
+        query_output = query_outputs.last_hidden_state
+
+        if query_output.dtype != image_embeds.dtype:
+            query_output = query_output.to(image_embeds.dtype)
+
+        language_model_inputs = self.language_projection(query_output)
+
+        if inputs_embeds is None:
+            if input_ids is None:
+                image_tokens = [self.config.image_token_index] * self.config.num_query_tokens
+                start_tokens = image_tokens + [self.config.text_config.bos_token_id]
+                input_ids = torch.tensor([start_tokens], dtype=torch.long, device=image_embeds.device)
+                input_ids = input_ids.repeat(batch_size, 1)
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        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
+
+        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        language_model_inputs = language_model_inputs.to(inputs_embeds.device, inputs_embeds.dtype)
+        inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, language_model_inputs)
+
+        inputs = {"inputs_embeds": inputs_embeds, "attention_mask": attention_mask}
+        if not self.language_model.config.is_encoder_decoder:
+            inputs["input_ids"] = input_ids
+
+        outputs = self.language_model.generate(**inputs, **generate_kwargs)
+
+        return outputs