optimum-rbln 0.7.3.post1__py3-none-any.whl → 0.7.4__py3-none-any.whl

optimum/rbln/diffusers/models/unets/unet_2d_condition.py

@@ -16,17 +16,18 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
-from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel, UNet2DConditionOutput
 from transformers import PretrainedConfig
 
+from ....configuration_utils import RBLNCompileConfig
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
-from ...modeling_diffusers import RBLNDiffusionMixin
+from ...configurations import RBLNUNet2DConditionModelConfig
+from ...modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
 
 logger = get_logger(__name__)
 
@@ -141,10 +142,13 @@ class _UNet_Kandinsky(torch.nn.Module):
 class RBLNUNet2DConditionModel(RBLNModel):
     hf_library_name = "diffusers"
     auto_model_class = UNet2DConditionModel
+    _rbln_config_class = RBLNUNet2DConditionModelConfig
+    output_class = UNet2DConditionOutput
+    output_key = "sample"
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
-        self.in_features = self.rbln_config.model_cfg.get("in_features", None)
+        self.in_features = self.rbln_config.in_features
         if self.in_features is not None:
 
             @dataclass
@@ -158,7 +162,9 @@ class RBLNUNet2DConditionModel(RBLNModel):
             self.add_embedding = ADDEMBEDDING(LINEAR1(self.in_features))
 
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+    def wrap_model_if_needed(
+        cls, model: torch.nn.Module, rbln_config: RBLNUNet2DConditionModelConfig
+    ) -> torch.nn.Module:
         if model.config.addition_embed_type == "text_time":
             return _UNet_SDXL(model).eval()
         elif model.config.addition_embed_type == "image":
@@ -168,117 +174,117 @@ class RBLNUNet2DConditionModel(RBLNModel):
 
     @classmethod
     def get_unet_sample_size(
-        cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]
-    ) -> Union[int, Tuple[int, int]]:
-        image_size = (rbln_config.get("img_height"), rbln_config.get("img_width"))
+        cls,
+        pipe: RBLNDiffusionMixin,
+        rbln_config: RBLNUNet2DConditionModelConfig,
+        image_size: Optional[Tuple[int, int]] = None,
+    ) -> Tuple[int, int]:
         scale_factor = pipe.movq_scale_factor if hasattr(pipe, "movq_scale_factor") else pipe.vae_scale_factor
-        if (image_size[0] is None) != (image_size[1] is None):
-            raise ValueError("Both image height and image width must be given or not given")
-        elif image_size[0] is None and image_size[1] is None:
-            if rbln_config["img2img_pipeline"]:
+
+        if image_size is None:
+            if "Img2Img" in pipe.__class__.__name__:
                 if hasattr(pipe, "vae"):
                     # In case of img2img, sample size of unet is determined by vae encoder.
                     vae_sample_size = pipe.vae.config.sample_size
                     if isinstance(vae_sample_size, int):
-                        sample_size = vae_sample_size // scale_factor
-                    else:
-                        sample_size = (
-                            vae_sample_size[0] // scale_factor,
-                            vae_sample_size[1] // scale_factor,
-                        )
+                        vae_sample_size = (vae_sample_size, vae_sample_size)
+
+                    sample_size = (
+                        vae_sample_size[0] // scale_factor,
+                        vae_sample_size[1] // scale_factor,
+                    )
                 elif hasattr(pipe, "movq"):
                     logger.warning(
-                        "RBLN config 'img_height' and 'img_width' should have been provided for this pipeline. "
+                        "RBLN config 'image_size' should have been provided for this pipeline. "
                         "Both variable will be set 512 by default."
                     )
                     sample_size = (512 // scale_factor, 512 // scale_factor)
             else:
                 sample_size = pipe.unet.config.sample_size
+                if isinstance(sample_size, int):
+                    sample_size = (sample_size, sample_size)
         else:
             sample_size = (image_size[0] // scale_factor, image_size[1] // scale_factor)
 
         return sample_size
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        text_model_hidden_size = pipe.text_encoder_2.config.hidden_size if hasattr(pipe, "text_encoder_2") else None
-        image_model_hidden_size = pipe.unet.config.encoder_hid_dim if hasattr(pipe, "unet") else None
-
-        batch_size = rbln_config.get("batch_size")
-        if not batch_size:
-            do_classifier_free_guidance = (
-                rbln_config.get("guidance_scale", 5.0) > 1.0 and pipe.unet.config.time_cond_proj_dim is None
-            )
-            batch_size = 2 if do_classifier_free_guidance else 1
-        else:
-            if rbln_config.get("guidance_scale"):
-                logger.warning(
-                    "guidance_scale is ignored because batch size is explicitly specified. "
-                    "To ensure consistent behavior, consider removing the guidance scale or "
-                    "adjusting the batch size configuration as needed."
-                )
+    def update_rbln_config_using_pipe(
+        cls, pipe: RBLNDiffusionMixin, rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
+        rbln_config.unet.text_model_hidden_size = (
+            pipe.text_encoder_2.config.hidden_size if hasattr(pipe, "text_encoder_2") else None
+        )
+        rbln_config.unet.image_model_hidden_size = pipe.unet.config.encoder_hid_dim if hasattr(pipe, "unet") else None
 
-        max_seq_len = pipe.text_encoder.config.max_position_embeddings if hasattr(pipe, "text_encoder") else None
-        rbln_config.update(
-            {
-                "max_seq_len": max_seq_len,
-                "text_model_hidden_size": text_model_hidden_size,
-                "image_model_hidden_size": image_model_hidden_size,
-                "sample_size": cls.get_unet_sample_size(pipe, rbln_config),
-                "batch_size": batch_size,
-                "is_controlnet": "controlnet" in pipe.config.keys(),
-            }
+        rbln_config.unet.max_seq_len = (
+            pipe.text_encoder.config.max_position_embeddings if hasattr(pipe, "text_encoder") else None
         )
 
+        rbln_config.unet.sample_size = cls.get_unet_sample_size(
+            pipe, rbln_config.unet, image_size=rbln_config.image_size
+        )
+        rbln_config.unet.use_additional_residuals = "controlnet" in pipe.config.keys()
+
         return rbln_config
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: "PreTrainedModel",
         model_config: "PretrainedConfig",
-        rbln_kwargs: Dict[str, Any] = {},
-    ) -> RBLNConfig:
-        batch_size = rbln_kwargs.get("batch_size")
-        max_seq_len = rbln_kwargs.get("max_seq_len")
-        sample_size = rbln_kwargs.get("sample_size")
-        is_controlnet = rbln_kwargs.get("is_controlnet")
-        rbln_in_features = None
-
-        if batch_size is None:
-            batch_size = 1
-
-        if sample_size is None:
-            sample_size = model_config.sample_size
+        rbln_config: RBLNUNet2DConditionModelConfig,
+    ) -> RBLNUNet2DConditionModelConfig:
+        if rbln_config.sample_size is None:
+            rbln_config.sample_size = model_config.sample_size
 
-        if isinstance(sample_size, int):
-            sample_size = (sample_size, sample_size)
+        if isinstance(rbln_config.sample_size, int):
+            rbln_config.sample_size = (rbln_config.sample_size, rbln_config.sample_size)
 
         input_info = [
-            ("sample", [batch_size, model_config.in_channels, sample_size[0], sample_size[1]], "float32"),
+            (
+                "sample",
+                [
+                    rbln_config.batch_size,
+                    model_config.in_channels,
+                    rbln_config.sample_size[0],
+                    rbln_config.sample_size[1],
+                ],
+                "float32",
+            ),
             ("timestep", [], "float32"),
         ]
 
-        if max_seq_len is not None:
+        if rbln_config.max_seq_len is not None:
             input_info.append(
-                ("encoder_hidden_states", [batch_size, max_seq_len, model_config.cross_attention_dim], "float32"),
+                (
+                    "encoder_hidden_states",
+                    [rbln_config.batch_size, rbln_config.max_seq_len, model_config.cross_attention_dim],
+                    "float32",
+                ),
             )
 
-        if is_controlnet:
+        if rbln_config.use_additional_residuals:
             # down block addtional residuals
-            first_shape = [batch_size, model_config.block_out_channels[0], sample_size[0], sample_size[1]]
-            height, width = sample_size[0], sample_size[1]
+            first_shape = [
+                rbln_config.batch_size,
+                model_config.block_out_channels[0],
+                rbln_config.sample_size[0],
+                rbln_config.sample_size[1],
+            ]
+            height, width = rbln_config.sample_size[0], rbln_config.sample_size[1]
             input_info.append(("down_block_additional_residuals_0", first_shape, "float32"))
             name_idx = 1
             for idx, _ in enumerate(model_config.down_block_types):
-                shape = [batch_size, model_config.block_out_channels[idx], height, width]
+                shape = [rbln_config.batch_size, model_config.block_out_channels[idx], height, width]
                 for _ in range(model_config.layers_per_block):
                     input_info.append((f"down_block_additional_residuals_{name_idx}", shape, "float32"))
                     name_idx += 1
                 if idx != len(model_config.down_block_types) - 1:
                     height = height // 2
                     width = width // 2
-                    shape = [batch_size, model_config.block_out_channels[idx], height, width]
+                    shape = [rbln_config.batch_size, model_config.block_out_channels[idx], height, width]
                     input_info.append((f"down_block_additional_residuals_{name_idx}", shape, "float32"))
                     name_idx += 1
 
@@ -286,33 +292,27 @@ class RBLNUNet2DConditionModel(RBLNModel):
             num_cross_attn_blocks = model_config.down_block_types.count("CrossAttnDownBlock2D")
             out_channels = model_config.block_out_channels[-1]
             shape = [
-                batch_size,
+                rbln_config.batch_size,
                 out_channels,
-                sample_size[0] // 2**num_cross_attn_blocks,
-                sample_size[1] // 2**num_cross_attn_blocks,
+                rbln_config.sample_size[0] // 2**num_cross_attn_blocks,
+                rbln_config.sample_size[1] // 2**num_cross_attn_blocks,
             ]
             input_info.append(("mid_block_additional_residual", shape, "float32"))
 
         if hasattr(model_config, "addition_embed_type"):
             if model_config.addition_embed_type == "text_time":
-                rbln_text_model_hidden_size = rbln_kwargs["text_model_hidden_size"]
-                rbln_in_features = model_config.projection_class_embeddings_input_dim
-                input_info.append(("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32"))
-                input_info.append(("time_ids", [batch_size, 6], "float32"))
+                rbln_config.in_features = model_config.projection_class_embeddings_input_dim
+                input_info.append(
+                    ("text_embeds", [rbln_config.batch_size, rbln_config.text_model_hidden_size], "float32")
+                )
+                input_info.append(("time_ids", [rbln_config.batch_size, 6], "float32"))
             elif model_config.addition_embed_type == "image":
-                rbln_image_model_hidden_size = rbln_kwargs["image_model_hidden_size"]
-                input_info.append(("image_embeds", [batch_size, rbln_image_model_hidden_size], "float32"))
+                input_info.append(
+                    ("image_embeds", [rbln_config.batch_size, rbln_config.image_model_hidden_size], "float32")
+                )
 
         rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-
-        rbln_config = RBLNConfig(
-            rbln_cls=cls.__name__,
-            compile_cfgs=[rbln_compile_config],
-            rbln_kwargs=rbln_kwargs,
-        )
-
-        if rbln_in_features is not None:
-            rbln_config.model_cfg["in_features"] = rbln_in_features
+        rbln_config.set_compile_cfgs([rbln_compile_config])
 
         return rbln_config
 
@@ -336,7 +336,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
         encoder_attention_mask: Optional[torch.Tensor] = None,
         return_dict: bool = True,
         **kwargs,
-    ):
+    ) -> Union[UNet2DConditionOutput, Tuple]:
         sample_batch_size = sample.size()[0]
         compiled_batch_size = self.compiled_batch_size
         if sample_batch_size != compiled_batch_size and (
@@ -344,40 +344,37 @@ class RBLNUNet2DConditionModel(RBLNModel):
         ):
             raise ValueError(
                 f"Mismatch between UNet's runtime batch size ({sample_batch_size}) and compiled batch size ({compiled_batch_size}). "
-                "This may be caused by the 'guidance scale' parameter, which doubles the runtime batch size in Stable Diffusion. "
-                "Adjust the batch size during compilation or modify the 'guidance scale' to match the compiled batch size.\n\n"
-                "For details, see: https://docs.rbln.ai/software/optimum/model_api.html#stable-diffusion"
+                "This may be caused by the 'guidance scale' parameter, which doubles the runtime batch size of UNet in Stable Diffusion. "
+                "Adjust the batch size of UNet during compilation to match the runtime batch size.\n\n"
+                "For details, see: https://docs.rbln.ai/software/optimum/model_api/diffusers/pipelines/stable_diffusion.html#important-batch-size-configuration-for-guidance-scale"
             )
 
         added_cond_kwargs = {} if added_cond_kwargs is None else added_cond_kwargs
 
         if down_block_additional_residuals is not None:
             down_block_additional_residuals = [t.contiguous() for t in down_block_additional_residuals]
