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

optimum/rbln/utils/depreacate_utils.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+import rebel
+
+from .logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def warn_deprecated_npu(npu: Optional[str] = None):
+    npu = npu or rebel.get_npu_name()
+    if npu == "RBLN-CA02":
+        logger.warning_once(
+            "Support for the RBLN-CA02 device is provided only up to optimum-rbln v0.8.0 and has reached end of life.",
+        )

optimum/rbln/utils/runtime_utils.py

@@ -14,7 +14,7 @@
 
 import re
 import threading
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, List, Optional, Union
 
 import rebel
 import torch
@@ -94,7 +94,7 @@ class RBLNPytorchRuntime:
     def __call__(self, *args: Any, **kwds: Any) -> Any:
         return self.forward(*args, **kwds)
 
-    def forward(self, *args: List["torch.Tensor"], **kwargs: Dict[str, "torch.Tensor"]):
+    def forward(self, *args: List["torch.Tensor"], **kwargs: "torch.Tensor"):
         # filtering useless args or kwarg such as None.
         args = list(filter(lambda arg: isinstance(arg, torch.Tensor), args))
         kwargs = dict(filter(lambda kwarg: isinstance(kwarg[1], torch.Tensor) or kwarg[0] == "out", kwargs.items()))
@@ -142,7 +142,7 @@ class UnavailableRuntime:
         """Returns an iterator with self as the only item."""
         return iter([self])
 
-    def forward(self, *args: List["torch.Tensor"], **kwargs: Dict[str, "torch.Tensor"]):
+    def forward(self, *args: List["torch.Tensor"], **kwargs: "torch.Tensor"):
         """Raises a detailed RuntimeError explaining why inference cannot be performed."""
         raise RuntimeError(
             "Cannot perform inference: RBLN runtime is not available.\n\n"
@@ -167,33 +167,44 @@ class ContextRblnConfig:
         device=None,
         device_map=None,
         create_runtimes=None,
-        optimize_host_mem=None,
         activate_profiler=None,
         timeout=None,
     ):
         self.device = device
         self.device_map = device_map
         self.create_runtimes = create_runtimes
-        self.optimize_host_mem = optimize_host_mem
         self.activate_profiler = activate_profiler
         self.timeout = timeout
+        self._previous_context = None
 
     def __enter__(self):
-        self._local.device = self.device
-        self._local.device_map = self.device_map
-        self._local.create_runtimes = self.create_runtimes
-        self._local.optimize_host_memory = self.optimize_host_mem
-        self._local.activate_profiler = self.activate_profiler
-        self._local.timeout = self.timeout
+        self._previous_context = {
+            "device": getattr(self._local, "device", None),
+            "device_map": getattr(self._local, "device_map", None),
+            "create_runtimes": getattr(self._local, "create_runtimes", None),
+            "activate_profiler": getattr(self._local, "activate_profiler", None),
+            "timeout": getattr(self._local, "timeout", None),
+        }
+
+        if self.device is not None:
+            self._local.device = self.device
+        if self.device_map is not None:
+            self._local.device_map = self.device_map
+        if self.create_runtimes is not None:
+            self._local.create_runtimes = self.create_runtimes
+        if self.activate_profiler is not None:
+            self._local.activate_profiler = self.activate_profiler
+        if self.timeout is not None:
+            self._local.timeout = self.timeout
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self._local.device = None
-        self._local.device_map = None
-        self._local.create_runtimes = None
-        self._local.optimize_host_memory = None
-        self._local.activate_profiler = None
-        self._local.timeout = None
+        if self._previous_context is not None:
+            self._local.device = self._previous_context["device"]
+            self._local.device_map = self._previous_context["device_map"]
+            self._local.create_runtimes = self._previous_context["create_runtimes"]
+            self._local.activate_profiler = self._previous_context["activate_profiler"]
+            self._local.timeout = self._previous_context["timeout"]
 
     @classmethod
     def get_current_context(cls):
@@ -201,7 +212,6 @@ class ContextRblnConfig:
             "device": getattr(cls._local, "device", None),
             "device_map": getattr(cls._local, "device_map", None),
             "create_runtimes": getattr(cls._local, "create_runtimes", None),
-            "optimize_host_memory": getattr(cls._local, "optimize_host_memory", None),
             "activate_profiler": getattr(cls._local, "activate_profiler", None),
             "timeout": getattr(cls._local, "timeout", None),
         }

optimum/rbln/utils/submodule.py

