optimum-rbln 0.1.12__py3-none-any.whl → 0.1.15__py3-none-any.whl

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet.py

@@ -26,203 +26,25 @@ from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
 import torch.nn.functional as F
-from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import StableDiffusionControlNetPipeline
 from diffusers.image_processor import PipelineImageInput
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.pipelines.controlnet.pipeline_controlnet import retrieve_timesteps
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
-from transformers import CLIPTextModel
 
-from ....modeling_config import use_rbln_config
-from ....transformers import RBLNCLIPTextModel
-from ....utils.runtime_utils import ContextRblnConfig
-from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
+from ....modeling_diffusers import RBLNDiffusionMixin
+from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
 
 logger = logging.get_logger(__name__)
 
 
-class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
-    @classmethod
-    @use_rbln_config
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for text-to-image generation using Stable Diffusion with ControlNet.
-
-        This model inherits from [`StableDiffusionControlNetPipeline`]. Check the superclass documentation for the generic methods
-        implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-        It implements the methods to convert a pre-trained Stable Diffusion Controlnet pipeline into a RBLNStableDiffusionControlNet pipeline by:
-        - transferring the checkpoint weights of the original into an optimized RBLN graph,
-        - compiling the resulting graph using the RBLN compiler.
-
-        Args:
-            model_id (`Union[str, Path]`):
-                Can be either:
-                - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
-                - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
-        """
-        export = kwargs.pop("export", None)
-        vae = kwargs.pop("vae", None)
-        unet = kwargs.pop("unet", None)
-        text_encoder = kwargs.pop("text_encoder", None)
-        controlnet = kwargs.pop("controlnet", None)
-        model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config = kwargs.pop("rbln_config", None)
-        rbln_config = {} if rbln_config is None else rbln_config
-
-        device = rbln_config.get("device", None)
-        device_map = rbln_config.get("device_map", None)
-        create_runtimes = rbln_config.get("create_runtimes", None)
-        optimize_host_memory = rbln_config.get("optimize_host_memory", None)
-
-        kwargs_dict = {
-            "pretrained_model_name_or_path": model_id,
-            **kwargs,
-        }
-
-        kwargs_dict.update(
-            {
-                **({"vae": vae} if vae is not None and isinstance(vae, AutoencoderKL) else {}),
-                **({"unet": unet} if unet is not None and isinstance(unet, UNet2DConditionModel) else {}),
-                **(
-                    {"text_encoder": text_encoder}
-                    if text_encoder is not None and isinstance(text_encoder, CLIPTextModel)
-                    else {}
-                ),
-                **(
-                    {"controlnet": controlnet}
-                    if controlnet is not None
-                    and (
-                        isinstance(controlnet, ControlNetModel)
-                        or all(isinstance(c, ControlNetModel) for c in controlnet)
-                    )
-                    else {}
-                ),
-            }
-        )
-
-        with ContextRblnConfig(
-            device=device,
-            device_map=device_map,
-            create_runtimes=create_runtimes,
-            optimze_host_mem=optimize_host_memory,
-        ):
-            model = super().from_pretrained(**{k: v for k, v in kwargs_dict.items() if v is not None})
-
-        if export is None or export is False:
-            return model
-
-        do_classifier_free_guidance = (
-            rbln_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        )
-
-        # compile model, create runtime
-        if not isinstance(vae, RBLNAutoencoderKL):
-            vae = RBLNAutoencoderKL.from_pretrained(
-                model_id=model_id,
-                subfolder="vae",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_unet_sample_size=model.unet.config.sample_size,
-                rbln_use_encode=False,
-                rbln_vae_scale_factor=model.vae_scale_factor,
-                rbln_config={**rbln_config},
-            )
-
-        if not isinstance(text_encoder, RBLNCLIPTextModel):
-            text_encoder = RBLNCLIPTextModel.from_pretrained(
-                model_id=model_id,
-                subfolder="text_encoder",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_config={**rbln_config},
-            )
-
-        batch_size = rbln_config.pop("batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        if not isinstance(unet, RBLNUNet2DConditionModel):
-            unet = RBLNUNet2DConditionModel.from_pretrained(
-                model_id=model_id,
-                subfolder="unet",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_max_seq_len=text_encoder.config.max_position_embeddings,
-                rbln_batch_size=unet_batch_size,
-                rbln_use_encode=False,
-                rbln_vae_scale_factor=model.vae_scale_factor,
-                rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
-                rbln_config={**rbln_config},
-            )
-
-        if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
-            if isinstance(controlnet, (list, tuple)):
-                multicontrolnet = []
-                for i, cid in enumerate(controlnet):
-                    subfolder_name = "controlnet" if i == 0 else f"controlnet_{i}"
-                    multicontrolnet.append(
-                        RBLNControlNetModel.from_model(
-                            model=cid,
-                            subfolder=subfolder_name,
-                            model_save_dir=model_save_dir,
-                            rbln_batch_size=unet_batch_size,
-                            rbln_vae_scale_factor=model.vae_scale_factor,
-                            rbln_config={**rbln_config},
-                        )
-                    )
-                controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
-                controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
-            else:
-                controlnet = RBLNControlNetModel.from_model(
-                    model=controlnet,
-                    subfolder="controlnet",
-                    model_save_dir=model_save_dir,
-                    rbln_batch_size=unet_batch_size,
-                    rbln_vae_scale_factor=model.vae_scale_factor,
-                    rbln_config={**rbln_config},
-                )
-                controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
-
-        if model_save_dir is not None:
-            # To skip saving original pytorch modules
-            del (model.vae, model.text_encoder, model.unet, model.controlnet)
-
-            # Direct calling of `save_pretrained` causes config.unet = (None, None).
-            # So config must be saved again, later.
-            model.save_pretrained(model_save_dir)
-
-        # replace modules
-        model.vae = vae
-        model.text_encoder = text_encoder
-        model.unet = unet
-        model.controlnet = controlnet
-
-        # update config to be able to load from file.
-        update_dict = {
-            "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
-            "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-            "controlnet": controlnet_dict,
-        }
-        model.register_to_config(**update_dict)
-
-        if model_save_dir is not None:
-            # overwrite to replace incorrect config
-            model.save_config(model_save_dir)
-
-        if optimize_host_memory is False:
-            model.compiled_models = [vae.compiled_models[0], text_encoder.compiled_models[0], unet.compiled_models[0]]
-            if isinstance(controlnet, RBLNMultiControlNetModel):
-                for c_model in controlnet.nets:
-                    model.compiled_models.append(c_model.compiled_models[0])
-            else:
-                model.compiled_models.append(controlnet.compiled_models[0])
-
-        return model
+class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionControlNetPipeline):
+    original_class = StableDiffusionControlNetPipeline
+    _submodules = ["text_encoder", "unet", "vae", "controlnet"]
 
     def check_inputs(
         self,
@@ -388,6 +210,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 )
 
     @torch.no_grad()
+    @remove_compile_time_kwargs
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
@@ -597,6 +420,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
         text_encoder_lora_scale = (
             self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
         )
+
         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
             prompt,
             device,

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py

@@ -26,207 +26,24 @@ from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
 import torch.nn.functional as F
-from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionControlNetImg2ImgPipeline
+from diffusers import StableDiffusionControlNetImg2ImgPipeline
 from diffusers.image_processor import PipelineImageInput
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
-from transformers import CLIPTextModel
 
-from ....modeling_config import use_rbln_config
-from ....transformers import RBLNCLIPTextModel
-from ....utils.runtime_utils import ContextRblnConfig
-from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
+from ....modeling_diffusers import RBLNDiffusionMixin
+from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
 
 logger = logging.get_logger(__name__)
 
 
-class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2ImgPipeline):
-    @classmethod
-    @use_rbln_config
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for image-to-image generation using Stable Diffusion with ControlNet.
-
-        This model inherits from [`StableDiffusionControlNetImg2ImgPipeline`]. Check the superclass documentation for the generic methods
-        implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-        It implements the methods to convert a pre-trained Stable Diffusion Controlnet pipeline into a RBLNStableDiffusionControlNetImg2Img pipeline by:
-        - transferring the checkpoint weights of the original into an optimized RBLN graph,
-        - compiling the resulting graph using the RBLN compiler.
-
-        Args:
-            model_id (`Union[str, Path]`):
-                Can be either:
-                - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
-                - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
-        """
-        export = kwargs.pop("export", None)
-        vae = kwargs.pop("vae", None)
-        unet = kwargs.pop("unet", None)
-        text_encoder = kwargs.pop("text_encoder", None)
-        controlnet = kwargs.pop("controlnet", None)
-        model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config = kwargs.pop("rbln_config", None)
-        rbln_config = {} if rbln_config is None else rbln_config
-
-        device = rbln_config.get("device", None)
-        device_map = rbln_config.get("device_map", None)
-        create_runtimes = rbln_config.get("create_runtimes", None)
-        optimize_host_memory = rbln_config.get("optimize_host_memory", None)
-
-        kwargs_dict = {
-            "pretrained_model_name_or_path": model_id,
-            **kwargs,
-        }
-
-        kwargs_dict.update(
-            {
-                **({"vae": vae} if vae is not None and isinstance(vae, AutoencoderKL) else {}),
-                **({"unet": unet} if unet is not None and isinstance(unet, UNet2DConditionModel) else {}),
-                **(
-                    {"text_encoder": text_encoder}
-                    if text_encoder is not None and isinstance(text_encoder, CLIPTextModel)
-                    else {}
-                ),
-                **(
-                    {"controlnet": controlnet}
-                    if controlnet is not None
-                    and (
-                        isinstance(controlnet, ControlNetModel)
-                        or all(isinstance(c, ControlNetModel) for c in controlnet)
-                    )
-                    else {}
-                ),
-            }
-        )
-
-        with ContextRblnConfig(
-            device=device,
-            device_map=device_map,
-            create_runtimes=create_runtimes,
-            optimze_host_mem=optimize_host_memory,
-        ):
-            model = super().from_pretrained(**{k: v for k, v in kwargs_dict.items() if v is not None})
-
-        if export is None or export is False:
-            return model
-
-        do_classifier_free_guidance = (
-            rbln_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        )
-
-        # compile model, create runtime
-        if not isinstance(vae, RBLNAutoencoderKL):
-            vae = RBLNAutoencoderKL.from_pretrained(
-                model_id=model_id,
-                subfolder="vae",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_unet_sample_size=model.unet.config.sample_size,
-                rbln_use_encode=True,
-                rbln_vae_scale_factor=model.vae_scale_factor,
-                rbln_config={**rbln_config},
-            )
-
-        if not isinstance(text_encoder, RBLNCLIPTextModel):
-            text_encoder = RBLNCLIPTextModel.from_pretrained(
-                model_id=model_id,
-                subfolder="text_encoder",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_config={**rbln_config},
-            )
-
-        batch_size = rbln_config.pop("batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        if not isinstance(unet, RBLNUNet2DConditionModel):
-            unet = RBLNUNet2DConditionModel.from_pretrained(
-                model_id=model_id,
-                subfolder="unet",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_max_seq_len=text_encoder.config.max_position_embeddings,
-                rbln_batch_size=unet_batch_size,
-                rbln_use_encode=True,
-                rbln_vae_scale_factor=model.vae_scale_factor,
-                rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
-                rbln_config={**rbln_config},
-            )
-
-        if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
-            if isinstance(controlnet, (list, tuple)):
-                multicontrolnet = []
-                for i, cid in enumerate(controlnet):
-                    subfolder_name = "controlnet" if i == 0 else f"controlnet_{i}"
-                    multicontrolnet.append(
-                        RBLNControlNetModel.from_model(
-                            model=cid,
-                            subfolder=subfolder_name,
-                            model_save_dir=model_save_dir,
-                            rbln_batch_size=unet_batch_size,
-                            rbln_vae_scale_factor=model.vae_scale_factor,
-                            rbln_config={**rbln_config},
-                        )
-                    )
-                controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
-                controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
-            else:
-                controlnet = RBLNControlNetModel.from_model(
-                    model=controlnet,
-                    subfolder="controlnet",
-                    model_save_dir=model_save_dir,
-                    rbln_batch_size=unet_batch_size,
-                    rbln_vae_scale_factor=model.vae_scale_factor,
-                    rbln_config={**rbln_config},
-                )
-                controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
-
-        if model_save_dir is not None:
-            # To skip saving original pytorch modules
-            del (model.vae, model.text_encoder, model.unet, model.controlnet)
-
-            # Direct calling of `save_pretrained` causes config.unet = (None, None).
-            # So config must be saved again, later.
-            model.save_pretrained(model_save_dir)
-
-        # replace modules
-        model.vae = vae
-        model.text_encoder = text_encoder
-        model.unet = unet
-        model.controlnet = controlnet
-
-        # update config to be able to load from file.
-        update_dict = {
-            "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
-            "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-            "controlnet": controlnet_dict,
-        }
-        model.register_to_config(**update_dict)
-
-        if model_save_dir is not None:
-            # overwrite to replace incorrect config
-            model.save_config(model_save_dir)
-
-        if optimize_host_memory is False:
-            model.compiled_models = [
-                vae.compiled_models[0],
-                vae.compiled_models[1],
-                text_encoder.compiled_models[0],
-                unet.compiled_models[0],
-            ]
-            if isinstance(controlnet, RBLNMultiControlNetModel):
-                for c_model in controlnet.nets:
-                    model.compiled_models.append(c_model.compiled_models[0])
-            else:
-                model.compiled_models.append(controlnet.compiled_models[0])
-
-        return model
+class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionControlNetImg2ImgPipeline):
+    original_class = StableDiffusionControlNetImg2ImgPipeline
+    _submodules = ["text_encoder", "unet", "vae", "controlnet"]
 
     def check_inputs(
         self,
@@ -386,6 +203,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                 )
 
     @torch.no_grad()