-            return (
-                super().forward(
-                    sample.contiguous(),
-                    timestep.float(),
-                    encoder_hidden_states,
-                    *down_block_additional_residuals,
-                    mid_block_additional_residual,
-                    **added_cond_kwargs,
-                ),
+            return super().forward(
+                sample.contiguous(),
+                timestep.float(),
+                encoder_hidden_states,
+                *down_block_additional_residuals,
+                mid_block_additional_residual,
+                **added_cond_kwargs,
+                return_dict=return_dict,
             )
 
         if "image_embeds" in added_cond_kwargs:
-            return (
-                super().forward(
-                    sample.contiguous(),
-                    timestep.float(),
-                    **added_cond_kwargs,
-                ),
-            )
-
-        return (
-            super().forward(
+            return super().forward(
                 sample.contiguous(),
                 timestep.float(),
-                encoder_hidden_states,
                 **added_cond_kwargs,
-            ),
+                return_dict=return_dict,
+            )
+
+        return super().forward(
+            sample.contiguous(),
+            timestep.float(),
+            encoder_hidden_states,
+            **added_cond_kwargs,
+            return_dict=return_dict,
         )

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

@@ -81,7 +81,7 @@ class RBLNMultiControlNetModel(RBLNModel):
             model.save_pretrained(real_save_path)
 
     @classmethod
-    def _get_rbln_config(cls, **rbln_config_kwargs):
+    def _update_rbln_config(cls, **rbln_config_kwargs):
         pass
 
     def forward(
@@ -100,16 +100,15 @@ class RBLNMultiControlNetModel(RBLNModel):
         return_dict: bool = True,
     ):
         for i, (image, scale, controlnet) in enumerate(zip(controlnet_cond, conditioning_scale, self.nets)):
-            output = controlnet.model[0](
+            down_samples, mid_sample = controlnet(
                 sample=sample.contiguous(),
                 timestep=timestep.float(),
                 encoder_hidden_states=encoder_hidden_states,
                 controlnet_cond=image,
                 conditioning_scale=torch.tensor(scale),
+                return_dict=return_dict,
             )
 
-            down_samples, mid_sample = output[:-1], output[-1]
-
             # merge samples
             if i == 0:
                 down_block_res_samples, mid_block_res_sample = down_samples, mid_sample

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

@@ -39,6 +39,7 @@ from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
 from ....utils.logging import get_logger
+from ...configurations import RBLNStableDiffusionControlNetPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
@@ -49,6 +50,7 @@ logger = get_logger(__name__)
 
 class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionControlNetPipeline):
     original_class = StableDiffusionControlNetPipeline
+    _rbln_config_class = RBLNStableDiffusionControlNetPipelineConfig
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]
 
     # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet.py

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

@@ -37,6 +37,7 @@ from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...configurations import RBLNStableDiffusionControlNetImg2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
@@ -48,6 +49,7 @@ logger = logging.get_logger(__name__)
 class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionControlNetImg2ImgPipeline):
     original_class = StableDiffusionControlNetImg2ImgPipeline
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]
+    _rbln_config_class = RBLNStableDiffusionControlNetImg2ImgPipelineConfig
 
     # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_img2img.py
     def check_inputs(

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

@@ -37,6 +37,7 @@ from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...configurations import RBLNStableDiffusionXLControlNetPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
@@ -47,6 +48,7 @@ logger = logging.get_logger(__name__)
 
 class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetPipeline):
     original_class = StableDiffusionXLControlNetPipeline
