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

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:
@@ -59,7 +59,7 @@ class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If both image_size and img_height/img_width are provided.
@@ -93,20 +93,27 @@ class RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig,
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelConfig",
+            batch_size=batch_size,
+        )
+        self.unet = self.initialize_submodule_config(
             unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,  # image size is equal to sample size in vae
         )
-        self.controlnet = self.init_submodule_config(RBLNControlNetModelConfig, controlnet)
+        self.controlnet = self.initialize_submodule_config(
+            controlnet,
+            cls_name="RBLNControlNetModelConfig",
+        )
 
         # Get default guidance scale from original class to set UNet and ControlNet batch size
         if guidance_scale is None:
@@ -178,7 +185,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:
@@ -201,7 +208,7 @@ class RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If both image_size and img_height/img_width are provided.
@@ -235,23 +242,32 @@ class RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
-        self.text_encoder_2 = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder_2, batch_size=batch_size
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelConfig",
+            batch_size=batch_size,
         )
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig,
+        self.text_encoder_2 = self.initialize_submodule_config(
+            text_encoder_2,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
+        )
+        self.unet = self.initialize_submodule_config(
             unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,  # image size is equal to sample size in vae
         )
-        self.controlnet = self.init_submodule_config(RBLNControlNetModelConfig, controlnet)
+        self.controlnet = self.initialize_submodule_config(
+            controlnet,
+            cls_name="RBLNControlNetModelConfig",
+        )
 
         # Get default guidance scale from original class to set UNet and ControlNet batch size
         guidance_scale = (

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 Any, Dict, 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: Dict[str, Any],
+        **kwargs: Any,
     ):
         """
         Args:
@@ -59,16 +59,19 @@ class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
             num_frames (Optional[int]): The number of frames in the generated video.
             fps (Optional[int]): The frames per second of the generated video.
             max_seq_len (Optional[int]): Maximum sequence length supported by the model.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
         """
         super().__init__(**kwargs)
 
-        self.text_encoder = self.init_submodule_config(
-            RBLNT5EncoderModelConfig, text_encoder, batch_size=batch_size, max_seq_len=max_seq_len
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNT5EncoderModelConfig",
+            batch_size=batch_size,
+            max_seq_len=max_seq_len,
         )
