optimum-rbln 0.9.5a4__py3-none-any.whl → 0.10.0.post1__py3-none-any.whl

optimum/rbln/__init__.py

@@ -86,6 +86,8 @@ _import_structure = {
         "RBLNDPTForDepthEstimationConfig",
         "RBLNDepthAnythingForDepthEstimationConfig",
         "RBLNDepthAnythingForDepthEstimation",
+        "RBLNDetrForObjectDetection",
+        "RBLNDetrForObjectDetectionConfig",
         "RBLNExaoneForCausalLM",
         "RBLNExaoneForCausalLMConfig",
         "RBLNGemmaModel",
@@ -120,6 +122,8 @@ _import_structure = {
         "RBLNLlamaForCausalLMConfig",
         "RBLNLlamaModel",
         "RBLNLlamaModelConfig",
+        "RBLNMixtralForCausalLM",
+        "RBLNMixtralForCausalLMConfig",
         "RBLNOPTForCausalLM",
         "RBLNOPTForCausalLMConfig",
         "RBLNLlavaForConditionalGeneration",
@@ -406,6 +410,8 @@ if TYPE_CHECKING:
         RBLNDecoderOnlyModelForCausalLMConfig,
         RBLNDepthAnythingForDepthEstimation,
         RBLNDepthAnythingForDepthEstimationConfig,
+        RBLNDetrForObjectDetection,
+        RBLNDetrForObjectDetectionConfig,
         RBLNDistilBertForQuestionAnswering,
         RBLNDistilBertForQuestionAnsweringConfig,
         RBLNDPTForDepthEstimation,
@@ -456,6 +462,8 @@ if TYPE_CHECKING:
         RBLNMistralForCausalLMConfig,
         RBLNMistralModel,
         RBLNMistralModelConfig,
+        RBLNMixtralForCausalLM,
+        RBLNMixtralForCausalLMConfig,
         RBLNOPTForCausalLM,
         RBLNOPTForCausalLMConfig,
         RBLNOPTModel,

optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.5a4'
-__version_tuple__ = version_tuple = (0, 9, 5, 'a4')
+__version__ = version = '0.10.0.post1'
+__version_tuple__ = version_tuple = (0, 10, 0, 'post1')
 
 __commit_id__ = commit_id = None

optimum/rbln/configuration_utils.py

@@ -24,7 +24,7 @@ import torch
 from packaging.version import Version
 
 from .__version__ import __version__
-from .utils.deprecation import deprecate_kwarg, warn_deprecated_npu
+from .utils.deprecation import deprecate_kwarg, deprecate_method, warn_deprecated_npu
 from .utils.logging import get_logger
 from .utils.runtime_utils import ContextRblnConfig
 
@@ -36,6 +36,30 @@ DEFAULT_COMPILED_MODEL_NAME = "compiled_model"
 TypeInputInfo = List[Tuple[str, Tuple[int], str]]
 
 
+def nested_update(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Recursively merge override dict into base dict.
+    For nested dicts, values are merged recursively instead of being replaced.
+    For non-dict values, override takes precedence.
+    Args:
+        base: The base dictionary to merge into (modified in-place).
+        override: The dictionary with values to merge.
+    Returns:
+        The merged base dictionary.
+    Example:
+        >>> base = {"a": 1, "nested": {"x": 10, "y": 20}}
+        >>> override = {"b": 2, "nested": {"y": 30, "z": 40}}
+        >>> nested_update(base, override)
+        {"a": 1, "b": 2, "nested": {"x": 10, "y": 30, "z": 40}}
+    """
+    for key, value in override.items():
+        if key in base and isinstance(base[key], dict) and isinstance(value, dict):
+            nested_update(base[key], value)
+        else:
+            base[key] = value
+    return base
+
+
 @runtime_checkable
 class RBLNSerializableConfigProtocol(Protocol):
     def _prepare_for_serialization(self) -> Dict[str, Any]: ...
@@ -216,8 +240,7 @@ class RBLNAutoConfig:
                 For example, the parsed contents of `rbln_config.json`.
 
         Returns:
-            RBLNModelConfig: A configuration instance. The specific subclass is
-            selected by `config_dict["cls_name"]`.
+            RBLNModelConfig: A configuration instance. The specific subclass is selected by `config_dict["cls_name"]`.
 
         Raises:
             ValueError: If `cls_name` is missing.
@@ -256,12 +279,13 @@ class RBLNAutoConfig:
 
         CONFIG_MAPPING[config.__name__] = config
 
-    @staticmethod
-    def load(
+    @classmethod
+    def from_pretrained(
+        cls,
         path: str,
-        passed_rbln_config: Optional["RBLNModelConfig"] = None,
-        kwargs: Optional[Dict[str, Any]] = None,
+        rbln_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
         return_unused_kwargs: bool = False,
+        **kwargs: Optional[Dict[str, Any]],
     ) -> Union["RBLNModelConfig", Tuple["RBLNModelConfig", Dict[str, Any]]]:
         """
         Load RBLNModelConfig from a path.
@@ -269,53 +293,58 @@ class RBLNAutoConfig:
 
         Args:
             path (str): Path to the RBLNModelConfig.
-            passed_rbln_config (Optional["RBLNModelConfig"]): RBLNModelConfig to pass its runtime options.
+            rbln_config (Optional[Dict[str, Any]]): Additional configuration to override.
+            return_unused_kwargs (bool): Whether to return unused kwargs.
+            kwargs: Additional keyword arguments to override configuration values.
 
         Returns:
             RBLNModelConfig: The loaded RBLNModelConfig.
-        """
-        if kwargs is None:
-            kwargs = {}
-        cls, config_file = load_config(path)
-
-        rbln_keys = [key for key in kwargs.keys() if key.startswith("rbln_")]
-        rbln_runtime_kwargs = {key[5:]: kwargs.pop(key) for key in rbln_keys if key[5:] in RUNTIME_KEYWORDS}
-        rbln_submodule_kwargs = {key[5:]: kwargs.pop(key) for key in rbln_keys if key[5:] in cls.submodules}
-
-        rbln_kwargs = {
-            key[5:]: kwargs.pop(key)
-            for key in rbln_keys
-            if key[5:] not in RUNTIME_KEYWORDS and key[5:] not in cls.submodules
-        }
 
-        # Process submodule's rbln_config
-        for submodule in cls.submodules:
-            if submodule not in config_file:
-                raise ValueError(f"Submodule {submodule} not found in rbln_config.json.")
-            submodule_config = config_file[submodule]
-            submodule_config.update(rbln_submodule_kwargs.pop(submodule, {}))
-            config_file[submodule] = RBLNAutoConfig.load_from_dict(submodule_config)
+        Examples:
+            ```python
+                config = RBLNAutoConfig.from_pretrained("/path/to/model")
+            ```
+        """
+        target_cls, _ = load_config(path)
+        return target_cls.from_pretrained(
+            path, rbln_config=rbln_config, return_unused_kwargs=return_unused_kwargs, **kwargs
+        )
 
-        if passed_rbln_config is not None:
-            config_file.update(passed_rbln_config._runtime_options)
-            # TODO(jongho): Reject if the passed_rbln_config has different attributes from the config_file
+    @classmethod
+    @deprecate_method(version="0.11.0", new_method="from_pretrained")
+    def load(
+        cls,
+        path: str,
+        rbln_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
+        return_unused_kwargs: bool = False,
+        **kwargs: Optional[Dict[str, Any]],
+    ) -> Union["RBLNModelConfig", Tuple["RBLNModelConfig", Dict[str, Any]]]:
+        """
+        Load RBLNModelConfig from a path.
+        Class name is automatically inferred from the `rbln_config.json` file.
 
-        config_file.update(rbln_runtime_kwargs)
+        Deprecated:
+            This method is deprecated and will be removed in version 0.11.0.
+            Use `from_pretrained` instead.
 
-        rbln_config = cls(**config_file)
+        Args:
+            path (str): Path to the RBLNModelConfig file or directory.
+            rbln_config (Optional[Dict[str, Any]]): Additional configuration to override.
+            return_unused_kwargs (bool): Whether to return unused kwargs.
+            kwargs: Additional keyword arguments to override configuration values.
 
-        if len(rbln_kwargs) > 0:
-            for key, value in rbln_kwargs.items():
-                if getattr(rbln_config, key) != value:
-                    raise ValueError(
-                        f"Cannot set the following arguments: {list(rbln_kwargs.keys())} "
-                        f"Since the value is already set to {getattr(rbln_config, key)}"
-                    )
+        Returns:
+            RBLNModelConfig: The loaded RBLNModelConfig.
 
-        if return_unused_kwargs:
-            return cls(**config_file), kwargs
-        else:
-            return cls(**config_file)
+        Examples:
+            ```python
+            # Deprecated usage:
+            config = RBLNAutoConfig.load("/path/to/model")
+            # Recommended usage:
+            config = RBLNAutoConfig.from_pretrained("/path/to/model")
+            ```
+        """
+        return cls.from_pretrained(path, rbln_config=rbln_config, return_unused_kwargs=return_unused_kwargs, **kwargs)
 
 
 class RBLNModelConfig(RBLNSerializableConfigProtocol):
@@ -866,15 +895,23 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             json.dump(serializable_data, jsonf, indent=2)
 
     @classmethod
-    def load(cls, path: str, **kwargs: Any) -> "RBLNModelConfig":
+    def from_pretrained(
+        cls,
+        path: str,
+        rbln_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
+        return_unused_kwargs: bool = False,
+        **kwargs: Optional[Dict[str, Any]],
+    ) -> Union["RBLNModelConfig", Tuple["RBLNModelConfig", Dict[str, Any]]]:
         """
         Load a RBLNModelConfig from a path.
 
         Args:
             path (str): Path to the RBLNModelConfig file or directory containing the config file.
+            rbln_config (Optional[Dict[str, Any]]): Additional configuration to override.
+            return_unused_kwargs (bool): Whether to return unused kwargs.
             kwargs: Additional keyword arguments to override configuration values.
-                      Keys starting with 'rbln_' will have the prefix removed and be used
-                      to update the configuration.
+                    Keys starting with 'rbln_' will have the prefix removed and be used
+                    to update the configuration.
 
         Returns:
             RBLNModelConfig: The loaded configuration instance.
@@ -883,17 +920,109 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             This method loads the configuration from the specified path and applies any
             provided overrides. If the loaded configuration class doesn't match the expected
             class, a warning will be logged.
+
+        Examples:
+            ```python
+            config = RBLNResNetForImageClassificationConfig.from_pretrained("/path/to/model")
+            ```
         """
         cls_reserved, config_file = load_config(path)
-
         if cls_reserved != cls:
             logger.warning(f"Expected {cls.__name__}, but got {cls_reserved.__name__}.")
 
+        if isinstance(rbln_config, dict):
+            for key, value in rbln_config.items():
+                if key not in kwargs:
+                    kwargs[f"rbln_{key}"] = value
+
         rbln_keys = [key for key in kwargs.keys() if key.startswith("rbln_")]
-        rbln_kwargs = {key[5:]: kwargs.pop(key) for key in rbln_keys}
-        config_file.update(rbln_kwargs)
+        rbln_runtime_kwargs = {key[5:]: kwargs.pop(key) for key in rbln_keys if key[5:] in RUNTIME_KEYWORDS}
+        rbln_submodule_kwargs = {key[5:]: kwargs.pop(key) for key in rbln_keys if key[5:] in cls.submodules}
+
+        rbln_kwargs = {
+            key[5:]: kwargs.pop(key)
+            for key in rbln_keys
+            if key[5:] not in RUNTIME_KEYWORDS and key[5:] not in cls.submodules
+        }
+
+        # Process submodule's rbln_config
+        for submodule in cls.submodules:
+            if submodule not in config_file:
+                raise ValueError(f"Submodule {submodule} not found in rbln_config.json.")
+            submodule_config = config_file[submodule]
+            submodule_config.update(rbln_runtime_kwargs)
 
-        return cls(**config_file)
+            update_dict = rbln_submodule_kwargs.pop(submodule, {})
+            if update_dict:
+                nested_update(submodule_config, update_dict)
+            config_file[submodule] = RBLNAutoConfig.load_from_dict(submodule_config)
+
+        if isinstance(rbln_config, RBLNModelConfig):
+            config_file.update(rbln_config._runtime_options)
+
+            # update submodule runtime
+            for submodule in rbln_config.submodules:
+                if str(config_file[submodule]) != str(getattr(rbln_config, submodule)):
+                    raise ValueError(
+                        f"Passed rbln_config has different attributes for submodule {submodule} than the config_file"
+                    )
+                config_file[submodule] = getattr(rbln_config, submodule)
+
+        config_file.update(rbln_runtime_kwargs)
+        rbln_config = cls(**config_file)
+        if len(rbln_kwargs) > 0:
+            for key, value in rbln_kwargs.items():
+                if getattr(rbln_config, key) != value:
+                    raise ValueError(
+                        f"Cannot set the following arguments: {list(rbln_kwargs.keys())} "
+                        f"Since the value is already set to {getattr(rbln_config, key)}"
+                    )
+        if return_unused_kwargs:
+            return rbln_config, kwargs
+        else:
+            return rbln_config
+
+    @classmethod
+    @deprecate_method(version="0.11.0", new_method="from_pretrained")
+    def load(
+        cls,
+        path: str,
+        rbln_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
+        return_unused_kwargs: bool = False,
+        **kwargs: Optional[Dict[str, Any]],
+    ) -> Union["RBLNModelConfig", Tuple["RBLNModelConfig", Dict[str, Any]]]:
+        """
+        Load a RBLNModelConfig from a path.
+
+        Deprecated:
+            This method is deprecated and will be removed in version 0.11.0.
+            Use `from_pretrained` instead.
+
+        Args:
+            path (str): Path to the RBLNModelConfig file or directory containing the config file.
+            rbln_config (Optional[Dict[str, Any]]): Additional configuration to override.
+            return_unused_kwargs (bool): Whether to return unused kwargs.
+            kwargs: Additional keyword arguments to override configuration values.
+                    Keys starting with 'rbln_' will have the prefix removed and be used
+                    to update the configuration.
+
+        Returns:
+            RBLNModelConfig: The loaded configuration instance.
+
+        Note:
+            This method loads the configuration from the specified path and applies any
+            provided overrides. If the loaded configuration class doesn't match the expected
+            class, a warning will be logged.
+
+        Examples:
+            ```python
+            # Deprecated usage:
+            config = RBLNResNetForImageClassificationConfig.load("/path/to/model")
+            # Recommended usage:
+            config = RBLNResNetForImageClassificationConfig.from_pretrained("/path/to/model")
+            ```
+        """
+        return cls.from_pretrained(path, rbln_config=rbln_config, return_unused_kwargs=return_unused_kwargs, **kwargs)
 
     @classmethod
     def initialize_from_kwargs(
@@ -993,3 +1122,18 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     @timeout.setter
     def timeout(self, timeout: int):
         self._runtime_options["timeout"] = timeout
+
+
+def convert_rbln_config_dict(
+    rbln_config: Optional[Union[Dict[str, Any], RBLNModelConfig]] = None, **kwargs
+) -> Tuple[Optional[Union[Dict[str, Any], RBLNModelConfig]], Dict[str, Any]]:
+    # Validate and merge rbln_ prefixed kwargs into rbln_config
+    kwargs_keys = list(kwargs.keys())
+    rbln_kwargs = {key[5:]: kwargs.pop(key) for key in kwargs_keys if key.startswith("rbln_")}
+
+    rbln_config = {} if rbln_config is None else rbln_config
+
+    if isinstance(rbln_config, dict) and len(rbln_kwargs) > 0:
+        rbln_config.update(rbln_kwargs)
+
+    return rbln_config, kwargs

optimum/rbln/diffusers/models/controlnet.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Dict, Optional, Union
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 
 import torch
 from diffusers import ControlNetModel
@@ -218,7 +218,7 @@ class RBLNControlNetModel(RBLNModel):
         added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
         return_dict: bool = True,
         **kwargs,
-    ):
+    ) -> Union[ControlNetOutput, Tuple]:
         """
         Forward pass for the RBLN-optimized ControlNetModel.
 

optimum/rbln/diffusers/models/transformers/prior_transformer.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 import torch
 from diffusers.models.transformers.prior_transformer import PriorTransformer, PriorTransformerOutput
@@ -134,7 +134,7 @@ class RBLNPriorTransformer(RBLNModel):
         encoder_hidden_states: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         return_dict: bool = True,
-    ):
+    ) -> Union[PriorTransformerOutput, Tuple]:
         """
         Forward pass for the RBLN-optimized PriorTransformer.
 

optimum/rbln/diffusers/models/transformers/transformer_cosmos.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 import rebel
 import torch
@@ -302,7 +302,7 @@ class RBLNCosmosTransformer3DModel(RBLNModel):
         condition_mask: Optional[torch.Tensor] = None,
         padding_mask: Optional[torch.Tensor] = None,
         return_dict: bool = True,
-    ):
+    ) -> Union[Transformer2DModelOutput, Tuple]:
         """
         Forward pass for the RBLN-optimized CosmosTransformer3DModel.
 

optimum/rbln/diffusers/models/transformers/transformer_sd3.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import torch
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
@@ -160,7 +160,7 @@ class RBLNSD3Transformer2DModel(RBLNModel):
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = True,
         **kwargs,
-    ):
+    ) -> Union[Transformer2DModelOutput, Tuple]:
         """
         Forward pass for the RBLN-optimized SD3Transformer2DModel.
 