@@ -13,16 +13,16 @@
 # limitations under the License.
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Type
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 from transformers import PretrainedConfig
 
-from ..configuration_utils import RBLNModelConfig
+from ..configuration_utils import RBLNModelConfig, get_rbln_config_class
 from ..utils.model_utils import get_rbln_model_cls
 
 
 if TYPE_CHECKING:
-    from transformers import PreTrainedModel
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
 
     from ..modeling import RBLNModel
 
@@ -42,7 +42,21 @@ class SubModulesMixin:
             setattr(self, submodule_meta["name"], submodule)
 
     @classmethod
-    def _update_submodule_config(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
+    def _get_submodule_config_class(
+        cls, cls_name: str, submodule_rbln_config: Dict[str, Any]
+    ) -> Type[RBLNModelConfig]:
+        if isinstance(submodule_rbln_config, dict) and "cls_name" in submodule_rbln_config:
+            config_cls_name = submodule_rbln_config["cls_name"]
+            return get_rbln_config_class(config_cls_name)
+        return get_rbln_config_class(f"RBLN{cls_name}Config")
+
+    @classmethod
+    def _update_submodule_config(
+        cls,
+        model: "PreTrainedModel",
+        rbln_config: RBLNModelConfig,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+    ):
         return rbln_config
 
     @classmethod
@@ -51,6 +65,7 @@ class SubModulesMixin:
     ) -> List["RBLNModel"]:
         rbln_submodules = []
         submodule_prefix = getattr(cls, "_rbln_submodule_prefix", None)
+        preprocessors = kwargs.pop("preprocessors", [])
 
         for submodule in cls._rbln_submodules:
             submodule_name = submodule["name"]
@@ -63,13 +78,20 @@ class SubModulesMixin:
             cls_name = torch_submodule.__class__.__name__
             submodule_cls: Type["RBLNModel"] = get_rbln_model_cls(f"RBLN{cls_name}")
             submodule_rbln_config = getattr(rbln_config, submodule_name) or {}
+            submodule_config_cls = cls._get_submodule_config_class(cls_name, submodule_rbln_config)
 
             if isinstance(submodule_rbln_config, dict):
-                submodule_rbln_config_class = submodule_cls.get_rbln_config_class()
-                submodule_rbln_config = submodule_rbln_config_class(**submodule_rbln_config)
-                setattr(rbln_config, submodule_name, submodule_rbln_config)
-
-            submodule_rbln_config = submodule_cls._update_submodule_config(model, submodule_rbln_config)
+                filtered_kwargs = rbln_config.filter_parameters(submodule_config_cls, submodule_rbln_config)
+                filtered_kwargs["cls_name"] = submodule_config_cls.__name__
+                submodule_rbln_config = submodule_config_cls(**filtered_kwargs)
+            elif not isinstance(submodule_rbln_config, submodule_config_cls):
+                config_dict = {k: v for k, v in submodule_rbln_config.__dict__.items() if not k.startswith("_")}
+                filtered_kwargs = rbln_config.filter_parameters(submodule_config_cls, config_dict)
+                filtered_kwargs["cls_name"] = submodule_config_cls.__name__
+                submodule_rbln_config = submodule_config_cls(**filtered_kwargs)
+
+            setattr(rbln_config, submodule_name, submodule_rbln_config)
+            submodule_rbln_config = submodule_cls._update_submodule_config(model, submodule_rbln_config, preprocessors)
 
             rbln_submodule = submodule_cls.from_model(
                 model=torch_submodule,

{optimum_rbln-0.8.2a4.dist-info → optimum_rbln-0.9.3rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.2a4
+Version: 0.9.3rc0
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai
@@ -20,15 +20,16 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Requires-Python: <3.13,>=3.9
+Requires-Python: <3.14,>=3.9
 Requires-Dist: accelerate>=1.0.1
-Requires-Dist: diffusers==0.34.0
+Requires-Dist: diffusers==0.35.1
 Requires-Dist: packaging>=24.1
-Requires-Dist: torch==2.7.0
-Requires-Dist: torchaudio<=2.7.0
-Requires-Dist: torchvision<=0.22.0
-Requires-Dist: transformers==4.51.3
+Requires-Dist: torch==2.8.0
+Requires-Dist: torchaudio<=2.8.0
+Requires-Dist: torchvision<=0.23.0
+Requires-Dist: transformers==4.57.1
 Description-Content-Type: text/markdown