optimum-rbln 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

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

@@ -26,224 +26,23 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
-from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline
+from diffusers import StableDiffusionXLControlNetImg2ImgPipeline
 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
-from transformers import CLIPTextModel
 
-from ....modeling_base import RBLNBaseModel
-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 RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNetImg2ImgPipeline):
-    @classmethod
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for image-to-image generation using Stable Diffusion XL with ControlNet.
-
-        This model inherits from [`StableDiffusionXLControlNetImg2ImgPipeline`]. 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 RBLNStableDiffusionXLControlNetImg2Img 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_kwargs, _ = RBLNBaseModel.resolve_rbln_config(rbln_config, kwargs)
-
-        device = rbln_kwargs.get("device", None)
-        device_map = rbln_kwargs.get("device_map", None)
-        create_runtimes = rbln_kwargs.get("create_runtimes", None)
-        optimize_host_memory = rbln_kwargs.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_kwargs.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=True,
-                rbln_vae_scale_factor=model.vae_scale_factor,
-                rbln_config={**rbln_kwargs},
-            )
-
-        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_kwargs},
-            )
-
-        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_kwargs},
-            )
-
-        batch_size = rbln_kwargs.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=True,
-                rbln_vae_scale_factor=model.vae_scale_factor,
-                rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
-                rbln_config={**rbln_kwargs},
-            )
-
-        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_pretrained(
-                            model_id=cid.config._name_or_path,
-                            subfolder=subfolder_name,
-                            export=True,
-                            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_kwargs},
-                        )
-                    )
-                controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
-                controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
-            else:
-                controlnet = RBLNControlNetModel.from_pretrained(
-                    model_id=controlnet.config._name_or_path,
-                    subfolder="controlnet",
-                    export=True,
-                    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_kwargs},
-                )
-                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", "RBLNCLIPTextModelWithProjection"),
-            "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)
-
-        # use for CI to access each compiled model
-        if optimize_host_memory is False:
-            model.compiled_models = [
-                vae.compiled_models[0],
-                vae.compiled_models[1],
-                text_encoder.compiled_models[0],
-                text_encoder_2.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 RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetImg2ImgPipeline):
+    _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
     def check_inputs(
         self,
@@ -447,6 +246,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 )
 
     @torch.no_grad()
+    @remove_compile_time_kwargs
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
@@ -733,6 +533,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
         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,

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -24,116 +24,8 @@
 
 from diffusers import StableDiffusionPipeline
 
-from ....modeling_base import RBLNBaseModel
-from ....transformers import RBLNCLIPTextModel
-from ....utils.runtime_utils import ContextRblnConfig
-from ...models import RBLNAutoencoderKL, RBLNUNet2DConditionModel
+from ....modeling_diffusers import RBLNDiffusionMixin
 
 
-class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
-    @classmethod
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for text-to-image generation using Stable Diffusion.
-
-        This model inherits from [`StableDiffusionPipeline`]. 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 pipeline into a RBLNStableDiffusion 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)
-        model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config = kwargs.pop("rbln_config", None)
-        rbln_kwargs, _ = RBLNBaseModel.resolve_rbln_config(rbln_config, kwargs)
-
-        device = rbln_kwargs.get("device", None)
-        device_map = rbln_kwargs.get("device_map", None)
-        create_runtimes = rbln_kwargs.get("create_runtimes", None)
-        optimize_host_memory = rbln_kwargs.get("optimize_host_memory", None)
-
-        with ContextRblnConfig(
-            device=device,
-            device_map=device_map,
-            create_runtimes=create_runtimes,
-            optimze_host_mem=optimize_host_memory,
-        ):
-            model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
-
-        if export is None or export is False:
-            return model
-
-        do_classifier_free_guidance = (
-            rbln_kwargs.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        )
-
-        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_config={**rbln_kwargs},
-        )
-        text_encoder = RBLNCLIPTextModel.from_pretrained(
-            model_id=model_id,
-            subfolder="text_encoder",
-            export=True,
-            model_save_dir=model_save_dir,
-            rbln_config={**rbln_kwargs},
-        )
-
-        batch_size = rbln_kwargs.pop("batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        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_is_controlnet=True if "controlnet" in model.config.keys() else False,
-            rbln_config={**rbln_kwargs},
-        )
-
-        if model_save_dir is not None:
-            # To skip saving original pytorch modules
-            del (model.vae, model.text_encoder, model.unet)
-
-            # 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
-
-        # update config to be able to load from file.
-        update_dict = {
-            "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
-            "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-        }
-        model.register_to_config(**update_dict)
-
-        if model_save_dir is not None:
-            # overwrite to replace incorrect config
-            model.save_config(model_save_dir)
-
-        model.models = [vae.model[0], text_encoder.model[0], unet.model[0]]
-
-        if optimize_host_memory is False:
-            model.compiled_models = [vae.compiled_models[0], text_encoder.compiled_models[0], unet.compiled_models[0]]
-
-        return model
+class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):
+    _submodules = ["text_encoder", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -24,125 +24,20 @@
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
-from ....modeling_base import RBLNBaseModel
-from ....transformers import RBLNCLIPTextModel
-from ....utils.runtime_utils import ContextRblnConfig
-from ...models import RBLNAutoencoderKL, RBLNUNet2DConditionModel
+from ....modeling_diffusers import RBLNDiffusionMixin
 
 
-class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
-    @classmethod
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for image-to-image generation using Stable Diffusion.
+class RBLNStableDiffusionImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionImg2ImgPipeline):
+    """
+    Pipeline for image-to-image generation using Stable Diffusion.
 
-        This model inherits from [`StableDiffusionPipeline`]. Check the superclass documentation for the generic methods
-        implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+    This model inherits from [`StableDiffusionPipeline`]. 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 pipeline into a RBLNStableDiffusion pipeline by:
-        - transferring the checkpoint weights of the original into an optimized RBLN graph,
-        - compiling the resulting graph using the RBLN compiler.
+    It implements the methods to convert a pre-trained Stable Diffusion pipeline into a RBLNStableDiffusion 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)
-        model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config = kwargs.pop("rbln_config", None)
-        rbln_kwargs, _ = RBLNBaseModel.resolve_rbln_config(rbln_config, kwargs)
-
-        device = rbln_kwargs.get("device", None)
-        device_map = rbln_kwargs.get("device_map", None)
-        create_runtimes = rbln_kwargs.get("create_runtimes", None)
-        optimize_host_memory = rbln_kwargs.get("optimize_host_memory", None)
-
-        with ContextRblnConfig(
-            device=device,
-            device_map=device_map,
-            create_runtimes=create_runtimes,
-            optimze_host_mem=optimize_host_memory,
-        ):
-            model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
-
-        if export is None or export is False:
-            return model
-
-        do_classifier_free_guidance = (
-            rbln_kwargs.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        )
-
-        # compile model, create runtime
-        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_kwargs},
-        )
-        text_encoder = RBLNCLIPTextModel.from_pretrained(
-            model_id=model_id,
-            subfolder="text_encoder",
-            export=True,
-            model_save_dir=model_save_dir,
-            rbln_config={**rbln_kwargs},
-        )
-
-        batch_size = rbln_kwargs.pop("batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        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_kwargs},
-        )
-
-        if model_save_dir is not None:
-            # To skip saving original pytorch modules
-            del (model.vae, model.text_encoder, model.unet)
-
-            # 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
-
-        # update config to be able to load from file.
-        update_dict = {
-            "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
-            "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-        }
-        model.register_to_config(**update_dict)
-
-        if model_save_dir is not None:
-            # overwrite to replace incorrect config
-            model.save_config(model_save_dir)
-
-        # vae encoder, vae decoder, text_encoder, unet
-        model.models = [vae.model[0], vae.model[1], text_encoder.model[0], unet.model[0]]
-
-        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],
-            ]
-
-        return model
+    original_class = StableDiffusionImg2ImgPipeline
+    _submodules = ["text_encoder", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

@@ -16,128 +16,9 @@
 
 from diffusers import StableDiffusionXLPipeline
 
-from ....modeling_base import RBLNBaseModel
-from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
-from ....utils.runtime_utils import ContextRblnConfig
-from ...models import RBLNAutoencoderKL, RBLNUNet2DConditionModel
+from ....modeling_diffusers import RBLNDiffusionMixin
 
 
-class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
-    @classmethod
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for text-to-image generation using Stable Diffusion XL.
-
-        This model inherits from [`StableDiffusionXLPipeline`]. Check the superclass documentation for the generic methods the
-        library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
-
-        It implements the methods to convert a pre-trained StableDiffusionXL pipeline into a RBLNStableDiffusionXL 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)
-        model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config = kwargs.pop("rbln_config", None)
-        rbln_kwargs, _ = RBLNBaseModel.resolve_rbln_config(rbln_config, kwargs)
-
-        device = rbln_kwargs.get("device", None)
-        device_map = rbln_kwargs.get("device_map", None)
-        create_runtimes = rbln_kwargs.get("create_runtimes", None)
-        optimize_host_memory = rbln_kwargs.get("optimize_host_memory", None)
-
-        with ContextRblnConfig(
-            device=device,
-            device_map=device_map,
-            create_runtimes=create_runtimes,
-            optimze_host_mem=optimize_host_memory,
-        ):
-            model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
-
-        if export is None or export is False:
-            return model
-
-        do_classifier_free_guidance = (
-            rbln_kwargs.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        )
-
-        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_config={**rbln_kwargs},
-        )
-        text_encoder = RBLNCLIPTextModel.from_pretrained(
-            model_id=model_id,
-            subfolder="text_encoder",
-            export=True,
-            model_save_dir=model_save_dir,
-            rbln_config={**rbln_kwargs},
-        )
-        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_kwargs},
-        )
-
-        batch_size = rbln_kwargs.pop("batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        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_is_controlnet=True if "controlnet" in model.config.keys() else False,
-            rbln_config={**rbln_kwargs},
-        )
-
-        if model_save_dir is not None:
-            # To skip saving original pytorch modules
-            del (model.vae, model.text_encoder, model.unet, model.text_encoder_2)
-
-            # Direct calling of `save_pretrained` causes config.unet = (None, None).
-            # So config must be saved again, later.
-            model.save_pretrained(model_save_dir)
-
-        model.vae = vae
-        model.text_encoder = text_encoder
-        model.unet = unet
-        model.text_encoder_2 = text_encoder_2
-        update_dict = {
-            "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
-            "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-            "text_encoder_2": ("optimum.rbln", "RBLNCLIPTextModelWithProjection"),
-        }
-        model.register_to_config(**update_dict)
-
-        if model_save_dir is not None:
-            # overwrite to replace incorrect config
-            model.save_config(model_save_dir)
-
-        model.models = [vae.model[0], unet.model[0], text_encoder.model[0], text_encoder_2.model[0]]
-
-        if optimize_host_memory is False:
-            model.compiled_models = [
-                vae.compiled_models[0],
-                unet.compiled_models[0],
-                text_encoder.compiled_models[0],
-                text_encoder_2.compiled_models[0],
-            ]
-
-        return model
+class RBLNStableDiffusionXLPipeline(RBLNDiffusionMixin, StableDiffusionXLPipeline):
+    original_class = StableDiffusionXLPipeline
+    _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]