optimum-rbln 0.8.1rc0__tar.gz → 0.8.2__tar.gz

optimum_rbln-0.8.2/.github/version.yaml

@@ -0,0 +1 @@
+rebel_compiler_version: 0.8.2.dev187+g9f5b6c9b

optimum_rbln-0.8.2/.github/workflows/rbln_optimum_pytest.yaml

@@ -0,0 +1,138 @@
+name: Optimum-rbln / PR / Pytest
+
+on:
+  workflow_call:
+    inputs:
+      ref:
+        description: "ref to checkout"
+        required: false
+        type: string
+      pr_number:
+        description: "PR number to run"
+        required: false
+        type: string
+      rebel_compiler_version:
+        description: "rebel_compiler version to run"
+        required: true
+        type: string
+      test_level:
+        description: "Test level for OPTIMUM_RBLN_TEST_LEVEL (default, full, essential)"
+        required: false
+        type: string
+        default: "default"
+      enable_hf_hub_tests:
+        description: "Whether to enable HF Hub tests (requires HF credentials)"
+        required: false
+        type: boolean
+        default: false
+      fail_fast:
+        description: "Whether to fail fast when one matrix job fails"
+        required: false
+        type: boolean
+        default: true
+
+env:
+  REBEL_PYPI_ENDPOINT: ${{ vars.REBEL_PYPI_INTERNAL_ENDPOINT }}
+  REBEL_PYPI_USERNAME: ${{ secrets.REBEL_PYPI_USERNAME }}
+  REBEL_PYPI_PASSWORD: ${{ secrets.REBEL_PYPI_PASSWORD }}
+  HF_HOME: ${{ secrets.HF_HOME }}
+  HF_USER_ID: ${{ inputs.enable_hf_hub_tests && secrets.HF_USER_ID || '' }}
+  HF_AUTH_TOKEN: ${{ inputs.enable_hf_hub_tests && secrets.HF_AUTH_TOKEN || '' }}
+
+jobs:
+  pytest:
+    name: Pytest (${{ matrix.test_type }})
+    runs-on: vm-ci-0
+    strategy:
+      fail-fast: ${{ inputs.fail_fast }}
+      matrix:
+        test_type: [config, transformers, diffusers, llm]
+    steps:
+      - name: Checkout the optimum-rbln repository
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_number && format('refs/pull/{0}/merge', inputs.pr_number) || inputs.ref }}
+          submodules: recursive
+          fetch-depth: 0
+
+      - name: Get commit message if not provided
+        id: get_commit_message
+        if: ${{ inputs.commit_message == '' }}
+        run: |
+          COMMIT_MESSAGE=$(git log -1 --pretty=%B)
+          echo "message<<EOF" >> $GITHUB_OUTPUT
+          echo "$COMMIT_MESSAGE" >> $GITHUB_OUTPUT
+          echo "EOF" >> $GITHUB_OUTPUT
+
+      - name: Check if test should be skipped
+        id: should_skip
+        run: |
+          COMMIT_MESSAGE="${{ steps.get_commit_message.outputs.message }}"
+          SKIP=false
+          
+          case "${{ matrix.test_type }}" in
+            "transformers")
+              if [[ "$COMMIT_MESSAGE" == *"[skip-transformers]"* ]]; then
+                SKIP=true
+              fi
+              ;;
+            "diffusers")
+              if [[ "$COMMIT_MESSAGE" == *"[skip-diffusers]"* ]]; then
+                SKIP=true
+              fi
+              ;;
+            "llm")
+              if [[ "$COMMIT_MESSAGE" == *"[skip-llms]"* ]]; then
+                SKIP=true
+              fi
+              ;;
+          esac
+          
+          echo "skip=$SKIP" >> $GITHUB_OUTPUT
+
+      - name: Setup uv
+        if: steps.should_skip.outputs.skip != 'true'
+        uses: astral-sh/setup-uv@v3
+        with:
+          enable-cache: true
+          cache-dependency-glob: "uv.lock"
+
+      - name: Setup Python
+        if: steps.should_skip.outputs.skip != 'true'
+        run: uv python install 3.9
+
+      - name: Install optimum-rbln with tests group dependencies
+        if: steps.should_skip.outputs.skip != 'true'
+        run: |
+          uv sync --frozen --group tests --reinstall-package optimum-rbln
+
+      - name: Install rebel-compiler
+        if: steps.should_skip.outputs.skip != 'true'
+        run: |
+          PYPI_URL=$(echo ${{ env.REBEL_PYPI_ENDPOINT }} | sed "s/\/\//\0${{ env.REBEL_PYPI_USERNAME}}:${{ env.REBEL_PYPI_PASSWORD}}@/")
+          uv pip install --extra-index-url $PYPI_URL rebel-compiler==${{ inputs.rebel_compiler_version }}
+
+      - name: Run pytest
+        if: steps.should_skip.outputs.skip != 'true'
+        env:
+          OPTIMUM_RBLN_TEST_LEVEL: ${{ inputs.test_level }}
+        run: |
+          case "${{ matrix.test_type }}" in
+            "config")
+              uv run --no-sync pytest tests/test_config.py -vv --durations 0
+              ;;
+            "transformers")
+              uv run --no-sync pytest tests/test_transformers.py -vv --durations 0
+              ;;
+            "diffusers")
+              uv run --no-sync pytest tests/test_diffusers.py -vv --durations 0
+              ;;
+            "llm")
+              uv run --no-sync pytest tests/test_llm.py -vv --durations 0
+              ;;
+          esac
+
+      - name: Skip message
+        if: steps.should_skip.outputs.skip == 'true'
+        run: |
+          echo "Found [skip-${{ matrix.test_type }}] in commit message, skipping CI"

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/.github/workflows/rbln_scheduled_test.yaml