+    _rbln_config_class = RBLNStableDiffusionXLControlNetPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
     # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.py

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

@@ -37,6 +37,7 @@ from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
 
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...configurations import RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
@@ -47,6 +48,7 @@ logger = logging.get_logger(__name__)
 
 class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetImg2ImgPipeline):
     original_class = StableDiffusionXLControlNetImg2ImgPipeline
+    _rbln_config_class = RBLNStableDiffusionXLControlNetImg2ImgPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
     # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl_img2img.py

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py

@@ -14,11 +14,13 @@
 
 from diffusers import KandinskyV22Pipeline
 
+from ...configurations import RBLNKandinskyV22PipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22Pipeline(RBLNDiffusionMixin, KandinskyV22Pipeline):
     original_class = KandinskyV22Pipeline
+    _rbln_config_class = RBLNKandinskyV22PipelineConfig
     _submodules = ["unet", "movq"]
 
     def get_compiled_image_size(self):

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py

@@ -29,6 +29,7 @@ from transformers import (
     CLIPVisionModelWithProjection,
 )
 
+from ...configurations import RBLNKandinskyV22CombinedPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 from .pipeline_kandinsky2_2 import RBLNKandinskyV22Pipeline
 from .pipeline_kandinsky2_2_img2img import RBLNKandinskyV22Img2ImgPipeline