optimum/rbln/diffusers/pipelines/auto_pipeline.py

@@ -176,7 +176,7 @@ class RBLNAutoPipelineBase:
         export: bool = None,
         rbln_config: Optional[Union[Dict[str, Any], RBLNModelConfig]] = None,
         **kwargs: Any,
-    ):
+    ) -> RBLNBaseModel:
         """
         Load an RBLN-accelerated Diffusers pipeline from a pretrained checkpoint or a compiled RBLN artifact.
 
@@ -201,8 +201,7 @@ class RBLNAutoPipelineBase:
                 - Remaining arguments are forwarded to the Diffusers loader.
 
         Returns:
-            RBLNBaseModel: An instantiated RBLN model wrapping the Diffusers pipeline, ready for
-            inference on RBLN NPUs.
+            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)

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py

@@ -26,7 +26,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -260,7 +260,7 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
         callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
         **kwargs,
-    ):
+    ) -> Union[StableDiffusionPipelineOutput, Tuple]:
         r"""
         The call function to the pipeline for generation.
 
@@ -321,14 +321,7 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
             output_type (`str`, *optional*, defaults to `"pil"`):
                 The output format of the generated image. Choose between `PIL.Image` or `np.array`.
             return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
-                plain tuple.
-            callback (`Callable`, *optional*):
-                A function that calls every `callback_steps` steps during inference. The function is called with the
-                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
-            callback_steps (`int`, *optional*, defaults to 1):
-                The frequency at which the `callback` function is called. If not specified, the callback is called at
-                every step.
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple.
             cross_attention_kwargs (`dict`, *optional*):
                 A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
                 [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
@@ -356,8 +349,6 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
                 will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
                 `._callback_tensor_inputs` attribute of your pipeine class.
 
-        Examples:
-
         Returns:
             [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
                 If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py

@@ -26,7 +26,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -253,7 +253,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
         callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
         **kwargs,
-    ):
+    ) -> Union[StableDiffusionPipelineOutput, Tuple]:
         r"""
         The call function to the pipeline for generation.
 
@@ -347,8 +347,6 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
                 will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
                 `._callback_tensor_inputs` attribute of your pipeine class.
 
-        Examples:
-
         Returns:
             [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
                 If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py

@@ -294,7 +294,7 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
         callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
         **kwargs,
-    ):
+    ) -> Union[StableDiffusionXLPipelineOutput, Tuple]:
         r"""
         The call function to the pipeline for generation.
 
