optimum-rbln 0.8.4a6__py3-none-any.whl → 0.8.4a8__py3-none-any.whl

optimum/rbln/diffusers/pipelines/auto_pipeline.py

@@ -15,7 +15,7 @@
 
 import importlib
 from pathlib import Path
-from typing import Type, Union
+from typing import Any, Dict, Type, Union
 
 from diffusers.models.controlnets import ControlNetUnionModel
 from diffusers.pipelines.auto_pipeline import (
@@ -29,6 +29,7 @@ from diffusers.pipelines.auto_pipeline import (
 )
 from huggingface_hub.utils import validate_hf_hub_args
 
+from optimum.rbln.configuration_utils import RBLNModelConfig
 from optimum.rbln.modeling_base import RBLNBaseModel
 from optimum.rbln.utils.model_utils import (
     MODEL_MAPPING,
@@ -168,14 +169,44 @@ class RBLNAutoPipelineBase:
 
     @classmethod
     @validate_hf_hub_args
-    def from_pretrained(cls, model_id, **kwargs):
-        rbln_cls = cls.get_rbln_cls(model_id, **kwargs)
-        return rbln_cls.from_pretrained(model_id, **kwargs)
+    def from_pretrained(
+        cls,
+        model_id: Union[str, Path],
+        *,
+        export: bool = None,
+        rbln_config: Union[Dict[str, Any], RBLNModelConfig] = {},
+        **kwargs: Any,
+    ):
+        """
+        Load an RBLN-accelerated Diffusers pipeline from a pretrained checkpoint or a compiled RBLN artifact.
 
-    @classmethod
-    def from_model(cls, model, **kwargs):
-        rbln_cls = get_rbln_model_cls(f"RBLN{model.__class__.__name__}")
-        return rbln_cls.from_model(model, **kwargs)
+        This method determines the concrete `RBLN*` model class that corresponds to the
+        underlying Diffusers pipeline architecture and dispatches to that class's
+        `from_pretrained()` implementation. If a compiled RBLN folder is detected at `model_id`
+        (or `export=False` is explicitly passed), it loads the compiled artifacts; otherwise it
+        compiles from the original Diffusers checkpoint.
+
+        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 Diffusers 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., `RBLNFluxPipelineConfig`).
+            kwargs: Additional keyword arguments.
+                - Arguments prefixed with `rbln_` are forwarded to the RBLN config.
+                - Remaining arguments are forwarded to the Diffusers loader.
+
+        Returns:
+            RBLNBaseModel: An instantiated RBLN model wrapping the Diffusers pipeline, 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)
 
     @staticmethod
     def register(rbln_cls: Type[RBLNBaseModel], exist_ok=False):
@@ -198,11 +229,15 @@ class RBLNAutoPipelineBase:
 
 
 class RBLNAutoPipelineForText2Image(RBLNAutoPipelineBase, AutoPipelineForText2Image):
+    """Text2Image AutoPipeline for RBLN NPUs."""
+
     _model_mapping = AUTO_TEXT2IMAGE_PIPELINES_MAPPING
     _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_TEXT2IMAGE_PIPELINES_MAPPING.items()}
 
 
 class RBLNAutoPipelineForImage2Image(RBLNAutoPipelineBase, AutoPipelineForImage2Image):
+    """Image2Image AutoPipeline for RBLN NPUs."""
+
     _model_mapping = AUTO_IMAGE2IMAGE_PIPELINES_MAPPING
     _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_IMAGE2IMAGE_PIPELINES_MAPPING.items()}
 
@@ -237,6 +272,8 @@ class RBLNAutoPipelineForImage2Image(RBLNAutoPipelineBase, AutoPipelineForImage2
 
 
 class RBLNAutoPipelineForInpainting(RBLNAutoPipelineBase, AutoPipelineForInpainting):
+    """Inpainting AutoPipeline for RBLN NPUs."""
+
     _model_mapping = AUTO_INPAINT_PIPELINES_MAPPING
     _model_mapping_names = {x[0]: x[1].__name__ for x in AUTO_INPAINT_PIPELINES_MAPPING.items()}
 

optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py

@@ -113,7 +113,7 @@ class RBLNCosmosTextToWorldPipeline(RBLNDiffusionMixin, CosmosTextToWorldPipelin
                 Configuration options for RBLN compilation. Can include settings for specific submodules
                 such as `text_encoder`, `unet`, and `vae`. Configuration can be tailored to the specific
                 pipeline being compiled.
-            **kwargs:
+            kwargs:
                 Additional arguments to pass to the underlying diffusion pipeline constructor or the
                 RBLN compilation process. These may include parameters specific to individual submodules
                 or the particular diffusion pipeline being used.

optimum/rbln/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py

@@ -113,7 +113,7 @@ class RBLNCosmosVideoToWorldPipeline(RBLNDiffusionMixin, CosmosVideoToWorldPipel
                 Configuration options for RBLN compilation. Can include settings for specific submodules
                 such as `text_encoder`, `unet`, and `vae`. Configuration can be tailored to the specific
                 pipeline being compiled.
-            **kwargs:
+            kwargs:
                 Additional arguments to pass to the underlying diffusion pipeline constructor or the
                 RBLN compilation process. These may include parameters specific to individual submodules
                 or the particular diffusion pipeline being used.

optimum/rbln/modeling.py

@@ -85,11 +85,13 @@ class RBLNModel(RBLNBaseModel):
         This method performs the actual model conversion and compilation process.
 
         Args:
-            model: The PyTorch model to be compiled. The object must be an instance of the HuggingFace transformers PreTrainedModel class.
-            rbln_config: Configuration for RBLN model compilation and runtime. This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
+            model (PreTrainedModel): The PyTorch model to be compiled.
+                The object must be an instance of the HuggingFace transformers PreTrainedModel class.
+            config (Optional[PretrainedConfig]): The configuration object associated with the model.
+            rbln_config (Optional[Union[RBLNModelConfig, Dict]]): Configuration for RBLN model compilation and runtime.
+                This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
                 For detailed configuration options, see the specific model's configuration class documentation.
-
-            kwargs: Additional keyword arguments. Arguments with the prefix 'rbln_' are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
+            kwargs: Additional keyword arguments. Arguments with the prefix `rbln_` are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
 
         The method performs the following steps:
 
@@ -99,7 +101,7 @@ class RBLNModel(RBLNBaseModel):
         4. Saves the compiled model and configurations
 
         Returns:
-            A RBLN model instance ready for inference on RBLN NPU devices.
+            (RBLNModel): A RBLN model instance ready for inference on RBLN NPU devices.
         """
         preprocessors = kwargs.pop("preprocessors", [])
         rbln_config, kwargs = cls.prepare_rbln_config(rbln_config=rbln_config, **kwargs)
@@ -241,29 +243,31 @@ class RBLNModel(RBLNBaseModel):
 
     def forward(self, *args: Any, return_dict: Optional[bool] = None, **kwargs: Any) -> Any:
         """
-        Defines the forward pass of the RBLN model, providing a drop-in replacement for HuggingFace PreTrainedModel.
+        Defines the forward pass of `RBLNModel`. The interface mirrors HuggingFace conventions so it can act as a drop-in
+        replacement in many cases.
 
-        This method executes the compiled RBLN model on RBLN NPU devices while maintaining full compatibility
-        with HuggingFace transformers and diffusers APIs. The RBLNModel can be used as a direct substitute
-        for any HuggingFace nn.Module/PreTrainedModel, enabling seamless integration into existing workflows.
+        This method executes the compiled RBLN model on RBLN NPU devices while remaining fully compatible with Hugging Face
+        Transformers and Diffusers APIs. In practice, `RBLNModel` can replace models built on `torch.nn.Module` — including
+        `transformers.PreTrainedModel` implementations and Diffusers components based on `diffusers.ModelMixin` — enabling
+        seamless integration into existing workflows.
 
         Args:
-            *args: Variable length argument list containing model inputs. The format matches the original
+            args: Variable length argument list containing model inputs. The format matches the original
                 HuggingFace model's forward method signature (e.g., input_ids, attention_mask for
                 transformers models, or sample, timestep for diffusers models).
             return_dict:
                 Whether to return outputs as a dictionary-like object or as a tuple. When `None`:
                 - For transformers models: Uses `self.config.use_return_dict` (typically `True`)
                 - For diffusers models: Defaults to `True`
-            **kwargs: Arbitrary keyword arguments containing additional model inputs and parameters,
+            kwargs: Arbitrary keyword arguments containing additional model inputs and parameters,
                 matching the original HuggingFace model's interface.
 
         Returns:
             Model outputs in the same format as the original HuggingFace model.
 
-            - If `return_dict=True`: Returns a dictionary-like object (e.g., BaseModelOutput,
+            If `return_dict=True`, Returns a dictionary-like object (e.g., BaseModelOutput,
                 CausalLMOutput) with named fields such as `logits`, `hidden_states`, etc.
-            - If `return_dict=False`: Returns a tuple containing the raw model outputs.
+            If `return_dict=False`, Returns a tuple containing the raw model outputs.
 
         Note:
             - This method maintains the exact same interface as the original HuggingFace model's forward method

optimum/rbln/modeling_base.py

@@ -373,7 +373,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
     def from_pretrained(
         cls: Type["RBLNBaseModel"],
         model_id: Union[str, Path],
-        export: bool = None,
+        export: Optional[bool] = None,
         rbln_config: Optional[Union[Dict, RBLNModelConfig]] = None,
         **kwargs: Any,
     ) -> "RBLNBaseModel":
@@ -382,15 +382,17 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         User can use this function to load a pre-trained model from the HuggingFace library and convert it to a RBLN model to be run on RBLN NPUs.
 
         Args:
-            model_id: The model id of the pre-trained model to be loaded. It can be downloaded from the HuggingFace model hub or a local path, or a model id of a compiled model using the RBLN Compiler.
-            export: A boolean flag to indicate whether the model should be compiled. If None, it will be determined based on the existence of the compiled model files in the model_id.
-            rbln_config: Configuration for RBLN model compilation and runtime. This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
+            model_id (Union[str, Path]): The model id of the pre-trained model to be loaded.
+                It can be downloaded from the HuggingFace model hub or a local path, or a model id of a compiled model using the RBLN Compiler.
+            export (Optional[bool]): A boolean flag to indicate whether the model should be compiled.
+                If None, it will be determined based on the existence of the compiled model files in the model_id.
+            rbln_config (Optional[Union[Dict, RBLNModelConfig]]): Configuration for RBLN model compilation and runtime.
+                This can be provided as a dictionary or an instance of the model's configuration class (e.g., `RBLNLlamaForCausalLMConfig` for Llama models).
                 For detailed configuration options, see the specific model's configuration class documentation.
-
-            kwargs: Additional keyword arguments. Arguments with the prefix 'rbln_' are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
+            kwargs: Additional keyword arguments. Arguments with the prefix `rbln_` are passed to rbln_config, while the remaining arguments are passed to the HuggingFace library.
 
         Returns:
-            A RBLN model instance ready for inference on RBLN NPU devices.
+            (RBLNModel): A RBLN model instance ready for inference on RBLN NPU devices.
         """
 
         if isinstance(model_id, Path):
@@ -536,9 +538,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         [`~optimum.rbln.modeling_base.RBLNBaseModel.from_pretrained`] class method.
 
         Args:
-            save_directory (`Union[str, Path]`):
+            save_directory (Union[str, Path]):
                 Directory where to save the model file.
-            push_to_hub (`bool`, *optional*, defaults to `False`):
+            push_to_hub (bool):
                 Whether or not to push your model to the HuggingFace model hub after saving it.
 
         """

optimum/rbln/transformers/configuration_generic.py

@@ -34,7 +34,7 @@ class RBLNTransformerEncoderConfig(RBLNModelConfig):
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
             model_input_names (Optional[List[str]]): Names of the input tensors for the model.
                 Defaults to class-specific rbln_model_input_names if not provided.
-            **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.
@@ -61,7 +61,7 @@ class RBLNImageModelConfig(RBLNModelConfig):
             image_size (Optional[Union[int, Tuple[int, int]]]): The size of input images.
                 Can be an integer for square images or a tuple (height, width).
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
-            **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.
@@ -133,7 +133,7 @@ class RBLNModelForAudioClassificationConfig(RBLNModelConfig):
             batch_size (Optional[int]): The batch size for inference. Defaults to 1.
             max_length (Optional[int]): Maximum length of the audio input in time dimension.
             num_mel_bins (Optional[int]): Number of Mel frequency bins for audio processing.
-            **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.

optimum/rbln/transformers/modeling_generic.py

@@ -253,6 +253,7 @@ class RBLNModelForAudioClassification(RBLNModel):
 
     A class to convert and run pre-trained transformers based AudioClassification models on RBLN devices.
     It implements the methods to convert a pre-trained transformers AudioClassification model into a RBLN transformer model by:
+
     - transferring the checkpoint weights of the original into an optimized RBLN graph,
     - compiling the resulting graph using the RBLN compiler.
 

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

@@ -15,13 +15,13 @@ import importlib
 import inspect
 import warnings
 from pathlib import Path
-from typing import Any, Type, Union
+from typing import Any, Dict, Optional, Type, Union
 
 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,
@@ -178,14 +178,74 @@ class _BaseAutoModelClass:
         return rbln_config.rbln_model_cls_name
 
     @classmethod
-    def from_pretrained(cls, model_id: Union[str, Path], *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: PreTrainedModel, *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: bool = False):

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

@@ -57,75 +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/blip_2/configuration_blip_2.py

@@ -42,11 +42,9 @@ class RBLNBlip2QFormerModelConfig(RBLNModelConfig):
     ):
         """
         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.num_query_tokens = num_query_tokens
@@ -68,8 +66,9 @@ class RBLNBlip2ForConditionalGenerationConfig(RBLNModelConfig):
         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.

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

@@ -112,7 +112,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,

optimum/rbln/transformers/models/clip/configuration_clip.py

@@ -22,10 +22,10 @@ class RBLNCLIPTextModelConfig(RBLNModelConfig):
         """
         Args:
             batch_size (Optional[int]): The batch size for text processing. Defaults to 1.
-            **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.
+            ValueError: If `batch_size` is not a positive integer.
         """
         super().__init__(**kwargs)
         self.batch_size = batch_size or 1
@@ -57,10 +57,13 @@ class RBLNCLIPVisionModelConfig(RBLNModelConfig):
             batch_size (Optional[int]): The batch size for image processing. Defaults to 1.
             image_size (Optional[int]): The size of input images. Can be an integer for square images,
                 a tuple/list (height, width), or a dictionary with 'height' and 'width' keys.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            interpolate_pos_encoding (Optional[bool]): Whether or not to interpolate pre-trained position encodings. Defaults to `False`.
+            output_hidden_states (Optional[bool]): Whether or not to return the hidden states of all layers.
+            output_attentions (Optional[bool]): Whether or not to return the attentions tensors of all attention layers
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
-            ValueError: If batch_size is not a positive integer.
+            ValueError: If `batch_size` is not a positive integer.
         """
         super().__init__(**kwargs)
         self.batch_size = batch_size or 1

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -83,7 +83,15 @@ class RBLNCLIPTextModel(RBLNModel):
         rbln_config.set_compile_cfgs([RBLNCompileConfig(input_info=input_info)])
         return rbln_config
 
-    def forward(self, input_ids: torch.LongTensor, return_dict: bool = None, **kwargs) -> torch.FloatTensor:
+    def forward(self, input_ids: torch.LongTensor, return_dict: Optional[bool] = None, **kwargs) -> torch.FloatTensor:
+        """
+        Forward pass for the RBLN-optimized CLIP text encoder model.
+
+        Args:
+            input_ids (torch.LongTensor): The input ids to the model.
+            return_dict (Optional[bool]): Whether to return a dictionary of outputs.
+        """
+
         # To ignore using attention_mask, we override forward method.
         output = super().forward(input_ids, return_dict=return_dict)
         return output
@@ -202,13 +210,24 @@ class RBLNCLIPVisionModel(RBLNModel):
 
     def forward(
         self,
-        pixel_values: Optional[torch.FloatTensor] = None,
+        pixel_values: torch.FloatTensor,
         return_dict: bool = True,
-        output_attentions: bool = None,
-        output_hidden_states: bool = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
         interpolate_pos_encoding: bool = False,
         **kwargs,
     ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        """
+        Forward pass for the RBLN-optimized CLIP vision encoder model.
+
+        Args:
+            pixel_values (torch.Tensor): The pixel values to the model.
+            return_dict (bool): Whether to return a dictionary of outputs.
+            output_attentions (Optional[bool]): Whether to return attentions.
+            output_hidden_states (Optional[bool]): Whether to return hidden states.
+            interpolate_pos_encoding (bool): Whether to interpolate position encoding.
+        """
+
         if len(kwargs) > 0 and any(value is not None for value in kwargs.values()):
             logger.warning(
                 f"Currently, optimum-rbln does not support kwargs {kwargs.keys()} for {self.__class__.__name__}."

optimum/rbln/transformers/models/colpali/configuration_colpali.py

@@ -54,11 +54,11 @@ class RBLNColPaliForRetrievalConfig(RBLNModelConfig):
     ):
         """
         Args:
-            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
             max_seq_lens (Union[int, List[int]]): The maximum sequence lengths for the language model.
                 This can be multiple values, and the model will be compiled for each max_seq_len, allowing selection of the most appropriate max_seq_len at inference time.
             output_hidden_states (Optional[bool]): Whether to output the hidden states of the language model.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            vision_tower (Optional[RBLNModelConfig]): Configuration for the vision encoder component.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
         Raises:
             ValueError: If batch_size is not a positive integer.
         """