optimum-rbln 0.7.4a4__py3-none-any.whl → 0.7.4a5__py3-none-any.whl

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

@@ -0,0 +1,143 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+from ....transformers import RBLNCLIPTextModelWithProjectionConfig, RBLNT5EncoderModelConfig
+from ....utils.logging import get_logger
+from ..models import RBLNAutoencoderKLConfig, RBLNSD3Transformer2DModelConfig
+
+
+logger = get_logger(__name__)
+
+
+class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
+    submodules = ["transformer", "text_encoder", "text_encoder_2", "text_encoder_3", "vae"]
+    _vae_uses_encoder = False
+
+    def __init__(
+        self,
+        transformer: Optional[RBLNSD3Transformer2DModelConfig] = None,
+        text_encoder: Optional[RBLNCLIPTextModelWithProjectionConfig] = None,
+        text_encoder_2: Optional[RBLNCLIPTextModelWithProjectionConfig] = None,
+        text_encoder_3: Optional[RBLNT5EncoderModelConfig] = None,
+        vae: Optional[RBLNAutoencoderKLConfig] = None,
+        *,
+        max_seq_len: Optional[int] = None,
+        sample_size: Optional[Tuple[int, int]] = None,
+        image_size: Optional[Tuple[int, int]] = None,
+        batch_size: Optional[int] = None,
+        img_height: Optional[int] = None,
+        img_width: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            transformer (Optional[RBLNSD3Transformer2DModelConfig]): Configuration for the transformer model component.
+                Initialized as RBLNSD3Transformer2DModelConfig if not provided.
+            text_encoder (Optional[RBLNCLIPTextModelWithProjectionConfig]): Configuration for the primary text encoder.
+                Initialized as RBLNCLIPTextModelWithProjectionConfig if not provided.
+            text_encoder_2 (Optional[RBLNCLIPTextModelWithProjectionConfig]): Configuration for the secondary text encoder.
+                Initialized as RBLNCLIPTextModelWithProjectionConfig if not provided.
+            text_encoder_3 (Optional[RBLNT5EncoderModelConfig]): Configuration for the tertiary text encoder.
+                Initialized as RBLNT5EncoderModelConfig if not provided.
+            vae (Optional[RBLNAutoencoderKLConfig]): Configuration for the VAE model component.
+                Initialized as RBLNAutoencoderKLConfig if not provided.
+            max_seq_len (Optional[int]): Maximum sequence length for text inputs. Defaults to 256.
+            sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the transformer model.
+            image_size (Optional[Tuple[int, int]]): Dimensions for the generated images.
+                Cannot be used together with img_height/img_width.
+            batch_size (Optional[int]): Batch size for inference, applied to all submodules.
+            img_height (Optional[int]): Height of the generated images.
+            img_width (Optional[int]): Width of the generated images.
+            guidance_scale (Optional[float]): Scale for classifier-free guidance. Deprecated parameter.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If both image_size and img_height/img_width are provided.
+
+        Note:
+            When guidance_scale > 1.0, the transformer batch size is automatically doubled to
+            accommodate classifier-free guidance.
+        """
+        super().__init__(**kwargs)
+        if image_size is not None and (img_height is not None or img_width is not None):
+            raise ValueError("image_size and img_height/img_width cannot both be provided")
+
+        if img_height is not None and img_width is not None:
+            image_size = (img_height, img_width)
+
+        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_2 = self.init_submodule_config(
+            RBLNCLIPTextModelWithProjectionConfig, text_encoder_2, batch_size=batch_size
+        )
+        self.text_encoder_3 = self.init_submodule_config(
+            RBLNT5EncoderModelConfig,
+            text_encoder_3,
+            batch_size=batch_size,
+            max_seq_len=max_seq_len,
+        )
+        self.transformer = self.init_submodule_config(
+            RBLNSD3Transformer2DModelConfig,
+            transformer,
+            batch_size=batch_size,
+            sample_size=sample_size,
+        )
+        self.vae = self.init_submodule_config(
+            RBLNAutoencoderKLConfig,
+            vae,
+            batch_size=batch_size,
+            uses_encoder=self.__class__._vae_uses_encoder,
+            sample_size=image_size,
+        )
+
+        if guidance_scale is not None:
+            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
+            do_classifier_free_guidance = guidance_scale > 1.0
+            if do_classifier_free_guidance:
+                self.transformer.batch_size = self.text_encoder.batch_size * 2
+
+    @property
+    def max_seq_len(self):
+        return self.text_encoder_3.max_seq_len
+
+    @property
+    def batch_size(self):
+        return self.vae.batch_size
+
+    @property
+    def sample_size(self):
+        return self.transformer.sample_size
+
+    @property
+    def image_size(self):
+        return self.vae.sample_size
+
+
+class RBLNStableDiffusion3PipelineConfig(_RBLNStableDiffusion3PipelineBaseConfig):
+    _vae_uses_encoder = False
+
+
+class RBLNStableDiffusion3Img2ImgPipelineConfig(_RBLNStableDiffusion3PipelineBaseConfig):
+    _vae_uses_encoder = True
+
+
+class RBLNStableDiffusion3InpaintPipelineConfig(_RBLNStableDiffusion3PipelineBaseConfig):
+    _vae_uses_encoder = True

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

@@ -0,0 +1,124 @@
+# Copyright 2025 Rebellions Inc. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple
+
+from ....configuration_utils import RBLNModelConfig
+from ....transformers import RBLNCLIPTextModelConfig, RBLNCLIPTextModelWithProjectionConfig
+from ....utils.logging import get_logger
+from ..models import RBLNAutoencoderKLConfig, RBLNUNet2DConditionModelConfig
+
+
+logger = get_logger(__name__)
+
+
+class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
+    submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]
+    _vae_uses_encoder = False
+
+    def __init__(
+        self,
+        text_encoder: Optional[RBLNCLIPTextModelConfig] = None,
+        text_encoder_2: Optional[RBLNCLIPTextModelWithProjectionConfig] = None,
+        unet: Optional[RBLNUNet2DConditionModelConfig] = None,
+        vae: Optional[RBLNAutoencoderKLConfig] = None,
+        *,
+        batch_size: Optional[int] = None,
+        img_height: Optional[int] = None,
+        img_width: Optional[int] = None,
+        sample_size: Optional[Tuple[int, int]] = None,
+        image_size: Optional[Tuple[int, int]] = None,
+        guidance_scale: Optional[float] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            text_encoder (Optional[RBLNCLIPTextModelConfig]): Configuration for the primary text encoder component.
+                Initialized as RBLNCLIPTextModelConfig if not provided.
+            text_encoder_2 (Optional[RBLNCLIPTextModelWithProjectionConfig]): Configuration for the secondary text encoder component.
+                Initialized as RBLNCLIPTextModelWithProjectionConfig if not provided.
+            unet (Optional[RBLNUNet2DConditionModelConfig]): Configuration for the UNet model component.
+                Initialized as RBLNUNet2DConditionModelConfig if not provided.
+            vae (Optional[RBLNAutoencoderKLConfig]): Configuration for the VAE model component.
+                Initialized as RBLNAutoencoderKLConfig if not provided.
+            batch_size (Optional[int]): Batch size for inference, applied to all submodules.
+            img_height (Optional[int]): Height of the generated images.
+            img_width (Optional[int]): Width of the generated images.
+            sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
+            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. Deprecated parameter.
+            **kwargs: Additional arguments passed to the parent RBLNModelConfig.
+
+        Raises:
+            ValueError: If both image_size and img_height/img_width are provided.
+
+        Note:
+            When guidance_scale > 1.0, the UNet batch size is automatically doubled to
+            accommodate classifier-free guidance.
+        """
+        super().__init__(**kwargs)
+        if image_size is not None and (img_height is not None or img_width is not None):
+            raise ValueError("image_size and img_height/img_width cannot both be provided")
+
+        if img_height is not None and img_width is not None:
+            image_size = (img_height, img_width)
+
+        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.unet = self.init_submodule_config(
+            RBLNUNet2DConditionModelConfig,
+            unet,
+            batch_size=batch_size,
+            sample_size=sample_size,
+        )
+        self.vae = self.init_submodule_config(
+            RBLNAutoencoderKLConfig,
+            vae,
+            batch_size=batch_size,
+            uses_encoder=self.__class__._vae_uses_encoder,
+            sample_size=image_size,  # image size is equal to sample size in vae
+        )
+
+        if guidance_scale is not None:
+            logger.warning("Specifying `guidance_scale` is deprecated. It will be removed in a future version.")
+            do_classifier_free_guidance = guidance_scale > 1.0
+            if do_classifier_free_guidance:
+                self.unet.batch_size = self.text_encoder.batch_size * 2
+
+    @property
+    def batch_size(self):
+        return self.vae.batch_size
+
+    @property
+    def sample_size(self):
+        return self.unet.sample_size
+
+    @property
+    def image_size(self):
+        return self.vae.sample_size
+
+
+class RBLNStableDiffusionXLPipelineConfig(_RBLNStableDiffusionXLPipelineBaseConfig):
+    _vae_uses_encoder = False
+
+
+class RBLNStableDiffusionXLImg2ImgPipelineConfig(_RBLNStableDiffusionXLPipelineBaseConfig):
+    _vae_uses_encoder = True
+
+
+class RBLNStableDiffusionXLInpaintPipelineConfig(_RBLNStableDiffusionXLPipelineBaseConfig):
+    _vae_uses_encoder = True

optimum/rbln/diffusers/modeling_diffusers.py

@@ -15,12 +15,14 @@
 import copy
 import importlib
 from os import PathLike
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 import torch
 
+from ..configuration_utils import ContextRblnConfig, RBLNModelConfig
 from ..modeling import RBLNModel
-from ..modeling_config import RUNTIME_KEYWORDS, ContextRblnConfig, use_rbln_config
+
+# from ..transformers import RBLNCLIPTextModelConfig
 from ..utils.decorator_utils import remove_compile_time_kwargs
 from ..utils.logging import get_logger
 
@@ -31,6 +33,10 @@ if TYPE_CHECKING:
     from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
 
 
+class RBLNDiffusionMixinConfig(RBLNModelConfig):
+    pass
+
+
 class RBLNDiffusionMixin:
     """
     RBLNDiffusionMixin provides essential functionalities for compiling Stable Diffusion pipeline components to run on RBLN NPUs.
@@ -69,6 +75,7 @@ class RBLNDiffusionMixin:
     _connected_classes = {}
     _submodules = []
     _prefix = {}
+    _rbln_config_class = None
 
     @classmethod
     def is_img2img_pipeline(cls):
@@ -78,35 +85,6 @@ class RBLNDiffusionMixin:
     def is_inpaint_pipeline(cls):
         return "Inpaint" in cls.__name__
 
-    @classmethod
-    def get_submodule_rbln_config(
-        cls, model: torch.nn.Module, submodule_name: str, rbln_config: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        submodule = getattr(model, submodule_name)
-        submodule_class_name = submodule.__class__.__name__
-        if isinstance(submodule, torch.nn.Module):
-            if submodule_class_name == "MultiControlNetModel":
-                submodule_class_name = "ControlNetModel"
-
-            submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), f"RBLN{submodule_class_name}")
-
-            submodule_config = rbln_config.get(submodule_name, {})
-            submodule_config = copy.deepcopy(submodule_config)
-
-            pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._submodules}
-
-            submodule_config.update({k: v for k, v in pipe_global_config.items() if k not in submodule_config})
-            submodule_config.update(
-                {
-                    "img2img_pipeline": cls.is_img2img_pipeline(),
-                    "inpaint_pipeline": cls.is_inpaint_pipeline(),
-                }
-            )
-            submodule_config = submodule_cls.update_rbln_config_using_pipe(model, submodule_config)
-        else:
-            raise ValueError(f"submodule {submodule_name} isn't supported")
-        return submodule_config
-
     @staticmethod
     def _maybe_apply_and_fuse_lora(
         model: torch.nn.Module,
@@ -146,7 +124,22 @@ class RBLNDiffusionMixin:
         return model
 
     @classmethod
-    @use_rbln_config
+    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:
+            rbln_config_class_name = cls.__name__ + "Config"
+            library = importlib.import_module("optimum.rbln")
+            cls._rbln_config_class = getattr(library, rbln_config_class_name, None)
+            if cls._rbln_config_class is None:
+                raise ValueError(
+                    f"RBLN config class {rbln_config_class_name} not found. This is an internal error. "
+                    "Please report it to the developers."
+                )
+        return cls._rbln_config_class
+
+    @classmethod
     def from_pretrained(
         cls,
         model_id: str,
@@ -159,6 +152,8 @@ class RBLNDiffusionMixin:
         lora_scales: Optional[Union[float, List[float]]] = None,
         **kwargs,
     ) -> RBLNModel:
+        rbln_config, kwargs = cls.get_rbln_config_class().initialize_from_kwargs(rbln_config, **kwargs)
+
         if export:
             # keep submodules if user passed any of them.
             passed_submodules = {
@@ -168,22 +163,12 @@ class RBLNDiffusionMixin:
         else:
             # raise error if any of submodules are torch module.
             model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
-            rbln_config = cls._flatten_rbln_config(rbln_config)
             for submodule_name in cls._submodules:
                 if isinstance(kwargs.get(submodule_name), torch.nn.Module):
                     raise AssertionError(
                         f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
                     )
 
-                submodule_config = rbln_config.get(submodule_name, {})
-
-                for key, value in rbln_config.items():
-                    if key in RUNTIME_KEYWORDS and key not in submodule_config:
-                        submodule_config[key] = value
-
-                if not any(kwd in submodule_config for kwd in RUNTIME_KEYWORDS):
-                    continue
-
                 module_name, class_name = model_index_config[submodule_name]
                 if module_name != "optimum.rbln":
                     raise ValueError(
@@ -192,19 +177,19 @@ class RBLNDiffusionMixin:
                         "Expected 'optimum.rbln'. Please check the model_index.json configuration."
                     )
 
-                submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), class_name)
-
+                submodule_cls: Type[RBLNModel] = getattr(importlib.import_module("optimum.rbln"), class_name)
+                submodule_config = getattr(rbln_config, submodule_name)
                 submodule = submodule_cls.from_pretrained(
                     model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
                 )
                 kwargs[submodule_name] = submodule
 
         with ContextRblnConfig(
-            device=rbln_config.get("device"),
-            device_map=rbln_config.get("device_map"),
-            create_runtimes=rbln_config.get("create_runtimes"),
-            optimize_host_mem=rbln_config.get("optimize_host_memory"),
-            activate_profiler=rbln_config.get("activate_profiler"),
+            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,
         ):
             model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
 
@@ -224,78 +209,27 @@ class RBLNDiffusionMixin:
             compiled_submodules = cls._compile_submodules(model, passed_submodules, model_save_dir, rbln_config)
         return cls._construct_pipe(model, compiled_submodules, model_save_dir, rbln_config)
 
-    @classmethod
-    def _prepare_rbln_config(
-        cls,
-        rbln_config,
-    ) -> Dict[str, Any]:
-        prepared_config = {}
-        for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
-            connected_pipe_config = rbln_config.pop(connected_pipe_name, {})
-            prefix = cls._prefix.get(connected_pipe_name, "")
-            guidance_scale = rbln_config.pop(f"{prefix}guidance_scale", None)
-            if "guidance_scale" not in connected_pipe_config and guidance_scale is not None:
-                connected_pipe_config["guidance_scale"] = guidance_scale
-            for submodule_name in connected_pipe_cls._submodules:
-                submodule_config = rbln_config.pop(prefix + submodule_name, {})
-                if submodule_name not in connected_pipe_config:
-                    connected_pipe_config[submodule_name] = {}
-                connected_pipe_config[submodule_name].update(
-                    {k: v for k, v in submodule_config.items() if k not in connected_pipe_config[submodule_name]}
-                )
-            prepared_config[connected_pipe_name] = connected_pipe_config
-        prepared_config.update(rbln_config)
-        return prepared_config
-
-    @classmethod
-    def _flatten_rbln_config(
-        cls,
-        rbln_config,
-    ) -> Dict[str, Any]:
-        prepared_config = cls._prepare_rbln_config(rbln_config)
-        flattened_config = {}
-        pipe_global_config = {k: v for k, v in prepared_config.items() if k not in cls._connected_classes.keys()}
-        for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
-            connected_pipe_config = prepared_config.pop(connected_pipe_name)
-            prefix = cls._prefix.get(connected_pipe_name, "")
-            connected_pipe_global_config = {
-                k: v for k, v in connected_pipe_config.items() if k not in connected_pipe_cls._submodules
-            }
-            for submodule_name in connected_pipe_cls._submodules:
-                flattened_config[prefix + submodule_name] = connected_pipe_config[submodule_name]
-                flattened_config[prefix + submodule_name].update(
-                    {
-                        k: v
-                        for k, v in connected_pipe_global_config.items()
-                        if k not in flattened_config[prefix + submodule_name]
-                    }
-                )
-        flattened_config.update(pipe_global_config)
-        return flattened_config
-
     @classmethod
     def _compile_pipelines(
         cls,
         model: torch.nn.Module,
         passed_submodules: Dict[str, RBLNModel],
         model_save_dir: Optional[PathLike],
-        rbln_config: Dict[str, Any],
+        rbln_config: "RBLNDiffusionMixinConfig",
     ) -> Dict[str, RBLNModel]:
         compiled_submodules = {}
-
-        rbln_config = cls._prepare_rbln_config(rbln_config)
-        pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._connected_classes.keys()}
         for connected_pipe_name, connected_pipe_cls in cls._connected_classes.items():
             connected_pipe_submodules = {}
             prefix = cls._prefix.get(connected_pipe_name, "")
             for submodule_name in connected_pipe_cls._submodules:
                 connected_pipe_submodules[submodule_name] = passed_submodules.get(prefix + submodule_name, None)
             connected_pipe = getattr(model, connected_pipe_name)
-            connected_pipe_config = {}
-            connected_pipe_config.update(pipe_global_config)
-            connected_pipe_config.update(rbln_config[connected_pipe_name])
             connected_pipe_compiled_submodules = connected_pipe_cls._compile_submodules(
-                connected_pipe, connected_pipe_submodules, model_save_dir, connected_pipe_config, prefix
+                connected_pipe,
+                connected_pipe_submodules,
+                model_save_dir,
+                getattr(rbln_config, connected_pipe_name),
+                prefix,
             )
             for submodule_name, compiled_submodule in connected_pipe_compiled_submodules.items():
                 compiled_submodules[prefix + submodule_name] = compiled_submodule
@@ -307,14 +241,19 @@ class RBLNDiffusionMixin:
         model: torch.nn.Module,
         passed_submodules: Dict[str, RBLNModel],
         model_save_dir: Optional[PathLike],
-        rbln_config: Dict[str, Any],
+        rbln_config: RBLNDiffusionMixinConfig,
         prefix: Optional[str] = "",
     ) -> Dict[str, RBLNModel]:
         compiled_submodules = {}
 
         for submodule_name in cls._submodules:
             submodule = passed_submodules.get(submodule_name) or getattr(model, submodule_name, None)
-            submodule_rbln_config = cls.get_submodule_rbln_config(model, submodule_name, rbln_config)
+
+            if getattr(rbln_config, submodule_name, None) is None:
+                raise ValueError(f"RBLN config for submodule {submodule_name} is not provided.")
+
+            submodule_rbln_cls: Type[RBLNModel] = getattr(rbln_config, submodule_name).rbln_model_cls
+            rbln_config = submodule_rbln_cls.update_rbln_config_using_pipe(model, rbln_config, submodule_name)
 
             if submodule is None:
                 raise ValueError(f"submodule ({submodule_name}) cannot be accessed since it is not provided.")
@@ -325,7 +264,7 @@ class RBLNDiffusionMixin:
                 submodule = cls._compile_multicontrolnet(
                     controlnets=submodule,
                     model_save_dir=model_save_dir,
-                    controlnet_rbln_config=submodule_rbln_config,
+                    controlnet_rbln_config=getattr(rbln_config, submodule_name),
                     prefix=prefix,
                 )
             elif isinstance(submodule, torch.nn.Module):
@@ -337,7 +276,7 @@ class RBLNDiffusionMixin:
                     model=submodule,
                     subfolder=subfolder,
                     model_save_dir=model_save_dir,
-                    rbln_config=submodule_rbln_config,
+                    rbln_config=getattr(rbln_config, submodule_name),
                 )
             else:
                 raise ValueError(f"Unknown class of submodule({submodule_name}) : {submodule.__class__.__name__} ")
@@ -350,22 +289,24 @@ class RBLNDiffusionMixin:
         cls,
         controlnets: "MultiControlNetModel",
         model_save_dir: Optional[PathLike],
-        controlnet_rbln_config: Dict[str, Any],
+        controlnet_rbln_config: RBLNModelConfig,
         prefix: Optional[str] = "",
     ):
         # Compile multiple ControlNet models for a MultiControlNet setup
         from .models.controlnet import RBLNControlNetModel
         from .pipelines.controlnet import RBLNMultiControlNetModel
 
-        compiled_controlnets = [
-            RBLNControlNetModel.from_model(
-                model=controlnet,
-                subfolder=f"{prefix}controlnet" if i == 0 else f"{prefix}controlnet_{i}",
-                model_save_dir=model_save_dir,
-                rbln_config=controlnet_rbln_config,
+        compiled_controlnets = []
+        for i, controlnet in enumerate(controlnets.nets):
+            _controlnet_rbln_config = copy.deepcopy(controlnet_rbln_config)
+            compiled_controlnets.append(
+                RBLNControlNetModel.from_model(
+                    model=controlnet,
+                    subfolder=f"{prefix}controlnet" if i == 0 else f"{prefix}controlnet_{i}",
+                    model_save_dir=model_save_dir,
+                    rbln_config=_controlnet_rbln_config,
+                )
             )
-            for i, controlnet in enumerate(controlnets.nets)
-        ]
         return RBLNMultiControlNetModel(compiled_controlnets)
 
     @classmethod
@@ -412,7 +353,7 @@ class RBLNDiffusionMixin:
             # overwrite to replace incorrect config
             model.save_config(model_save_dir)
 
-        if rbln_config.get("optimize_host_memory") is False:
+        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:
@@ -441,9 +382,9 @@ class RBLNDiffusionMixin:
                 kwargs["height"] = compiled_image_size[0]
                 kwargs["width"] = compiled_image_size[1]
 
-            compiled_num_frames = self.unet.rbln_config.model_cfg.get("num_frames", None)
+            compiled_num_frames = self.unet.rbln_config.num_frames
             if compiled_num_frames is not None:
-                kwargs["num_frames"] = self.unet.rbln_config.model_cfg.get("num_frames")
+                kwargs["num_frames"] = compiled_num_frames
             return kwargs
             ```
         """