@@ -47,6 +47,8 @@ jobs:
       ref: main
       rebel_compiler_version: ${{ needs.check-compiler.outputs.compiler_version }}
       test_level: "full"
+      enable_hf_hub_tests: true
+      fail_fast: false
     secrets: inherit
 
   optimum-rbln-inference-test:

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/.github/workflows/rbln_trigger_on_pr.yaml

@@ -92,13 +92,4 @@ jobs:
       ref: ${{ github.event.pull_request.head.sha }}
       rebel_compiler_version: ${{ needs.check-compiler.outputs.compiler_version }}
       test_level: "default"
-    secrets: inherit
-
-  optimum-rbln-inference-test:
-    needs: [check-skip-ci, check-code-quality, test-docstrings, check-compiler, check-team-member]
-    if: ${{ needs.check-skip-ci.outputs.should_skip != 'true' && needs.check-team-member.outputs.is_team_member == 'true' }}
-    uses: ./.github/workflows/rbln_optimum_inference_test.yaml
-    with:
-      ref: ${{ github.event.pull_request.head.sha }}
-      rebel_compiler_version: ${{ needs.check-compiler.outputs.compiler_version }}
     secrets: inherit

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.8.1rc0
+Version: 0.8.2
 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

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/pyproject.toml

@@ -55,6 +55,7 @@ tests = [
     "protobuf>=5.27.2",
     "soundfile>=0.13.1",
     "librosa>=0.11.0",
+    "simphile>=1.0.2",
 ]
 quality = [
     "ruff>=0.3.3",

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/__init__.py

@@ -72,6 +72,8 @@ _import_structure = {
         "RBLNCLIPVisionModelWithProjectionConfig",
         "RBLNColPaliForRetrieval",
         "RBLNColPaliForRetrievalConfig",
+        "RBLNDecoderOnlyModelConfig",
+        "RBLNDecoderOnlyModel",
         "RBLNDecoderOnlyModelForCausalLM",
         "RBLNDecoderOnlyModelForCausalLMConfig",
         "RBLNDistilBertForQuestionAnswering",
@@ -80,12 +82,16 @@ _import_structure = {
         "RBLNDPTForDepthEstimationConfig",
         "RBLNExaoneForCausalLM",
         "RBLNExaoneForCausalLMConfig",
+        "RBLNGemmaModel",
+        "RBLNGemmaModelConfig",
         "RBLNGemmaForCausalLM",
         "RBLNGemmaForCausalLMConfig",
         "RBLNGemma3ForCausalLM",
         "RBLNGemma3ForCausalLMConfig",
         "RBLNGemma3ForConditionalGeneration",
         "RBLNGemma3ForConditionalGenerationConfig",
+        "RBLNGPT2Model",
+        "RBLNGPT2ModelConfig",
         "RBLNGPT2LMHeadModel",
         "RBLNGPT2LMHeadModelConfig",
         "RBLNIdefics3VisionTransformer",
@@ -94,22 +100,44 @@ _import_structure = {
         "RBLNIdefics3VisionTransformerConfig",
         "RBLNLlamaForCausalLM",
         "RBLNLlamaForCausalLMConfig",
+        "RBLNLlamaModel",
+        "RBLNLlamaModelConfig",
         "RBLNOPTForCausalLM",
         "RBLNOPTForCausalLMConfig",
+        "RBLNLlavaForConditionalGeneration",
+        "RBLNLlavaForConditionalGenerationConfig",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNLlavaNextForConditionalGenerationConfig",
         "RBLNMidmLMHeadModel",
         "RBLNMidmLMHeadModelConfig",
+        "RBLNMistralModel",
+        "RBLNMistralModelConfig",
         "RBLNMistralForCausalLM",
         "RBLNMistralForCausalLMConfig",
+        "RBLNOPTModel",
+        "RBLNOPTModelConfig",
+        "RBLNPegasusForConditionalGeneration",
+        "RBLNPegasusForConditionalGenerationConfig",
+        "RBLNPegasusModel",
+        "RBLNPegasusModelConfig",
         "RBLNPhiForCausalLM",
         "RBLNPhiForCausalLMConfig",
+        "RBLNPixtralVisionModel",
+        "RBLNPixtralVisionModelConfig",
+        "RBLNPhiModel",
+        "RBLNPhiModelConfig",
         "RBLNQwen2ForCausalLM",
         "RBLNQwen2ForCausalLMConfig",
         "RBLNQwen2_5_VisionTransformerPretrainedModel",
         "RBLNQwen2_5_VisionTransformerPretrainedModelConfig",
         "RBLNQwen2_5_VLForConditionalGeneration",
         "RBLNQwen2_5_VLForConditionalGenerationConfig",
+        "RBLNQwen2Model",
+        "RBLNQwen2ModelConfig",
+        "RBLNQwen3ForCausalLM",
+        "RBLNQwen3ForCausalLMConfig",
+        "RBLNQwen3Model",
+        "RBLNQwen3ModelConfig",
         "RBLNResNetForImageClassification",
         "RBLNResNetForImageClassificationConfig",
         "RBLNRobertaForMaskedLM",
@@ -203,11 +231,7 @@ _import_structure = {
 }
 
 if TYPE_CHECKING:
-    from .configuration_utils import (
-        RBLNAutoConfig,
-        RBLNCompileConfig,
-        RBLNModelConfig,
-    )
+    from .configuration_utils import RBLNAutoConfig, RBLNCompileConfig, RBLNModelConfig
     from .diffusers import (
         RBLNAutoencoderKL,
         RBLNAutoencoderKLConfig,
@@ -273,10 +297,7 @@ if TYPE_CHECKING:
         RBLNVQModel,
         RBLNVQModelConfig,
     )
-    from .modeling import (
-        RBLNBaseModel,
-        RBLNModel,
-    )
+    from .modeling import RBLNBaseModel, RBLNModel
     from .transformers import (
         RBLNASTForAudioClassification,
         RBLNASTForAudioClassificationConfig,
@@ -319,6 +340,8 @@ if TYPE_CHECKING:
         RBLNCLIPVisionModelWithProjectionConfig,
         RBLNColPaliForRetrieval,
         RBLNColPaliForRetrievalConfig,
+        RBLNDecoderOnlyModel,
+        RBLNDecoderOnlyModelConfig,
         RBLNDecoderOnlyModelForCausalLM,
         RBLNDecoderOnlyModelForCausalLMConfig,
         RBLNDistilBertForQuestionAnswering,
@@ -333,30 +356,56 @@ if TYPE_CHECKING:
         RBLNGemma3ForConditionalGenerationConfig,
         RBLNGemmaForCausalLM,
         RBLNGemmaForCausalLMConfig,
+        RBLNGemmaModel,
+        RBLNGemmaModelConfig,
         RBLNGPT2LMHeadModel,
         RBLNGPT2LMHeadModelConfig,
+        RBLNGPT2Model,
+        RBLNGPT2ModelConfig,
         RBLNIdefics3ForConditionalGeneration,
         RBLNIdefics3ForConditionalGenerationConfig,
         RBLNIdefics3VisionTransformer,
         RBLNIdefics3VisionTransformerConfig,
         RBLNLlamaForCausalLM,
         RBLNLlamaForCausalLMConfig,
+        RBLNLlamaModel,
+        RBLNLlamaModelConfig,
+        RBLNLlavaForConditionalGeneration,
+        RBLNLlavaForConditionalGenerationConfig,
         RBLNLlavaNextForConditionalGeneration,
         RBLNLlavaNextForConditionalGenerationConfig,
         RBLNMidmLMHeadModel,
         RBLNMidmLMHeadModelConfig,
         RBLNMistralForCausalLM,
         RBLNMistralForCausalLMConfig,
+        RBLNMistralModel,
+        RBLNMistralModelConfig,
         RBLNOPTForCausalLM,
         RBLNOPTForCausalLMConfig,
+        RBLNOPTModel,
+        RBLNOPTModelConfig,
+        RBLNPegasusForConditionalGeneration,
+        RBLNPegasusForConditionalGenerationConfig,
+        RBLNPegasusModel,
+        RBLNPegasusModelConfig,
         RBLNPhiForCausalLM,
         RBLNPhiForCausalLMConfig,
+        RBLNPhiModel,
+        RBLNPhiModelConfig,
+        RBLNPixtralVisionModel,
+        RBLNPixtralVisionModelConfig,
         RBLNQwen2_5_VisionTransformerPretrainedModel,
         RBLNQwen2_5_VisionTransformerPretrainedModelConfig,
         RBLNQwen2_5_VLForConditionalGeneration,
         RBLNQwen2_5_VLForConditionalGenerationConfig,
         RBLNQwen2ForCausalLM,
         RBLNQwen2ForCausalLMConfig,
+        RBLNQwen2Model,
+        RBLNQwen2ModelConfig,
+        RBLNQwen3ForCausalLM,
+        RBLNQwen3ForCausalLMConfig,
+        RBLNQwen3Model,
+        RBLNQwen3ModelConfig,
         RBLNResNetForImageClassification,
         RBLNResNetForImageClassificationConfig,
         RBLNRobertaForMaskedLM,

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.8.1rc0'
-__version_tuple__ = version_tuple = (0, 8, 1, 'rc0')
+__version__ = version = '0.8.2'
+__version_tuple__ = version_tuple = (0, 8, 2)

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/configuration_utils.py

@@ -23,6 +23,7 @@ import numpy as np
 import torch
 
 from .__version__ import __version__
+from .utils.depreacate_utils import warn_deprecated_npu
 from .utils.logging import get_logger
 from .utils.runtime_utils import ContextRblnConfig
 
@@ -147,7 +148,7 @@ class RBLNCompileConfig:
         return asdict(self)
 
 
-RUNTIME_KEYWORDS = ["create_runtimes", "optimize_host_memory", "device", "device_map", "activate_profiler"]
+RUNTIME_KEYWORDS = ["create_runtimes", "optimize_host_memory", "device", "device_map", "activate_profiler", "timeout"]
 CONFIG_MAPPING: Dict[str, Type["RBLNModelConfig"]] = {}
 
 
@@ -481,6 +482,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         "device",
         "device_map",
         "activate_profiler",
+        "timeout",
     ]
     submodules: List[str] = []
     subclass_non_save_attributes = []
@@ -489,7 +491,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         self,
         submodule_config_cls: Type["RBLNModelConfig"],
         submodule_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> "RBLNModelConfig":
         # Initialize a submodule config from a dict or a RBLNModelConfig.
         # kwargs is specified from the predecessor config.
@@ -561,9 +563,10 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         activate_profiler: Optional[bool] = None,
         npu: Optional[str] = None,
         tensor_parallel_size: Optional[int] = None,
+        timeout: Optional[int] = None,
         optimum_rbln_version: Optional[str] = None,
         _compile_cfgs: List[RBLNCompileConfig] = [],
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Initialize a RBLN model configuration with runtime options and compile configurations.
@@ -577,6 +580,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             activate_profiler (Optional[bool]): Whether to activate the profiler for performance analysis.
             npu (Optional[str]): The NPU device name to use for compilation.
             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.
             _compile_cfgs (List[RBLNCompileConfig]): List of compilation configurations for the model.
             **kwargs: Additional keyword arguments.
@@ -599,6 +603,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         self._runtime_options["device"] = device
         self._runtime_options["device_map"] = device_map
         self._runtime_options["activate_profiler"] = activate_profiler
+        self._runtime_options["timeout"] = timeout
 
         # Automatically pass npu, tensor_parallel_size to compile_cfgs
         self.npu = npu
@@ -671,6 +676,9 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             compile_cfg.npu = self.npu
             compile_cfg.tensor_parallel_size = self.tensor_parallel_size
 
+        target_npu = self.npu or next((cfg.npu for cfg in self._compile_cfgs if cfg.npu is not None), None)
+        warn_deprecated_npu(target_npu)
+
     def freeze(self):
         if self._frozen:
             raise RuntimeError(f"`{self.__class__.__name__}` is already frozen.")
@@ -709,7 +717,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
             json.dump(serializable_data, jsonf, indent=2)
 
     @classmethod
-    def load(cls, path: str, **kwargs: Dict[str, Any]) -> "RBLNModelConfig":
+    def load(cls, path: str, **kwargs: Any) -> "RBLNModelConfig":
         """
         Load a RBLNModelConfig from a path.
 
@@ -742,7 +750,7 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     def initialize_from_kwargs(
         cls: Type["RBLNModelConfig"],
         rbln_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Tuple["RBLNModelConfig", Dict[str, Any]]:
         # Initialize RBLNModelConfig from kwargs.
         kwargs_keys = list(kwargs.keys())
@@ -838,3 +846,14 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     @activate_profiler.setter
     def activate_profiler(self, activate_profiler: bool):
         self._runtime_options["activate_profiler"] = activate_profiler
+
+    @property
+    def timeout(self):
+        context = ContextRblnConfig.get_current_context()["timeout"]
+        if context is not None:
+            return context
+        return self._runtime_options["timeout"]
+
+    @timeout.setter
+    def timeout(self, timeout: int):
+        self._runtime_options["timeout"] = timeout

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/diffusers/configurations/models/__init__.py

@@ -1,8 +1,8 @@
 from .configuration_autoencoder_kl import RBLNAutoencoderKLConfig
 from .configuration_autoencoder_kl_cosmos import RBLNAutoencoderKLCosmosConfig
 from .configuration_controlnet import RBLNControlNetModelConfig
-from .configuration_cosmos_transformer import RBLNCosmosTransformer3DModelConfig
 from .configuration_prior_transformer import RBLNPriorTransformerConfig
+from .configuration_transformer_cosmos import RBLNCosmosTransformer3DModelConfig
 from .configuration_transformer_sd3 import RBLNSD3Transformer2DModelConfig
 from .configuration_unet_2d_condition import RBLNUNet2DConditionModelConfig
 from .configuration_vq_model import RBLNVQModelConfig

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -33,7 +33,7 @@ class RBLNAutoencoderKLConfig(RBLNModelConfig):
         vae_scale_factor: Optional[float] = None,  # TODO: rename to scaling_factor
         in_channels: Optional[int] = None,
         latent_channels: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/diffusers/configurations/models/configuration_autoencoder_kl_cosmos.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 from ....utils.logging import get_logger
@@ -22,6 +22,8 @@ logger = get_logger(__name__)
 
 
 class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
+    """Configuration class for RBLN Cosmos Variational Autoencoder (VAE) models."""
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -33,7 +35,7 @@ class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
         vae_scale_factor_temporal: Optional[int] = None,
         vae_scale_factor_spatial: Optional[int] = None,
         use_slicing: Optional[bool] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Args:
@@ -48,7 +50,7 @@ class RBLNAutoencoderKLCosmosConfig(RBLNModelConfig):
                 Determines how much shorter the latent representations are compared to the original videos.
             vae_scale_factor_spatial (Optional[int]): The scaling factor between pixel space and latent space.
                 Determines how much smaller the latent representations are compared to the original videos.
-            use_slicing (Optional[Bool]): Enable sliced VAE encoding and decoding.
+            use_slicing (Optional[bool]): Enable sliced VAE encoding and decoding.
                 If True, the VAE will split the input tensor in slices to compute encoding or decoding in several steps.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/diffusers/configurations/models/configuration_controlnet.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -29,7 +29,7 @@ class RBLNControlNetModelConfig(RBLNModelConfig):
         unet_sample_size: Optional[Tuple[int, int]] = None,
         vae_sample_size: Optional[Tuple[int, int]] = None,
         text_model_hidden_size: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/diffusers/configurations/models/configuration_prior_transformer.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -32,7 +32,7 @@ class RBLNPriorTransformerConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         embedding_dim: Optional[int] = None,
         num_embeddings: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

optimum_rbln-0.8.1rc0/src/optimum/rbln/diffusers/configurations/models/configuration_cosmos_transformer.py → optimum_rbln-0.8.2/src/optimum/rbln/diffusers/configurations/models/configuration_transformer_cosmos.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Any, Optional
 
 from ....configuration_utils import RBLNModelConfig
 
 
 class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
+    """Configuration class for RBLN Cosmos Transformer models."""
+
     def __init__(
         self,
         batch_size: Optional[int] = None,
@@ -31,7 +33,7 @@ class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
         num_latent_frames: Optional[int] = None,
         latent_height: Optional[int] = None,
         latent_width: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Args:
@@ -50,6 +52,9 @@ class RBLNCosmosTransformer3DModelConfig(RBLNModelConfig):
         Raises:
             ValueError: If batch_size is not a positive integer.
         """
+        if kwargs.get("timeout") is None:
+            kwargs["timeout"] = 80
+
         super().__init__(**kwargs)
         self.batch_size = batch_size or 1
         self.num_frames = num_frames or 121

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/diffusers/configurations/models/configuration_transformer_sd3.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -27,7 +27,7 @@ class RBLNSD3Transformer2DModelConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         sample_size: Optional[Union[int, Tuple[int, int]]] = None,
         prompt_embed_length: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/diffusers/configurations/models/configuration_unet_2d_condition.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -38,7 +38,7 @@ class RBLNUNet2DConditionModelConfig(RBLNModelConfig):
         in_features: Optional[int] = None,
         text_model_hidden_size: Optional[int] = None,
         image_model_hidden_size: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

{optimum_rbln-0.8.1rc0 → optimum_rbln-0.8.2}/src/optimum/rbln/diffusers/configurations/models/configuration_vq_model.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNModelConfig
 
@@ -33,7 +33,7 @@ class RBLNVQModelConfig(RBLNModelConfig):
         vqmodel_scale_factor: Optional[float] = None,  # TODO: rename to scaling_factor
         in_channels: Optional[int] = None,
         latent_channels: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args: