optimum-rbln 0.8.4a8__py3-none-any.whl → 0.9.2__py3-none-any.whl

optimum/rbln/configuration_utils.py

@@ -33,7 +33,6 @@ logger = get_logger(__name__)
 
 
 DEFAULT_COMPILED_MODEL_NAME = "compiled_model"
-DEFAULT_MOD_NAME = "default"
 TypeInputInfo = List[Tuple[str, Tuple[int], str]]
 
 
@@ -41,6 +40,9 @@ TypeInputInfo = List[Tuple[str, Tuple[int], str]]
 class RBLNSerializableConfigProtocol(Protocol):
     def _prepare_for_serialization(self) -> Dict[str, Any]: ...
 
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self._prepare_for_serialization()})"
+
 
 @dataclass
 class RBLNCompileConfig:
@@ -49,17 +51,13 @@ class RBLNCompileConfig:
 
     Attributes:
         compiled_model_name (str): Name of the compiled model.
-        mod_name (str): Name of the RBLN module.
         input_info (Union[List[TypeInputInfo], TypeInputInfo]): Information about input tensors.
-        fusion (Optional[bool]): Whether to use fusion optimization.
         npu (Optional[str]): NPU configuration.
         tensor_parallel_size (Optional[int]): Size for tensor parallelism.
     """
 
     compiled_model_name: str = DEFAULT_COMPILED_MODEL_NAME
-    mod_name: str = DEFAULT_MOD_NAME
     input_info: Union[List[TypeInputInfo], TypeInputInfo] = None
-    fusion: Optional[bool] = None
     npu: Optional[str] = None
     tensor_parallel_size: Optional[int] = None
 
@@ -113,9 +111,7 @@ class RBLNCompileConfig:
 
     def update(self, kwargs: Dict[str, Any]):
         self.compiled_model_name = kwargs.get("compiled_model_name", self.compiled_model_name)
-        self.mod_name = kwargs.get("mod_name", self.mod_name)
         self.input_info = kwargs.get("input_info", self.input_info)
-        self.fusion = kwargs.get("fusion", self.fusion)
         self.npu = kwargs.get("npu", self.npu)
         self.tensor_parallel_size = kwargs.get("tensor_parallel_size", self.tensor_parallel_size)
         return self
@@ -149,7 +145,7 @@ class RBLNCompileConfig:
         return asdict(self)
 
 
-RUNTIME_KEYWORDS = ["create_runtimes", "optimize_host_memory", "device", "device_map", "activate_profiler", "timeout"]
+RUNTIME_KEYWORDS = ["create_runtimes", "device", "device_map", "activate_profiler", "timeout"]
 CONFIG_MAPPING: Dict[str, Type["RBLNModelConfig"]] = {}
 
 
@@ -525,7 +521,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         "npu",
         "tensor_parallel_size",
         "create_runtimes",
-        "optimize_host_memory",
         "device",
         "device_map",
         "activate_profiler",
@@ -534,18 +529,18 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     submodules: List[str] = []
     subclass_non_save_attributes = []
 
-    def init_submodule_config(
+    def initialize_submodule_config(
         self,
-        submodule_config_cls: Type["RBLNModelConfig"],
         submodule_config: Optional[Union[Dict[str, Any], "RBLNModelConfig"]] = None,
+        force_kwargs: bool = False,
         **kwargs: Any,
     ) -> "RBLNModelConfig":
-        # Initialize a submodule config from a dict or a RBLNModelConfig.
-        # kwargs is specified from the predecessor config.
-
         if submodule_config is None:
             submodule_config = {}
 
+        if isinstance(submodule_config, RBLNModelConfig):
+            return submodule_config
+
         if isinstance(submodule_config, dict):
             from_predecessor = self._runtime_options.copy()
             from_predecessor.update(
@@ -559,13 +554,60 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
 
             init_kwargs = from_predecessor
             init_kwargs.update(submodule_config)
-            submodule_config = submodule_config_cls(**init_kwargs)
 
-        if not isinstance(submodule_config, submodule_config_cls):
+            if force_kwargs:
+                for key, value in kwargs.items():
+                    if key in init_kwargs:
+                        if init_kwargs[key] != value:
+                            raise ValueError(
+                                f"Parameter conflict for '{key}': submodule_config has {init_kwargs[key]}, "
+                                f"but kwargs has {value}. Using kwargs value: {value}"
+                            )
+                        init_kwargs[key] = value
+
+            if "cls_name" in init_kwargs:
+                config_cls = get_rbln_config_class(init_kwargs["cls_name"])
+            else:
+                return init_kwargs
+
+            submodule_config = config_cls(**init_kwargs)
+
+        if not isinstance(submodule_config, RBLNModelConfig):
             raise TypeError(f"Invalid submodule config type: {type(submodule_config)}")
 
         return submodule_config
 
+    def filter_parameters(self, config_cls: Type["RBLNModelConfig"], parameters: Dict[str, Any]) -> Dict[str, Any]:
+        import importlib
+
+        model_cls_name = config_cls.__name__.replace("Config", "")
+        modeling_module_name = config_cls.__module__.replace("configuration_", "modeling_")
+
+        model_cls = None
+        try:
+            modeling_module = importlib.import_module(modeling_module_name)
+            if hasattr(modeling_module, model_cls_name):
+                model_cls = getattr(modeling_module, model_cls_name)
+        except ImportError:
+            logger.debug(f"Could not import modeling module: {modeling_module_name}")
+
+        filtered_out_params = set()
+
+        if model_cls is not None:
+            if not getattr(model_cls, "_tp_support", False):
+                filtered_out_params.add("tensor_parallel_size")
+
+        filtered_params = {}
+        for key, value in parameters.items():
+            if key in filtered_out_params:
+                logger.debug(
+                    f"Parameter '{key}' filtered out for {config_cls.__name__} (not supported by model flags)."
+                )
+            else:
+                filtered_params[key] = value
+
+        return filtered_params
+
     def __setattr__(self, key, value):
         if (
             key != "_attributes_map"
@@ -604,7 +646,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         self,
         cls_name: Optional[str] = None,
         create_runtimes: Optional[bool] = None,
-        optimize_host_memory: Optional[bool] = None,
         device: Optional[Union[int, List[int]]] = None,
         device_map: Optional[Dict[str, Union[int, List[int]]]] = None,
         activate_profiler: Optional[bool] = None,
@@ -614,6 +655,8 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         optimum_rbln_version: Optional[str] = None,
         _torch_dtype: Optional[str] = None,
         _compile_cfgs: List[RBLNCompileConfig] = [],
+        *,
+        optimize_host_memory: Optional[bool] = None,
         **kwargs: Any,
     ):
         """