@@ -38,6 +39,7 @@ from .pipeline_kandinsky2_2_prior import RBLNKandinskyV22PriorPipeline
 
 class RBLNKandinskyV22CombinedPipeline(RBLNDiffusionMixin, KandinskyV22CombinedPipeline):
     original_class = KandinskyV22CombinedPipeline
+    _rbln_config_class = RBLNKandinskyV22CombinedPipelineConfig
     _connected_classes = {"prior_pipe": RBLNKandinskyV22PriorPipeline, "decoder_pipe": RBLNKandinskyV22Pipeline}
     _submodules = ["prior_image_encoder", "prior_text_encoder", "prior_prior", "unet", "movq"]
     _prefix = {"prior_pipe": "prior_"}

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py

@@ -14,11 +14,13 @@
 
 from diffusers import KandinskyV22Img2ImgPipeline
 
+from ...configurations import RBLNKandinskyV22Img2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22Img2ImgPipeline(RBLNDiffusionMixin, KandinskyV22Img2ImgPipeline):
     original_class = KandinskyV22Img2ImgPipeline
+    _rbln_config_class = RBLNKandinskyV22Img2ImgPipelineConfig
     _submodules = ["unet", "movq"]
 
     def get_compiled_image_size(self):

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpaint.py

@@ -14,11 +14,13 @@
 
 from diffusers import KandinskyV22InpaintPipeline
 
+from ...configurations import RBLNKandinskyV22InpaintPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22InpaintPipeline(RBLNDiffusionMixin, KandinskyV22InpaintPipeline):
     original_class = KandinskyV22InpaintPipeline
