optimum-rbln 0.8.1rc1__py3-none-any.whl → 0.8.2__py3-none-any.whl

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/__version__.py

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

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/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/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
@@ -35,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:

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/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/diffusers/configurations/models/configuration_transformer_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
 
@@ -33,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:
@@ -52,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/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/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/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:

optimum/rbln/diffusers/configurations/pipelines/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
 from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
@@ -38,7 +38,7 @@ class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -178,7 +178,7 @@ class RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

optimum/rbln/diffusers/configurations/pipelines/configuration_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 ....transformers import RBLNT5EncoderModelConfig
@@ -41,7 +41,7 @@ class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
         num_frames: Optional[int] = None,
         fps: Optional[int] = None,
         max_seq_len: Optional[int] = None,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Args:

optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.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
 from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNCLIPVisionModelWithProjectionConfig
@@ -37,7 +37,7 @@ class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         img_width: Optional[int] = None,
         height: Optional[int] = None,
         width: Optional[int] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -148,7 +148,7 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
         *,
         batch_size: Optional[int] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Initialize a configuration for Kandinsky 2.2 prior pipeline optimized for RBLN NPU.
@@ -226,7 +226,7 @@ class RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
         prior_text_encoder: Optional[RBLNCLIPTextModelWithProjectionConfig] = None,
         unet: Optional[RBLNUNet2DConditionModelConfig] = None,
         movq: Optional[RBLNVQModelConfig] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Initialize a configuration for combined Kandinsky 2.2 pipelines optimized for RBLN NPU.

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.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
 from ....transformers import RBLNCLIPTextModelConfig
@@ -37,7 +37,7 @@ class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.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
 from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNT5EncoderModelConfig
@@ -40,7 +40,7 @@ class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
         height: Optional[int] = None,
         width: Optional[int] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.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
 from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
@@ -38,7 +38,7 @@ class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:

optimum/rbln/diffusers/modeling_diffusers.py

@@ -70,8 +70,6 @@ class RBLNDiffusionMixin:
     _submodules = []
     _optional_submodules = []
     _prefix = {}
-    _rbln_config_class = None
-    _hf_class = None
 
     @staticmethod
     def _maybe_apply_and_fuse_lora(
@@ -114,14 +112,14 @@ class RBLNDiffusionMixin:
     @classmethod
     def get_rbln_config_class(cls) -> Type[RBLNModelConfig]:
         # Lazily loads and caches the corresponding RBLN model config class.
-        if cls._rbln_config_class is None:
+        if "_rbln_config_class" not in cls.__dict__ or cls._rbln_config_class is None:
             rbln_config_class_name = cls.__name__ + "Config"
             cls._rbln_config_class = get_rbln_config_class(rbln_config_class_name)
         return cls._rbln_config_class
 
     @classmethod
     def get_hf_class(cls):
-        if cls._hf_class is None:
+        if "_hf_class" not in cls.__dict__ or cls._hf_class is None:
             hf_cls_name = cls.__name__[4:]
             library = importlib.import_module("diffusers")
             cls._hf_class = getattr(library, hf_cls_name, None)
@@ -138,7 +136,7 @@ class RBLNDiffusionMixin:
         lora_ids: Optional[Union[str, List[str]]] = None,
         lora_weights_names: Optional[Union[str, List[str]]] = None,
         lora_scales: Optional[Union[float, List[float]]] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> "RBLNDiffusionMixin":
         """
         Load a pretrained diffusion pipeline from a model checkpoint, with optional compilation for RBLN NPUs.
@@ -230,6 +228,7 @@ class RBLNDiffusionMixin:
             create_runtimes=rbln_config.create_runtimes,
             optimize_host_mem=rbln_config.optimize_host_memory,
             activate_profiler=rbln_config.activate_profiler,
+            timeout=rbln_config.timeout,
         ):
             model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
 

optimum/rbln/diffusers/models/__init__.py

@@ -35,20 +35,10 @@ _import_structure = {
 }
 
 if TYPE_CHECKING:
-    from .autoencoders import (
-        RBLNAutoencoderKL,
-        RBLNAutoencoderKLCosmos,
-        RBLNVQModel,
-    )
+    from .autoencoders import RBLNAutoencoderKL, RBLNAutoencoderKLCosmos, RBLNVQModel
     from .controlnet import RBLNControlNetModel
-    from .transformers import (
-        RBLNCosmosTransformer3DModel,
-        RBLNPriorTransformer,
-        RBLNSD3Transformer2DModel,
-    )
-    from .unets import (
-        RBLNUNet2DConditionModel,
-    )
+    from .transformers import RBLNCosmosTransformer3DModel, RBLNPriorTransformer, RBLNSD3Transformer2DModel
+    from .unets import RBLNUNet2DConditionModel
 else:
     import sys
 

optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py

@@ -209,6 +209,7 @@ class RBLNAutoencoderKL(RBLNModel):
                 tensor_type="pt",
                 device=device_val,
                 activate_profiler=rbln_config.activate_profiler,
+                timeout=rbln_config.timeout,
             )
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]

optimum/rbln/diffusers/models/autoencoders/autoencoder_kl_cosmos.py

@@ -200,6 +200,7 @@ class RBLNAutoencoderKLCosmos(RBLNModel):
                 tensor_type="pt",
                 device=device_val,
                 activate_profiler=rbln_config.activate_profiler,
+                timeout=rbln_config.timeout,
             )
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]

optimum/rbln/diffusers/models/autoencoders/vq_model.py

@@ -165,6 +165,7 @@ class RBLNVQModel(RBLNModel):
                 tensor_type="pt",
                 device=device_val,
                 activate_profiler=rbln_config.activate_profiler,
+                timeout=rbln_config.timeout,
             )
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]