@@ -622,7 +665,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
         Args:
             cls_name (Optional[str]): The class name of the configuration. Defaults to the current class name.
             create_runtimes (Optional[bool]): Whether to create RBLN runtimes. Defaults to True.
-            optimize_host_memory (Optional[bool]): Whether to optimize host memory usage. Defaults to True.
             device (Optional[Union[int, List[int]]]): The device(s) to load the model onto. Can be a single device ID or a list.
             device_map (Optional[Dict[str, Union[int, List[int]]]]): Mapping from compiled model names to device IDs.
             activate_profiler (Optional[bool]): Whether to activate the profiler for performance analysis.
@@ -648,12 +690,14 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
 
         self._runtime_options = {}
         self._runtime_options["create_runtimes"] = create_runtimes
-        self._runtime_options["optimize_host_memory"] = optimize_host_memory
         self._runtime_options["device"] = device
         self._runtime_options["device_map"] = device_map
         self._runtime_options["activate_profiler"] = activate_profiler
         self._runtime_options["timeout"] = timeout
 
+        if optimize_host_memory is not None:
+            logger.warning("`optimize_host_memory` is deprecated and will be removed in future versions.")
+
         # Automatically pass npu, tensor_parallel_size to compile_cfgs
         self.npu = npu
         self.tensor_parallel_size = tensor_parallel_size
@@ -871,19 +915,6 @@ class RBLNModelConfig(RBLNSerializableConfigProtocol):
     def create_runtimes(self, create_runtimes: bool):
         self._runtime_options["create_runtimes"] = create_runtimes
 
-    @property
-    def optimize_host_memory(self):
-        context = ContextRblnConfig.get_current_context()["optimize_host_memory"]
-        if context is not None:
-            return context
-        elif self._runtime_options["optimize_host_memory"] is None:
-            return True
-        return self._runtime_options["optimize_host_memory"]
-
-    @optimize_host_memory.setter
-    def optimize_host_memory(self, optimize_host_memory: bool):
-        self._runtime_options["optimize_host_memory"] = optimize_host_memory
-
     @property
     def device(self):
         context = ContextRblnConfig.get_current_context()["device"]

optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py

@@ -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:
@@ -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

@@ -63,12 +63,15 @@ class RBLNCosmosPipelineBaseConfig(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

@@ -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,
@@ -173,14 +177,20 @@ class RBLNKandinskyV22PriorPipelineConfig(RBLNModelConfig):
             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:
@@ -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

@@ -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

@@ -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

@@ -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

@@ -244,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,
         ):
@@ -412,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/pipelines/cosmos/configuration_cosmos_guardrail.py

@@ -15,7 +15,7 @@
 from typing import Any, Optional, Tuple
 
 from ....configuration_utils import RBLNAutoConfig, RBLNModelConfig
-from ....transformers import RBLNLlamaForCausalLMConfig, RBLNSiglipVisionModelConfig
+from ....transformers import RBLNSiglipVisionModelConfig
 
 
 class RBLNVideoSafetyModelConfig(RBLNModelConfig):
@@ -81,30 +81,28 @@ class RBLNCosmosSafetyCheckerConfig(RBLNModelConfig):
 
         tensor_parallel_size = kwargs.get("tensor_parallel_size")
 