@@ -431,8 +431,6 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
                 will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
                 `._callback_tensor_inputs` attribute of your pipeine class.
 
-        Examples:
-
         Returns:
             [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
                 If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py

@@ -309,7 +309,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
         callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
         **kwargs,
-    ):
+    ) -> Union[StableDiffusionXLPipelineOutput, Tuple]:
         r"""
         Function invoked when calling the pipeline for generation.
 
@@ -465,8 +465,6 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
                 will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
                 `._callback_tensor_inputs` attribute of your pipeine class.
 
-        Examples:
-
         Returns:
             [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
             [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple`

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

@@ -203,7 +203,7 @@ class RBLNRetinaFaceFilter(RetinaFaceFilter):
                 f"If you only need to compile the model without loading it to NPU, you can use:\n"
                 f"  from_pretrained(..., rbln_create_runtimes=False) or\n"
                 f"  from_pretrained(..., rbln_config={{..., 'create_runtimes': False}})\n\n"
-                f"To check your NPU status, run the 'rbln-stat' command in your terminal.\n"
+                f"To check your NPU status, run the 'rbln-smi' command in your terminal.\n"
                 f"Make sure your NPU is properly installed and operational."
             )
             raise rebel.core.exception.RBLNRuntimeError(error_msg) from e
@@ -278,7 +278,7 @@ class RBLNVideoSafetyModel(VideoSafetyModel):
                 f"If you only need to compile the model without loading it to NPU, you can use:\n"
                 f"  from_pretrained(..., rbln_create_runtimes=False) or\n"
                 f"  from_pretrained(..., rbln_config={{..., 'create_runtimes': False}})\n\n"
-                f"To check your NPU status, run the 'rbln-stat' command in your terminal.\n"
+                f"To check your NPU status, run the 'rbln-smi' command in your terminal.\n"
                 f"Make sure your NPU is properly installed and operational."
             )
             raise rebel.core.exception.RBLNRuntimeError(error_msg) from e

optimum/rbln/modeling_base.py

@@ -24,7 +24,7 @@ import torch
 from transformers import AutoConfig, AutoModel, GenerationConfig, PretrainedConfig
 from transformers.utils.hub import PushToHubMixin
 
-from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig, get_rbln_config_class
+from .configuration_utils import RBLNCompileConfig, RBLNModelConfig, get_rbln_config_class
 from .utils.hub import pull_compiled_model_from_hub, validate_files
 from .utils.logging import get_logger
 from .utils.runtime_utils import UnavailableRuntime, tp_and_devices_are_ok
@@ -206,8 +206,9 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                     f"does not match the expected model class name ({cls.__name__})."
                 )
 
-            rbln_config, kwargs = RBLNAutoConfig.load(
-                model_path_subfolder, passed_rbln_config=rbln_config, kwargs=kwargs, return_unused_kwargs=True
+            config_cls = cls.get_rbln_config_class()
+            rbln_config, kwargs = config_cls.from_pretrained(
+                model_path_subfolder, rbln_config=rbln_config, return_unused_kwargs=True, **kwargs
             )
 
             if rbln_config.rbln_model_cls_name != cls.__name__:
@@ -306,7 +307,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                 f"If you only need to compile the model without loading it to NPU, you can use:\n"
                 f"  from_pretrained(..., rbln_create_runtimes=False) or\n"
                 f"  from_pretrained(..., rbln_config={{..., 'create_runtimes': False}})\n\n"
-                f"To check your NPU status, run the 'rbln-stat' command in your terminal.\n"
+                f"To check your NPU status, run the 'rbln-smi' command in your terminal.\n"
                 f"Make sure your NPU is properly installed and operational."
             )
             raise rebel.core.exception.RBLNRuntimeError(error_msg) from e

optimum/rbln/transformers/__init__.py

@@ -68,6 +68,8 @@ _import_structure = {
         "RBLNDecoderOnlyModelForCausalLMConfig",
         "RBLNDecoderOnlyModelConfig",
         "RBLNDecoderOnlyModel",
+        "RBLNDetrForObjectDetection",
+        "RBLNDetrForObjectDetectionConfig",
         "RBLNDistilBertForQuestionAnswering",
         "RBLNDistilBertForQuestionAnsweringConfig",
         "RBLNDPTForDepthEstimation",
@@ -130,6 +132,8 @@ _import_structure = {
         "RBLNMistralForCausalLMConfig",
         "RBLNMistralModel",
         "RBLNMistralModelConfig",
+        "RBLNMixtralForCausalLM",
+        "RBLNMixtralForCausalLMConfig",
         "RBLNOPTForCausalLM",
         "RBLNOPTForCausalLMConfig",
         "RBLNOPTModel",
@@ -246,6 +250,8 @@ if TYPE_CHECKING:
         RBLNDecoderOnlyModelForCausalLMConfig,
         RBLNDepthAnythingForDepthEstimation,
         RBLNDepthAnythingForDepthEstimationConfig,
+        RBLNDetrForObjectDetection,
+        RBLNDetrForObjectDetectionConfig,
         RBLNDistilBertForQuestionAnswering,
         RBLNDistilBertForQuestionAnsweringConfig,
         RBLNDPTForDepthEstimation,
@@ -296,6 +302,8 @@ if TYPE_CHECKING:
         RBLNMistralForCausalLMConfig,
         RBLNMistralModel,
         RBLNMistralModelConfig,
+        RBLNMixtralForCausalLM,
+        RBLNMixtralForCausalLMConfig,
         RBLNOPTForCausalLM,
         RBLNOPTForCausalLMConfig,
         RBLNOPTModel,