-        self.transformer = self.init_submodule_config(
-            RBLNCosmosTransformer3DModelConfig,
+        self.transformer = self.initialize_submodule_config(
             transformer,
+            cls_name="RBLNCosmosTransformer3DModelConfig",
             batch_size=batch_size,
             max_seq_len=max_seq_len,
             height=height,
@@ -76,18 +79,18 @@ class RBLNCosmosPipelineBaseConfig(RBLNModelConfig):
             num_frames=num_frames,
             fps=fps,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLCosmosConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLCosmosConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             height=height,
             width=width,
             num_frames=num_frames,
         )
-        self.safety_checker = self.init_submodule_config(
-            RBLNCosmosSafetyCheckerConfig,
+        self.safety_checker = self.initialize_submodule_config(
             safety_checker,
+            cls_name="RBLNCosmosSafetyCheckerConfig",
             batch_size=batch_size,
             height=height,
             width=width,

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:
@@ -54,7 +54,7 @@ class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
             img_width (Optional[int]): Width of the generated images.
             height (Optional[int]): Height of the generated images.
             width (Optional[int]): Width of the generated images.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If both image_size and img_height/img_width are provided.
@@ -88,10 +88,14 @@ class RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.unet = self.init_submodule_config(RBLNUNet2DConditionModelConfig, unet, sample_size=sample_size)
-        self.movq = self.init_submodule_config(
-            RBLNVQModelConfig,
+        self.unet = self.initialize_submodule_config(
+            unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
+            sample_size=sample_size,
+        )
+        self.movq = self.initialize_submodule_config(
             movq,
+            cls_name="RBLNVQModelConfig",
             batch_size=batch_size,
             sample_size=image_size,  # image size is equal to sample size in vae
             uses_encoder=self._movq_uses_encoder,
@@ -148,7 +152,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.
@@ -166,21 +170,27 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
                 Initialized as RBLNPriorTransformerConfig if not provided.
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Note:
             When guidance_scale > 1.0, the prior batch size is automatically doubled to
             accommodate classifier-free guidance.
         """
         super().__init__(**kwargs)
-        self.text_encoder = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder, batch_size=batch_size
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
         )
-        self.image_encoder = self.init_submodule_config(
-            RBLNCLIPVisionModelWithProjectionConfig, image_encoder, batch_size=batch_size
+        self.image_encoder = self.initialize_submodule_config(
+            image_encoder,
+            cls_name="RBLNCLIPVisionModelWithProjectionConfig",
+            batch_size=batch_size,
+        )
+        self.prior = self.initialize_submodule_config(
+            prior,
+            cls_name="RBLNPriorTransformerConfig",
         )
-
-        self.prior = self.init_submodule_config(RBLNPriorTransformerConfig, prior)
 
         # Get default guidance scale from original class to set UNet batch size
         if guidance_scale is None:
@@ -226,7 +236,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.
@@ -259,7 +269,7 @@ class RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
                 Used if decoder_pipe is not provided.
             movq (Optional[RBLNVQModelConfig]): Direct configuration for the MoVQ (VQ-GAN) model.
                 Used if decoder_pipe is not provided.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
         """
         super().__init__(**kwargs)
 
@@ -286,18 +296,18 @@ class RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.prior_pipe = self.init_submodule_config(
-            RBLNKandinskyV22PriorPipelineConfig,
+        self.prior_pipe = self.initialize_submodule_config(
             prior_pipe,
+            cls_name="RBLNKandinskyV22PriorPipelineConfig",
             prior=prior_prior,
             image_encoder=prior_image_encoder,
             text_encoder=prior_text_encoder,
             batch_size=batch_size,
             guidance_scale=guidance_scale,
         )
-        self.decoder_pipe = self.init_submodule_config(
-            self._decoder_pipe_cls,
+        self.decoder_pipe = self.initialize_submodule_config(
             decoder_pipe,
+            cls_name=self._decoder_pipe_cls.__name__,
             unet=unet,
             movq=movq,
             batch_size=batch_size,

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:
@@ -56,7 +56,7 @@ class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If both image_size and img_height/img_width are provided.
@@ -90,18 +90,22 @@ class RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig,
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelConfig",
+            batch_size=batch_size,
+        )
+        self.unet = self.initialize_submodule_config(
             unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
-            sample_size=image_size,  # image size is equal to sample size in vae
+            sample_size=image_size,
         )
 
         # Get default guidance scale from original class to set UNet batch size

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:
@@ -64,7 +64,7 @@ class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
             height (Optional[int]): Height of the generated images.
             width (Optional[int]): Width of the generated images.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If both image_size and img_height/img_width are provided.
@@ -100,27 +100,31 @@ class RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
 
         max_seq_len = max_seq_len or 256
 
-        self.text_encoder = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder, batch_size=batch_size
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
         )
-        self.text_encoder_2 = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder_2, batch_size=batch_size
+        self.text_encoder_2 = self.initialize_submodule_config(
+            text_encoder_2,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
         )
-        self.text_encoder_3 = self.init_submodule_config(
-            RBLNT5EncoderModelConfig,
+        self.text_encoder_3 = self.initialize_submodule_config(
             text_encoder_3,
+            cls_name="RBLNT5EncoderModelConfig",
             batch_size=batch_size,
             max_seq_len=max_seq_len,
             model_input_names=["input_ids"],
         )
-        self.transformer = self.init_submodule_config(
-            RBLNSD3Transformer2DModelConfig,
+        self.transformer = self.initialize_submodule_config(
             transformer,
+            cls_name="RBLNSD3Transformer2DModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,

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:
@@ -59,7 +59,7 @@ class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
-            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+            kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
             ValueError: If both image_size and img_height/img_width are provided.
@@ -93,18 +93,25 @@ class RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
         elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
             raise ValueError("Both img_height and img_width must be provided together if used")
 
-        self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
-        self.text_encoder_2 = self.init_submodule_config(
-            RBLNCLIPTextModelWithProjectionConfig, text_encoder_2, batch_size=batch_size
+        self.text_encoder = self.initialize_submodule_config(
+            text_encoder,
+            cls_name="RBLNCLIPTextModelConfig",
+            batch_size=batch_size,
+        )
+        self.text_encoder_2 = self.initialize_submodule_config(
+            text_encoder_2,
+            cls_name="RBLNCLIPTextModelWithProjectionConfig",
+            batch_size=batch_size,
         )
-        self.unet = self.init_submodule_config(
-            RBLNUNet2DConditionModelConfig,
+
+        self.unet = self.initialize_submodule_config(
             unet,
+            cls_name="RBLNUNet2DConditionModelConfig",
             sample_size=sample_size,
         )
-        self.vae = self.init_submodule_config(
-            RBLNAutoencoderKLConfig,
+        self.vae = self.initialize_submodule_config(
             vae,
+            cls_name="RBLNAutoencoderKLConfig",
             batch_size=batch_size,
             uses_encoder=self.__class__._vae_uses_encoder,
             sample_size=image_size,  # image size is equal to sample size in vae

optimum/rbln/diffusers/modeling_diffusers.py

@@ -33,6 +33,10 @@ if TYPE_CHECKING:
 
 
 class RBLNDiffusionMixinConfig(RBLNModelConfig):
+    """
+    Configuration class for RBLN diffusion pipelines.
+    """
+
     pass
 
 
@@ -54,8 +58,8 @@ class RBLNDiffusionMixin:
         ```
 
     Class Variables:
-        _submodules: List of submodule names that should be compiled (typically ["text_encoder", "unet", "vae"])
-        _optional_submodules: List of submodule names compiled without inheriting RBLNModel (typically ["safety_checker"])
+        - `_submodules`: List of submodule names that should be compiled (typically ["text_encoder", "unet", "vae"])
+        - `_optional_submodules`: List of submodule names compiled without inheriting RBLNModel (typically ["safety_checker"])
 
     Methods:
         from_pretrained: Creates and optionally compiles a model from a pretrained checkpoint
@@ -130,20 +134,20 @@ class RBLNDiffusionMixin:
         cls,
         model_id: str,
         *,
-        export: bool = False,
+        export: bool = None,
         model_save_dir: Optional[PathLike] = None,
         rbln_config: Dict[str, Any] = {},
         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.
 
         This method has two distinct operating modes:
-        - When `export=True`: Takes a PyTorch-based diffusion model, compiles it for RBLN NPUs, and loads the compiled model
-        - When `export=False`: Loads an already compiled RBLN model from `model_id` without recompilation
+            - When `export=True`: Takes a PyTorch-based diffusion model, compiles it for RBLN NPUs, and loads the compiled model
+            - When `export=False`: Loads an already compiled RBLN model from `model_id` without recompilation
 
         It supports various diffusion pipelines including Stable Diffusion, Kandinsky, ControlNet, and other diffusers-based models.
 
@@ -170,7 +174,7 @@ class RBLNDiffusionMixin:
                 Names of specific LoRA weight files to load, corresponding to lora_ids. Only used when `export=True`.
             lora_scales:
                 Scaling factor(s) to apply to the LoRA adapter(s). Only used when `export=True`.
-            **kwargs:
+            kwargs:
                 Additional arguments to pass to the underlying diffusion pipeline constructor or the
                 RBLN compilation process. These may include parameters specific to individual submodules
                 or the particular diffusion pipeline being used.
@@ -181,6 +185,20 @@ class RBLNDiffusionMixin:
         """
         rbln_config, kwargs = cls.get_rbln_config_class().initialize_from_kwargs(rbln_config, **kwargs)
 
+        if export is None:
+            export = any(
+                not RBLNModel._is_compiled(
+                    model_id,
+                    token=kwargs.get("token"),
+                    revision=kwargs.get("revision"),
+                    force_download=kwargs.get("force_download", False),
+                    cache_dir=kwargs.get("cache_dir"),
+                    subfolder=submodule_name,
+                    local_files_only=kwargs.get("local_files_only", False),
+                )
+                for submodule_name in cls._submodules
+            )
+
         if export:
             # keep submodules if user passed any of them.
             passed_submodules = {
@@ -226,7 +244,6 @@ class RBLNDiffusionMixin:
             device=rbln_config.device,
             device_map=rbln_config.device_map,
             create_runtimes=rbln_config.create_runtimes,
-            optimize_host_mem=rbln_config.optimize_host_memory,
             activate_profiler=rbln_config.activate_profiler,
             timeout=rbln_config.timeout,
         ):
@@ -394,12 +411,11 @@ class RBLNDiffusionMixin:
             # overwrite to replace incorrect config
             model.save_config(model_save_dir)
 
-        if rbln_config.optimize_host_memory is False:
-            # Keep compiled_model objs to further analysis. -> TODO: remove soon...
-            model.compiled_models = []
-            for name in cls._submodules:
-                submodule = getattr(model, name)
-                model.compiled_models.extend(submodule.compiled_models)
+        # Keep compiled_model objs to further analysis. -> TODO: remove soon...
+        model.compiled_models = []
+        for name in cls._submodules:
+            submodule = getattr(model, name)
+            model.compiled_models.extend(submodule.compiled_models)
 
         return model
 

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