+    _rbln_config_class = RBLNKandinskyV22InpaintPipelineConfig
     _submodules = ["unet", "movq"]
 
     def get_compiled_image_size(self):

optimum/rbln/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py

@@ -14,9 +14,11 @@
 
 from diffusers import KandinskyV22PriorPipeline
 
+from ...configurations import RBLNKandinskyV22PriorPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNKandinskyV22PriorPipeline(RBLNDiffusionMixin, KandinskyV22PriorPipeline):
     original_class = KandinskyV22PriorPipeline
+    _rbln_config_class = RBLNKandinskyV22PriorPipelineConfig
     _submodules = ["text_encoder", "image_encoder", "prior"]

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

@@ -12,11 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 from diffusers import StableDiffusionPipeline
 
+from ...configurations import RBLNStableDiffusionPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):
     original_class = StableDiffusionPipeline
-    _submodules = ["text_encoder", "unet", "vae"]
+    _rbln_config_class = RBLNStableDiffusionPipelineConfig
+    _submodules = ["vae", "text_encoder", "unet"]

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

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
+from ...configurations import RBLNStableDiffusionImg2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionImg2ImgPipeline):
     original_class = StableDiffusionImg2ImgPipeline
+    _rbln_config_class = RBLNStableDiffusionImg2ImgPipelineConfig
     _submodules = ["text_encoder", "unet", "vae"]

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

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionInpaintPipeline
 
+from ...configurations import RBLNStableDiffusionInpaintPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionInpaintPipeline(RBLNDiffusionMixin, StableDiffusionInpaintPipeline):
     original_class = StableDiffusionInpaintPipeline
+    _rbln_config_class = RBLNStableDiffusionInpaintPipelineConfig
     _submodules = ["text_encoder", "unet", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusion3Pipeline
 
+from ...configurations import RBLNStableDiffusion3PipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Pipeline(RBLNDiffusionMixin, StableDiffusion3Pipeline):
     original_class = StableDiffusion3Pipeline
+    _rbln_config_class = RBLNStableDiffusion3PipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusion3Img2ImgPipeline
 
+from ...configurations import RBLNStableDiffusion3Img2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3Img2ImgPipeline(RBLNDiffusionMixin, StableDiffusion3Img2ImgPipeline):
     original_class = StableDiffusion3Img2ImgPipeline
+    _rbln_config_class = RBLNStableDiffusion3Img2ImgPipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

optimum/rbln/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusion3InpaintPipeline
 
+from ...configurations import RBLNStableDiffusion3InpaintPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusion3InpaintPipeline(RBLNDiffusionMixin, StableDiffusion3InpaintPipeline):
     original_class = StableDiffusion3InpaintPipeline
+    _rbln_config_class = RBLNStableDiffusion3InpaintPipelineConfig
     _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]

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

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionXLPipeline
 
+from ...configurations import RBLNStableDiffusionXLPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLPipeline(RBLNDiffusionMixin, StableDiffusionXLPipeline):
     original_class = StableDiffusionXLPipeline
+    _rbln_config_class = RBLNStableDiffusionXLPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

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

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionXLImg2ImgPipeline
 
+from ...configurations import RBLNStableDiffusionXLImg2ImgPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLImg2ImgPipeline):
     original_class = StableDiffusionXLImg2ImgPipeline
+    _rbln_config_class = RBLNStableDiffusionXLImg2ImgPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

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

@@ -14,9 +14,11 @@
 
 from diffusers import StableDiffusionXLInpaintPipeline
 
+from ...configurations import RBLNStableDiffusionXLInpaintPipelineConfig
 from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionXLInpaintPipeline(RBLNDiffusionMixin, StableDiffusionXLInpaintPipeline):
     original_class = StableDiffusionXLInpaintPipeline
+    _rbln_config_class = RBLNStableDiffusionXLInpaintPipelineConfig
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]