+    @remove_compile_time_kwargs
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
@@ -594,6 +412,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
         text_encoder_lora_scale = (
             self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
         )
+
         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
             prompt,
             device,

optimum/rbln/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py

@@ -26,208 +26,24 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
-from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionXLControlNetPipeline
+from diffusers import StableDiffusionXLControlNetPipeline
 from diffusers.image_processor import PipelineImageInput
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
 from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
-from transformers import CLIPTextModel
 
-from ....modeling_config import use_rbln_config
-from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
-from ....utils.runtime_utils import ContextRblnConfig
-from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
+from ....modeling_diffusers import RBLNDiffusionMixin
+from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
 
 logger = logging.get_logger(__name__)
 
 
-class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipeline):
-    @classmethod
-    @use_rbln_config
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for text-to-image generation using Stable Diffusion XL with ControlNet.
-
-        This model inherits from [`StableDiffusionXLControlNetPipeline`]. Check the superclass documentation for the generic methods
-        implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-        It implements the methods to convert a pre-trained Stable Diffusion XL Controlnet pipeline into a RBLNStableDiffusionXLControlNet pipeline by:
-        - transferring the checkpoint weights of the original into an optimized RBLN graph,
-        - compiling the resulting graph using the RBLN compiler.
-
-        Args:
-            model_id (`Union[str, Path]`):
-                Can be either:
-                - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
-                - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
-        """
-        export = kwargs.pop("export", None)
-        vae = kwargs.pop("vae", None)
-        unet = kwargs.pop("unet", None)
-        text_encoder = kwargs.pop("text_encoder", None)
-        text_encoder_2 = kwargs.pop("text_encoder_2", None)
-        controlnet = kwargs.pop("controlnet", None)
-        model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config = kwargs.pop("rbln_config", None)
-        rbln_config = {} if rbln_config is None else rbln_config
-
-        device = rbln_config.get("device", None)
-        device_map = rbln_config.get("device_map", None)
-        create_runtimes = rbln_config.get("create_runtimes", None)
-        optimize_host_memory = rbln_config.get("optimize_host_memory", None)
-
-        kwargs_dict = {
-            "pretrained_model_name_or_path": model_id,
-            **kwargs,
-        }
-
-        kwargs_dict.update(
-            {
-                **({"vae": vae} if vae is not None and isinstance(vae, AutoencoderKL) else {}),
-                **({"unet": unet} if unet is not None and isinstance(unet, UNet2DConditionModel) else {}),
-                **(
-                    {"text_encoder": text_encoder}
-                    if text_encoder is not None and isinstance(text_encoder, CLIPTextModel)
-                    else {}
-                ),
-                **(
-                    {"controlnet": controlnet}
-                    if controlnet is not None
-                    and (
-                        isinstance(controlnet, ControlNetModel)
-                        or all(isinstance(c, ControlNetModel) for c in controlnet)
-                    )
-                    else {}
-                ),
-            }
-        )
-
-        with ContextRblnConfig(
-            device=device,
-            device_map=device_map,
-            create_runtimes=create_runtimes,
-            optimze_host_mem=optimize_host_memory,
-        ):
-            model = super().from_pretrained(**{k: v for k, v in kwargs_dict.items() if v is not None})
-
-        if export is None or export is False:
-            return model
-
-        do_classifier_free_guidance = (
-            rbln_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        )
-
-        if not isinstance(vae, RBLNAutoencoderKL):
-            vae = RBLNAutoencoderKL.from_pretrained(
-                model_id=model_id,
-                subfolder="vae",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_unet_sample_size=model.unet.config.sample_size,
-                rbln_use_encode=False,
-                rbln_vae_scale_factor=model.vae_scale_factor,
-                rbln_config={**rbln_config},
-            )
-
-        if not isinstance(text_encoder, RBLNCLIPTextModel):
-            text_encoder = RBLNCLIPTextModel.from_pretrained(
-                model_id=model_id,
-                subfolder="text_encoder",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_config={**rbln_config},
-            )
-
-        if not isinstance(text_encoder_2, RBLNCLIPTextModel):
-            text_encoder_2 = RBLNCLIPTextModelWithProjection.from_pretrained(
-                model_id=model_id,
-                subfolder="text_encoder_2",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_config={**rbln_config},
-            )
-
-        batch_size = rbln_config.pop("batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        if not isinstance(unet, RBLNUNet2DConditionModel):
-            unet = RBLNUNet2DConditionModel.from_pretrained(
-                model_id=model_id,
-                subfolder="unet",
-                export=True,
-                model_save_dir=model_save_dir,
-                rbln_max_seq_len=model.text_encoder.config.max_position_embeddings,
-                rbln_text_model_hidden_size=model.text_encoder_2.config.hidden_size,
-                rbln_batch_size=unet_batch_size,
-                rbln_use_encode=False,
-                rbln_vae_scale_factor=model.vae_scale_factor,
-                rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
-                rbln_config={**rbln_config},
-            )
-
-        if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
-            if isinstance(controlnet, (list, tuple)):
-                multicontrolnet = []
-                for i, cid in enumerate(controlnet):
-                    subfolder_name = "controlnet" if i == 0 else f"controlnet_{i}"
-                    multicontrolnet.append(
-                        RBLNControlNetModel.from_model(
-                            model=cid,
-                            subfolder=subfolder_name,
-                            model_save_dir=model_save_dir,
-                            rbln_batch_size=unet_batch_size,
-                            rbln_text_model_hidden_size=model.text_encoder_2.config.hidden_size,
-                            rbln_vae_scale_factor=model.vae_scale_factor,
-                            rbln_config={**rbln_config},
-                        )
-                    )
-                controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
-                controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
-            else:
-                controlnet = RBLNControlNetModel.from_model(
-                    model=controlnet,
-                    subfolder="controlnet",
-                    model_save_dir=model_save_dir,
-                    rbln_batch_size=unet_batch_size,
-                    rbln_text_model_hidden_size=model.text_encoder_2.config.hidden_size,
-                    rbln_vae_scale_factor=model.vae_scale_factor,
-                    rbln_config={**rbln_config},
-                )
-                controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
-
-        if model_save_dir is not None:
-            # To skip saving original pytorch modules
-            del (model.vae, model.text_encoder, model.unet, model.controlnet)
-
-            # Direct calling of `save_pretrained` causes config.unet = (None, None).
-            # So config must be saved again, later.
-            model.save_pretrained(model_save_dir)
-
-        # replace modules
-        model.vae = vae
-        model.text_encoder = text_encoder
-        model.unet = unet
-        model.text_encoder_2 = text_encoder_2
-        model.controlnet = controlnet
-
-        # update config to be able to load from file
-        update_dict = {
-            "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
-            "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-            "text_encoder_2": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "controlnet": controlnet_dict,
-        }
-        model.register_to_config(**update_dict)
-
-        if model_save_dir is not None:
-            # overwrite to replace incorrect config
-            model.save_config(model_save_dir)
-
-        return model
+class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetPipeline):
+    original_class = StableDiffusionXLControlNetPipeline
+    _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
     def check_inputs(
         self,
@@ -419,6 +235,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 )
 
     @torch.no_grad()
+    @remove_compile_time_kwargs
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
@@ -682,6 +499,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
         text_encoder_lora_scale = (
             self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
         )
+
         (
             prompt_embeds,
             negative_prompt_embeds,