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

optimum/rbln/utils/hub.py

@@ -16,7 +16,8 @@ import json
 from pathlib import Path
 from typing import List, Optional, Union
 
-from huggingface_hub import HfApi, get_token, hf_hub_download
+from huggingface_hub import HfApi, get_token, hf_hub_download, try_to_load_from_cache
+from huggingface_hub.errors import LocalEntryNotFoundError
 
 
 def pull_compiled_model_from_hub(
@@ -29,6 +30,97 @@ def pull_compiled_model_from_hub(
     local_files_only: bool,
 ) -> Path:
     """Pull model files from the HuggingFace Hub."""
+    config_filename = "rbln_config.json" if subfolder == "" else f"{subfolder}/rbln_config.json"
+
+    # Try to find config file in cache first.
+    config_cache_path = try_to_load_from_cache(
+        repo_id=str(model_id),
+        filename=config_filename,
+        revision=revision,
+        cache_dir=cache_dir,
+    )
+
+    # If config is cached and we're not forcing download, try to use cached files
+    if config_cache_path and isinstance(config_cache_path, str) and not force_download:
+        config_path = Path(config_cache_path)
+        if config_path.exists():
+            cache_dir_path = config_path.parent
+
+            # Look for .rbln files in the same directory
+            pattern_rbln = "*.rbln"
+            rbln_files = list(cache_dir_path.glob(pattern_rbln))
+
+            # Validate files found in cache
+            rbln_config_filenames = [config_path] if config_path.exists() else []
+            validate_files(rbln_files, rbln_config_filenames, f"cached repository {model_id}")
+
+            # If local_files_only is True, return cached directory without API call
+            if local_files_only:
+                return cache_dir_path
+
+            # If local_files_only is False, ensure all files are downloaded
+            # Download config file (will use cache if available, download if missing)
+            rbln_config_cache_path = hf_hub_download(
+                repo_id=model_id,
+                filename=config_filename,
+                token=token,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                local_files_only=False,
+            )
+            cache_dir_path = Path(rbln_config_cache_path).parent
+
+            # Download all .rbln files found in cache (hf_hub_download will use cache if available)
+            for rbln_file in rbln_files:
+                filename = rbln_file.name if subfolder == "" else f"{subfolder}/{rbln_file.name}"
+                try:
+                    hf_hub_download(
+                        repo_id=model_id,
+                        filename=filename,
+                        token=token,
+                        revision=revision,
+                        cache_dir=cache_dir,
+                        force_download=force_download,
+                        local_files_only=False,
+                    )
+                except LocalEntryNotFoundError:
+                    # File might not exist in repo, skip it
+                    pass
+
+            # Note: We skip the API call here since we're using cached files
+            # If there are additional files in the repo that aren't cached,
+            # they won't be downloaded.
+            # If the user needs all files, they should use force_download=True
+            return cache_dir_path
+
+    # If local_files_only is True and config not found in cache, try to download with local_files_only
+    if local_files_only:
+        try:
+            rbln_config_cache_path = hf_hub_download(
+                repo_id=model_id,
+                filename=config_filename,
+                token=token,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                local_files_only=True,
+            )
+            cache_dir_path = Path(rbln_config_cache_path).parent
+            rbln_files = list(cache_dir_path.glob("*.rbln"))
+            rbln_config_filenames = [Path(rbln_config_cache_path)] if Path(rbln_config_cache_path).exists() else []
+            validate_files(rbln_files, rbln_config_filenames, f"cached repository {model_id}")
+            return cache_dir_path
+        except LocalEntryNotFoundError as err:
+            raise FileNotFoundError(
+                f"Could not find compiled model files for {model_id} in local cache. "
+                f"Set local_files_only=False to download from HuggingFace Hub."
+            ) from err
+
+    # List files from repository. This only happens when:
+    # 1. Config is not cached, OR
+    # 2. force_download=True, OR
+    # 3. local_files_only=False and we need to discover all files in the repo
     huggingface_token = _get_huggingface_token(token)
     repo_files = list(
         map(
@@ -51,7 +143,6 @@ def pull_compiled_model_from_hub(
         rbln_config_cache_path = hf_hub_download(
             repo_id=model_id,
             filename=filename,
-            subfolder=subfolder,
             token=token,
             revision=revision,
             cache_dir=cache_dir,

optimum/rbln/utils/import_utils.py

@@ -136,7 +136,22 @@ def is_rbln_available() -> bool:
 
 def check_version_compats() -> None:
     warnings.filterwarnings(action="always", category=ImportWarning, module="optimum.rbln")
-    my_version = importlib.metadata.version("optimum-rbln")
+    try:
+        my_version = importlib.metadata.version("optimum-rbln")
+    except importlib.metadata.PackageNotFoundError:
+        # Common dev case: running from source (e.g. PYTHONPATH=src) without installing the package.
+        # package metadata doesn't exist, so fall back to the in-repo version file.
+        try:
+            from optimum.rbln.__version__ import __version__ as my_version  # type: ignore
+        except Exception:
+            warnings.warn(
+                "Could not determine optimum-rbln version (package metadata missing). "
+                "If you are running from source, consider `pip install -e .` to install metadata.",
+                ImportWarning,
+                stacklevel=2,
+            )
+            return
+
     target_version = list(filter(lambda v: Version(my_version) >= Version(v), RBLN_VERSION_COMPATS.keys()))[0]
     for compat in RBLN_VERSION_COMPATS[target_version]:
         try:

optimum/rbln/utils/runtime_utils.py

@@ -20,6 +20,10 @@ import rebel
 import torch
 
 
+def is_compiler_supports_buffer_resize() -> bool:
+    return hasattr(rebel.RBLNCompiledModel, "exp_multiply_buffer_size")
+
+
 def get_available_dram(npu: Optional[str] = None) -> int:
     """
     Get the available DRAM size of the specified NPU.
@@ -75,12 +79,6 @@ def tp_and_devices_are_ok(
     if tensor_parallel_size is None:
         tensor_parallel_size = 1
 
-    if rebel.device_count() < tensor_parallel_size:
-        return (
-            f"Tensor parallel size {tensor_parallel_size} is greater than "
-            f"the number of available devices {rebel.device_count()}."
-        )
-
     if device is None:
         device = list(range(tensor_parallel_size))
     elif isinstance(device, int):
@@ -100,9 +98,15 @@ def tp_and_devices_are_ok(
             return None
         if rebel.get_npu_name(device_id) is None:
             return (
-                f"Device {device_id} is not a valid NPU device. Please check your NPU status with 'rbln-stat' command."
+                f"Device {device_id} is not a valid NPU device. Please check your NPU status with 'rbln-smi' command."
             )
 
+    if rebel.device_count() < tensor_parallel_size:
+        return (
+            f"Tensor parallel size {tensor_parallel_size} is greater than "
+            f"the number of available devices {rebel.device_count()}."
+        )
+
     if npu is not None:
         for device_id in device:
             npu_name = rebel.get_npu_name(device_id)
@@ -181,7 +185,7 @@ class UnavailableRuntime:
             "This model was loaded with create_runtimes=False. To use this model for inference:\n"
             "1. Load the model with runtime creation enabled:\n"
             "   model = RBLNModel.from_pretrained(..., rbln_create_runtimes=True)\n"
-            "2. Ensure your NPU hardware is properly configured (check with 'rbln-stat' command)\n"
+            "2. Ensure your NPU hardware is properly configured (check with 'rbln-smi' command)\n"
             "3. If you're on a machine without NPU hardware, you need to transfer the model files\n"
             "   to a compatible system with NPU support."
         )

optimum/rbln/utils/submodule.py

@@ -61,12 +61,25 @@ class SubModulesMixin:
     ):
         return rbln_config
 
+    @classmethod
+    def _update_submodule_rbln_config(
+        cls,
+        submodule_name: str,
+        submodule_cls: Type["RBLNModel"],
+        model: "PreTrainedModel",
+        submodule_config: PretrainedConfig,
+        submodule_rbln_config: RBLNModelConfig,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+    ):
+        return submodule_rbln_config
+
     @classmethod
     def _export_submodules_from_model(
         cls, model: "PreTrainedModel", model_save_dir: str, rbln_config: RBLNModelConfig, **kwargs
     ) -> List["RBLNModel"]:
         rbln_submodules = []
         submodule_prefix = getattr(cls, "_rbln_submodule_prefix", None)
+        submodule_postfix = getattr(cls, "_rbln_submodule_postfix", None)
         preprocessors = kwargs.pop("preprocessors", [])
 
         for submodule in cls._rbln_submodules:
@@ -74,6 +87,9 @@ class SubModulesMixin:
             if submodule_prefix is not None:
                 torch_submodule: PreTrainedModel = getattr(model, submodule_prefix)
                 torch_submodule = getattr(torch_submodule, submodule_name)
+            elif submodule_postfix is not None:
+                torch_submodule: PreTrainedModel = getattr(model, submodule_name)
+                torch_submodule = getattr(torch_submodule, submodule_postfix)
             else:
                 torch_submodule: PreTrainedModel = getattr(model, submodule_name)
 
@@ -92,6 +108,14 @@ class SubModulesMixin:
                 filtered_kwargs["cls_name"] = submodule_config_cls.__name__
                 submodule_rbln_config = submodule_config_cls(**filtered_kwargs)
 
+            submodule_rbln_config = cls._update_submodule_rbln_config(
+                submodule_name=submodule_name,
+                submodule_cls=submodule_cls,
+                model=model,
+                submodule_config=torch_submodule.config,
+                submodule_rbln_config=submodule_rbln_config,
+                preprocessors=preprocessors,
+            )
             setattr(rbln_config, submodule_name, submodule_rbln_config)
             submodule_rbln_config = submodule_cls._update_submodule_config(model, submodule_rbln_config, preprocessors)
 

{optimum_rbln-0.9.4a2.dist-info → optimum_rbln-0.10.0.post1.dist-info}/METADATA

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.9.4a2
+Version: 0.10.0.post1
 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
-Project-URL: Repository, https://github.com/rebellions-sw/optimum-rbln
+Project-URL: Repository, https://github.com/rbln-sw/optimum-rbln
 Author-email: "Rebellions Inc." <support@rebellions.ai>
 License-Expression: Apache-2.0
 License-File: LICENSE
@@ -24,12 +24,12 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: <3.14,>=3.9
 Requires-Dist: accelerate>=1.0.1
-Requires-Dist: diffusers==0.35.2
+Requires-Dist: diffusers==0.36.0
 Requires-Dist: packaging>=24.1
 Requires-Dist: torch==2.8.0
 Requires-Dist: torchaudio<=2.8.0
 Requires-Dist: torchvision<=0.23.0
-Requires-Dist: transformers==4.57.1
+Requires-Dist: transformers==4.57.3
 Description-Content-Type: text/markdown
 
 
@@ -40,7 +40,7 @@ Description-Content-Type: text/markdown
 <img src="assets/rbln_logo.png" width="60%"/>
 
 [![PyPI version](https://badge.fury.io/py/optimum-rbln.svg)](https://badge.fury.io/py/optimum-rbln)
-[![License](https://img.shields.io/github/license/rebellions-sw/optimum-rbln)](https://github.com/rebellions-sw/optimum-rbln/blob/main/LICENSE)
+[![License](https://img.shields.io/github/license/rbln-sw/optimum-rbln)](https://github.com/rbln-sw/optimum-rbln/blob/main/LICENSE)
 [![Documentation](https://img.shields.io/badge/docs-available-brightgreen)](https://docs.rbln.ai/software/optimum/optimum_rbln.html)
 [![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-2.1-4baaaa.svg)](CODE_OF_CONDUCT.md)
 
@@ -113,7 +113,7 @@ pip install optimum-rbln --extra-index-url https://download.pytorch.org/whl/cpu
 The below command installs `optimum-rbln` along with its dependencies.
 
 ```bash
-git clone https://github.com/rebellions-sw/optimum-rbln.git
+git clone https://github.com/rbln-sw/optimum-rbln.git
 cd optimum-rbln
 ./scripts/uv-sync.sh
 ```