-        self.llamaguard3 = self.init_submodule_config(
-            RBLNLlamaForCausalLMConfig,
+        self.llamaguard3 = self.initialize_submodule_config(
             llamaguard3,
+            cls_name="RBLNLlamaForCausalLMConfig",
             batch_size=batch_size,
             tensor_parallel_size=tensor_parallel_size,
             max_seq_len=max_seq_len,
         )
-
-        self.siglip_encoder = self.init_submodule_config(
-            RBLNSiglipVisionModelConfig,
+        self.siglip_encoder = self.initialize_submodule_config(
             siglip_encoder,
+            cls_name="RBLNSiglipVisionModelConfig",
             batch_size=batch_size,
             image_size=(384, 384),
         )
-
-        self.video_safety_model = self.init_submodule_config(
-            RBLNVideoSafetyModelConfig,
+        self.video_safety_model = self.initialize_submodule_config(
             video_safety_model,
+            cls_name="RBLNVideoSafetyModelConfig",
             batch_size=batch_size,
             input_size=1152,
         )
-        self.face_blur_filter = self.init_submodule_config(
-            RBLNRetinaFaceFilterConfig,
+        self.face_blur_filter = self.initialize_submodule_config(
             face_blur_filter,
+            cls_name="RBLNRetinaFaceFilterConfig",
             batch_size=batch_size,
             image_size=image_size,
         )

optimum/rbln/modeling.py

@@ -34,6 +34,49 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
+def _get_dtype(
+    cls,
+    dtype: Optional[Union[str, torch.dtype, dict]],
+    config: PretrainedConfig,
+) -> tuple[PretrainedConfig, Optional[torch.dtype], Optional[torch.dtype]]:
+    dtype_orig = None
+
+    if dtype is not None:
+        if isinstance(dtype, str):
+            if dtype == "auto":
+                if hasattr(config, "dtype") and config.dtype is not None:
+                    dtype = config.dtype
+                else:
+                    dtype = torch.get_default_dtype()
+            elif hasattr(torch, dtype):
+                dtype = getattr(torch, dtype)
+                config.dtype = dtype
+        elif isinstance(dtype, torch.dtype):
+            config.dtype = dtype
+        elif isinstance(dtype, dict):
+            for key, curr_dtype in dtype.items():
+                if hasattr(config, key):
+                    value = getattr(config, key)
+                    curr_dtype = curr_dtype if not isinstance(curr_dtype, str) else getattr(torch, curr_dtype)
+                    value.dtype = curr_dtype
+            # main torch dtype for modules that aren't part of any sub-config
+            dtype = dtype.get("")
+            dtype = dtype if not isinstance(dtype, str) else getattr(torch, dtype)
+            config.dtype = dtype
+            if dtype is None:
+                dtype = torch.float32
+        else:
+            raise ValueError(f"Invalid dtype: {dtype}")
+
+        dtype_orig = cls._set_default_dtype(dtype)
+    else:
+        # Use default dtype
+        default_dtype = torch.get_default_dtype()
+        config.dtype = default_dtype
+
+    return config, dtype, dtype_orig
+
+
 class RBLNModel(RBLNBaseModel):
     @classmethod
     def update_kwargs(cls, kwargs):
@@ -70,6 +113,10 @@ class RBLNModel(RBLNBaseModel):
         )
         return compiled_model
 
+    @classmethod
+    def _reconstruct_model_if_needed(cls, model: "PreTrainedModel"):
+        return model
+
     @classmethod
     def from_model(
         cls,
@@ -103,6 +150,8 @@ class RBLNModel(RBLNBaseModel):
         Returns:
             (RBLNModel): A RBLN model instance ready for inference on RBLN NPU devices.
         """
+
+        model = cls._reconstruct_model_if_needed(model)
         preprocessors = kwargs.pop("preprocessors", [])
         rbln_config, kwargs = cls.prepare_rbln_config(rbln_config=rbln_config, **kwargs)
 
@@ -209,6 +258,7 @@ class RBLNModel(RBLNBaseModel):
         **kwargs,
     ) -> "PreTrainedModel":
         kwargs = cls.update_kwargs(kwargs)
+
         return cls.get_hf_class().from_pretrained(
             model_id,
             subfolder=subfolder,

optimum/rbln/modeling_base.py

@@ -315,7 +315,7 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
             rbln_config,
             model_save_dir=model_save_dir,
             subfolder=subfolder,
-            rbln_compiled_models=(None if rbln_config.optimize_host_memory else rbln_compiled_models),
+            rbln_compiled_models=rbln_compiled_models,
             rbln_submodules=rbln_submodules,
             **kwargs,
         )
@@ -433,7 +433,6 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         compiled_model = rebel.compile_from_torch(
             model,
             input_info=rbln_compile_config.input_info,
-            fusion=rbln_compile_config.fusion,
             npu=rbln_compile_config.npu,
             tensor_parallel_size=rbln_compile_config.tensor_parallel_size,
             **kwargs,