optimum-rbln 0.9.4a2__py3-none-any.whl → 0.10.0.post1__py3-none-any.whl

optimum/rbln/__init__.py

@@ -86,11 +86,17 @@ _import_structure = {
         "RBLNDPTForDepthEstimationConfig",
         "RBLNDepthAnythingForDepthEstimationConfig",
         "RBLNDepthAnythingForDepthEstimation",
+        "RBLNDetrForObjectDetection",
+        "RBLNDetrForObjectDetectionConfig",
         "RBLNExaoneForCausalLM",
         "RBLNExaoneForCausalLMConfig",
         "RBLNGemmaModel",
         "RBLNGemmaModelConfig",
         "RBLNGemmaForCausalLM",
+        "RBLNGemma2ForCausalLM",
+        "RBLNGemma2ForCausalLMConfig",
+        "RBLNGemma2Model",
+        "RBLNGemma2ModelConfig",
         "RBLNGemmaForCausalLMConfig",
         "RBLNGemma3ForCausalLM",
         "RBLNGemma3ForCausalLMConfig",
@@ -100,6 +106,8 @@ _import_structure = {
         "RBLNGPT2ModelConfig",
         "RBLNGPT2LMHeadModel",
         "RBLNGPT2LMHeadModelConfig",
+        "RBLNGptOssForCausalLM",
+        "RBLNGptOssForCausalLMConfig",
         "RBLNGroundingDinoDecoder",
         "RBLNGroundingDinoDecoderConfig",
         "RBLNGroundingDinoForObjectDetection",
@@ -114,6 +122,8 @@ _import_structure = {
         "RBLNLlamaForCausalLMConfig",
         "RBLNLlamaModel",
         "RBLNLlamaModelConfig",
+        "RBLNMixtralForCausalLM",
+        "RBLNMixtralForCausalLMConfig",
         "RBLNOPTForCausalLM",
         "RBLNOPTForCausalLMConfig",
         "RBLNLlavaForConditionalGeneration",
@@ -140,14 +150,24 @@ _import_structure = {
         "RBLNPixtralVisionModelConfig",
         "RBLNPhiModel",
         "RBLNPhiModelConfig",
+        "RBLNPaliGemmaForConditionalGeneration",
+        "RBLNPaliGemmaForConditionalGenerationConfig",
+        "RBLNPaliGemmaModel",
+        "RBLNPaliGemmaModelConfig",
         "RBLNQwen2ForCausalLM",
         "RBLNQwen2ForCausalLMConfig",
         "RBLNQwen2_5_VisionTransformerPretrainedModel",
         "RBLNQwen2_5_VisionTransformerPretrainedModelConfig",
         "RBLNQwen2_5_VLForConditionalGeneration",
         "RBLNQwen2_5_VLForConditionalGenerationConfig",
+        "RBLNQwen3MoeForCausalLM",
+        "RBLNQwen3MoeForCausalLMConfig",
+        "RBLNQwen2_5_VLModel",
+        "RBLNQwen2_5_VLModelConfig",
         "RBLNQwen2Model",
         "RBLNQwen2ModelConfig",
+        "RBLNQwen2MoeForCausalLM",
+        "RBLNQwen2MoeForCausalLMConfig",
         "RBLNQwen3ForCausalLM",
         "RBLNQwen3ForCausalLMConfig",
         "RBLNQwen3Model",
@@ -156,6 +176,8 @@ _import_structure = {
         "RBLNQwen2VisionTransformerPretrainedModelConfig",
         "RBLNQwen2VLForConditionalGeneration",
         "RBLNQwen2VLForConditionalGenerationConfig",
+        "RBLNQwen2VLModel",
+        "RBLNQwen2VLModelConfig",
         "RBLNResNetForImageClassification",
         "RBLNResNetForImageClassificationConfig",
         "RBLNRobertaForMaskedLM",
@@ -388,12 +410,18 @@ if TYPE_CHECKING:
         RBLNDecoderOnlyModelForCausalLMConfig,
         RBLNDepthAnythingForDepthEstimation,
         RBLNDepthAnythingForDepthEstimationConfig,
+        RBLNDetrForObjectDetection,
+        RBLNDetrForObjectDetectionConfig,
         RBLNDistilBertForQuestionAnswering,
         RBLNDistilBertForQuestionAnsweringConfig,
         RBLNDPTForDepthEstimation,
         RBLNDPTForDepthEstimationConfig,
         RBLNExaoneForCausalLM,
         RBLNExaoneForCausalLMConfig,
+        RBLNGemma2ForCausalLM,
+        RBLNGemma2ForCausalLMConfig,
+        RBLNGemma2Model,
+        RBLNGemma2ModelConfig,
         RBLNGemma3ForCausalLM,
         RBLNGemma3ForCausalLMConfig,
         RBLNGemma3ForConditionalGeneration,
@@ -406,6 +434,8 @@ if TYPE_CHECKING:
         RBLNGPT2LMHeadModelConfig,
         RBLNGPT2Model,
         RBLNGPT2ModelConfig,
+        RBLNGptOssForCausalLM,
+        RBLNGptOssForCausalLMConfig,
         RBLNGroundingDinoDecoder,
         RBLNGroundingDinoDecoderConfig,
         RBLNGroundingDinoEncoder,
@@ -432,10 +462,16 @@ if TYPE_CHECKING:
         RBLNMistralForCausalLMConfig,
         RBLNMistralModel,
         RBLNMistralModelConfig,
+        RBLNMixtralForCausalLM,
+        RBLNMixtralForCausalLMConfig,
         RBLNOPTForCausalLM,
         RBLNOPTForCausalLMConfig,
         RBLNOPTModel,
         RBLNOPTModelConfig,
+        RBLNPaliGemmaForConditionalGeneration,
+        RBLNPaliGemmaForConditionalGenerationConfig,
+        RBLNPaliGemmaModel,
+        RBLNPaliGemmaModelConfig,
         RBLNPegasusForConditionalGeneration,
         RBLNPegasusForConditionalGenerationConfig,
         RBLNPegasusModel,
@@ -450,18 +486,26 @@ if TYPE_CHECKING:
         RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
         RBLNQwen2_5_VLForConditionalGeneration,
         RBLNQwen2_5_VLForConditionalGenerationConfig,
+        RBLNQwen2_5_VLModel,
+        RBLNQwen2_5_VLModelConfig,
         RBLNQwen2ForCausalLM,
         RBLNQwen2ForCausalLMConfig,
         RBLNQwen2Model,
         RBLNQwen2ModelConfig,
+        RBLNQwen2MoeForCausalLM,
+        RBLNQwen2MoeForCausalLMConfig,
         RBLNQwen2VisionTransformerPretrainedModel,
         RBLNQwen2VisionTransformerPretrainedModelConfig,
         RBLNQwen2VLForConditionalGeneration,
         RBLNQwen2VLForConditionalGenerationConfig,
+        RBLNQwen2VLModel,
+        RBLNQwen2VLModelConfig,
         RBLNQwen3ForCausalLM,
         RBLNQwen3ForCausalLMConfig,
         RBLNQwen3Model,
         RBLNQwen3ModelConfig,
+        RBLNQwen3MoeForCausalLM,
+        RBLNQwen3MoeForCausalLMConfig,
         RBLNResNetForImageClassification,
         RBLNResNetForImageClassificationConfig,
         RBLNRobertaForMaskedLM,

optimum/rbln/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.9.4a2'
-__version_tuple__ = version_tuple = (0, 9, 4, 'a2')
+__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 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]: ...
@@ -92,7 +116,7 @@ class RBLNCompileConfig:
                 and isinstance(item[0], str)  # name
                 and isinstance(item[1], (tuple, list))  # shape
                 and all(isinstance(x, int) for x in item[1])
-                and isinstance(item[2], str)  # dtype
+                and (isinstance(item[2], str) or isinstance(item[2], torch.dtype))  # dtype
                 for item in input_info
             )
 
@@ -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):
@@ -524,8 +553,8 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     non_save_attributes = [
         "_frozen",
         "_runtime_options",
-        "torch_dtype",
         "npu",
+        "dtype",
         "tensor_parallel_size",
         "create_runtimes",
         "device",
@@ -650,6 +679,14 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
 
         super().__setattr__(key, value)
 
+    @deprecate_kwarg(
+        old_name="_torch_dtype",
+        new_name="dtype",
+        version="0.12.0",
+        deprecated_type=torch.dtype,
+        value_replacer=RBLNCompileConfig.normalize_dtype,
+        raise_if_greater_or_equal_version=False,
+    )
     def __init__(
         self,
         cls_name: Optional[str] = None,
@@ -661,7 +698,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         tensor_parallel_size: Optional[int] = None,
         timeout: Optional[int] = None,
         optimum_rbln_version: Optional[str] = None,
-        _torch_dtype: Optional[str] = None,
+        dtype: Optional[Union[str, torch.dtype]] = None,
         _compile_cfgs: Optional[List[RBLNCompileConfig]] = None,
         *,
         optimize_host_memory: Optional[bool] = None,
@@ -680,7 +717,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             tensor_parallel_size (Optional[int]): Size for tensor parallelism to distribute the model across devices.
             timeout (Optional[int]): The timeout for the runtime in seconds. If it isn't provided, it will be set to 60 by default.
             optimum_rbln_version (Optional[str]): The optimum-rbln version used for this configuration.
-            _torch_dtype (Optional[str]): The data type to use for the model.
+            dtype (Optional[Union[str, torch.dtype]]): The data type to use for the model.
             _compile_cfgs (List[RBLNCompileConfig]): List of compilation configurations for the model.
             kwargs: Additional keyword arguments.
 
@@ -710,7 +747,9 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         self.npu = npu
         self.tensor_parallel_size = tensor_parallel_size
 
-        self._torch_dtype = _torch_dtype or "float32"
+        if dtype is not None and isinstance(dtype, torch.dtype):
+            dtype = RBLNCompileConfig.normalize_dtype(dtype)
+        self._dtype = dtype or "float32"
         self.optimum_rbln_version = optimum_rbln_version
         if self.optimum_rbln_version is None:
             self.optimum_rbln_version = __version__
@@ -743,14 +782,24 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
 
     @property
     def torch_dtype(self):
-        return getattr(torch, self._torch_dtype)
+        logger.warning_once("`torch_dtype` is deprecated. Use `dtype` instead.")
+        return self.dtype
 
     @torch_dtype.setter
     def torch_dtype(self, torch_dtype: Union[str, torch.dtype]):
-        if isinstance(torch_dtype, torch.dtype):
-            torch_dtype = RBLNCompileConfig.normalize_dtype(torch_dtype)
+        logger.warning_once("`torch_dtype` is deprecated. Use `dtype` instead.")
+        self.dtype = torch_dtype
+
+    @property
+    def dtype(self):
+        return getattr(torch, self._dtype)
 
-        self._torch_dtype = torch_dtype
+    @dtype.setter
+    def dtype(self, dtype: Union[str, torch.dtype]):
+        if isinstance(dtype, torch.dtype):
+            dtype = RBLNCompileConfig.normalize_dtype(dtype)
+
+        self._dtype = dtype
 
     @property
     def rbln_model_cls_name(self) -> str:
@@ -774,10 +823,15 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             if isinstance(value, RBLNSerializableConfigProtocol):
                 # Convert nested RBLNModelConfig to its serializable form
                 serializable_map[key] = value._prepare_for_serialization()
+            elif key == "_dtype":
+                serializable_map["dtype"] = value
+            elif isinstance(value, list) and all(isinstance(item, RBLNSerializableConfigProtocol) for item in value):
+                serializable_map[key] = [item._prepare_for_serialization() for item in value]
             elif key == "_compile_cfgs":
                 serializable_map[key] = [cfg.asdict() for cfg in value]
             else:
                 serializable_map[key] = value
+
         return serializable_map
 
     def __repr__(self):
@@ -825,18 +879,12 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             if not isinstance(submodule_config, RBLNModelConfig):
                 raise ValueError(f"`{submodule_name}` must be an instance of `RBLNModelConfig` before freezing.")
 
-            if not submodule_config.is_frozen():
-                raise ValueError(f"`{submodule_name}` config must be frozen before freezing super config.")
-
         self._frozen = True
 
     def is_frozen(self):
         return self._frozen
 
     def save(self, path: str):
-        if not self._frozen:
-            raise RuntimeError("`RBLNModelConfig` is not frozen. Please call `set_compile_cfgs` first.")
-
         # save as json file without runtime attributes
         path = Path(path)
         if path.is_dir():
@@ -847,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.
@@ -864,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)
+
+            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
 
-        return cls(**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 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(
@@ -974,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.