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

optimum/rbln/diffusers/models/{unet_2d_condition.py → unets/unet_2d_condition.py}

@@ -29,9 +29,9 @@ import torch
 from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from transformers import PretrainedConfig
 
-from ...modeling_base import RBLNModel
-from ...modeling_config import RBLNCompileConfig, RBLNConfig
-from ...utils.context import override_auto_classes
+from ....modeling import RBLNModel
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 if TYPE_CHECKING:
@@ -125,6 +125,9 @@ class _UNet_SDXL(torch.nn.Module):
 
 
 class RBLNUNet2DConditionModel(RBLNModel):
+    hf_library_name = "diffusers"
+    auto_model_class = UNet2DConditionModel
+
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
         self.in_features = self.rbln_config.model_cfg.get("in_features", None)
@@ -140,15 +143,6 @@ class RBLNUNet2DConditionModel(RBLNModel):
 
             self.add_embedding = ADDEMBEDDING(LINEAR1(self.in_features))
 
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        with override_auto_classes(
-            config_func=UNet2DConditionModel.load_config,
-            model_func=UNet2DConditionModel.from_pretrained,
-        ):
-            rt = super().from_pretrained(*args, **kwargs)
-        return rt
-
     @classmethod
     def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
         if model.config.addition_embed_type == "text_time":
@@ -156,6 +150,61 @@ class RBLNUNet2DConditionModel(RBLNModel):
         else:
             return _UNet_SD(model).eval()
 
+    @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"))
+        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"]:
+                # 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 // pipe.vae_scale_factor
+                else:
+                    sample_size = (
+                        vae_sample_size[0] // pipe.vae_scale_factor,
+                        vae_sample_size[1] // pipe.vae_scale_factor,
+                    )
+            else:
+                sample_size = pipe.unet.config.sample_size
+        else:
+            sample_size = (image_size[0] // pipe.vae_scale_factor, image_size[1] // pipe.vae_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
+
+        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."
+                )
+
+        rbln_config.update(
+            {
+                "max_seq_len": pipe.text_encoder.config.max_position_embeddings,
+                "text_model_hidden_size": text_model_hidden_size,
+                "sample_size": cls.get_unet_sample_size(pipe, rbln_config),
+                "batch_size": batch_size,
+                "is_controlnet": "controlnet" in pipe.config.keys(),
+            }
+        )
+
+        return rbln_config
+
     @classmethod
     def _get_rbln_config(
         cls,
@@ -179,7 +228,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
             sample_size = (sample_size, sample_size)
 
         if max_seq_len is None:
-            raise ValueError("`rbln_max_seq_len` (ex. text_encoder's max_position_embeddings )must be specified")
+            raise ValueError("`rbln_max_seq_len` (ex. text_encoder's max_position_embeddings) must be specified.")
 
         input_info = [
             ("sample", [batch_size, model_config.in_channels, sample_size[0], sample_size[1]], "float32"),
@@ -216,15 +265,13 @@ class RBLNUNet2DConditionModel(RBLNModel):
             ]
             input_info.append(("mid_block_additional_residual", shape, "float32"))
 
-        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
-
         if hasattr(model_config, "addition_embed_type") and 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
-            rbln_compile_config.input_info.append(
-                ("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32")
-            )
-            rbln_compile_config.input_info.append(("time_ids", [batch_size, 6], "float32"))
+            input_info.append(("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32"))
+            input_info.append(("time_ids", [batch_size, 6], "float32"))
+
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
@@ -237,6 +284,10 @@ class RBLNUNet2DConditionModel(RBLNModel):
 
         return rbln_config
 
+    @property
+    def compiled_batch_size(self):
+        return self.rbln_config.compile_cfgs[0].input_info[0][1][0]
+
     def forward(
         self,
         sample: torch.Tensor,
@@ -254,9 +305,18 @@ class RBLNUNet2DConditionModel(RBLNModel):
         return_dict: bool = True,
         **kwargs,
     ):
-        """
-        arg order : latent_model_input, t, prompt_embeds
-        """
+        sample_batch_size = sample.size()[0]
+        compiled_batch_size = self.compiled_batch_size
+        if sample_batch_size != compiled_batch_size and (
+            sample_batch_size * 2 == compiled_batch_size or sample_batch_size == compiled_batch_size * 2
+        ):
+            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"
+            )
+
         added_cond_kwargs = {} if added_cond_kwargs is None else added_cond_kwargs
 
         if down_block_additional_residuals is not None:

optimum/rbln/diffusers/pipelines/__init__.py

@@ -20,6 +20,7 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
+
 from typing import TYPE_CHECKING
 
 from transformers.utils import _LazyModule
@@ -36,8 +37,18 @@ _import_structure = {
     "stable_diffusion": [
         "RBLNStableDiffusionImg2ImgPipeline",
         "RBLNStableDiffusionPipeline",
+        "RBLNStableDiffusionInpaintPipeline",
+    ],
+    "stable_diffusion_xl": [
+        "RBLNStableDiffusionXLImg2ImgPipeline",
+        "RBLNStableDiffusionXLPipeline",
+        "RBLNStableDiffusionXLInpaintPipeline",
+    ],
+    "stable_diffusion_3": [
+        "RBLNStableDiffusion3Pipeline",
+        "RBLNStableDiffusion3Img2ImgPipeline",
+        "RBLNStableDiffusion3InpaintPipeline",
     ],
-    "stable_diffusion_xl": ["RBLNStableDiffusionXLImg2ImgPipeline", "RBLNStableDiffusionXLPipeline"],
 }
 if TYPE_CHECKING:
     from .controlnet import (
@@ -49,9 +60,19 @@ if TYPE_CHECKING:
     )
     from .stable_diffusion import (
         RBLNStableDiffusionImg2ImgPipeline,
+        RBLNStableDiffusionInpaintPipeline,
         RBLNStableDiffusionPipeline,
     )
-    from .stable_diffusion_xl import RBLNStableDiffusionXLImg2ImgPipeline, RBLNStableDiffusionXLPipeline
+    from .stable_diffusion_3 import (
+        RBLNStableDiffusion3Img2ImgPipeline,
+        RBLNStableDiffusion3InpaintPipeline,
+        RBLNStableDiffusion3Pipeline,
+    )
+    from .stable_diffusion_xl import (
+        RBLNStableDiffusionXLImg2ImgPipeline,
+        RBLNStableDiffusionXLInpaintPipeline,
+        RBLNStableDiffusionXLPipeline,
+    )
 else:
     import sys
 

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

@@ -27,13 +27,9 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import torch
-from diffusers import ControlNetModel
 from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
-from optimum.exporters import TasksManager
-from transformers import AutoConfig, AutoModel
 
-from ....modeling_base import RBLNModel
-from ....modeling_config import RBLNConfig
+from ....modeling import RBLNModel
 from ...models.controlnet import RBLNControlNetModel
 
 
@@ -44,10 +40,12 @@ logger = logging.getLogger(__name__)
 
 
 class RBLNMultiControlNetModel(RBLNModel):
+    hf_library_name = "diffusers"
+    _hf_class = MultiControlNetModel
+
     def __init__(
         self,
         models: List[RBLNControlNetModel],
-        **kwargs,
     ):
         self.nets = models
         self.dtype = torch.float32
@@ -59,27 +57,6 @@ class RBLNMultiControlNetModel(RBLNModel):
             cm.extend(net.compiled_models)
         return cm
 
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        def get_model_from_task(
-            task: str,
-            model_name_or_path: Union[str, Path],
-            **kwargs,
-        ):
-            return MultiControlNetModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path, **kwargs)
-
-        tasktmp = TasksManager.get_model_from_task
-        configtmp = AutoConfig.from_pretrained
-        modeltmp = AutoModel.from_pretrained
-        TasksManager.get_model_from_task = get_model_from_task
-        AutoConfig.from_pretrained = ControlNetModel.load_config
-        AutoModel.from_pretrained = MultiControlNetModel.from_pretrained
-        rt = super().from_pretrained(*args, **kwargs)
-        AutoConfig.from_pretrained = configtmp
-        AutoModel.from_pretrained = modeltmp
-        TasksManager.get_model_from_task = tasktmp
-        return rt
-
     @classmethod
     def _from_pretrained(
         cls,
@@ -88,19 +65,22 @@ class RBLNMultiControlNetModel(RBLNModel):
     ) -> RBLNModel:
         idx = 0
         controlnets = []
-        model_path_to_load = model_id
+        subfolder_name = kwargs.pop("subfolder", None)
+        if subfolder_name is not None:
+            model_path_to_load = model_id + "/" + subfolder_name
+        else:
+            model_path_to_load = model_id
+
+        base_model_path_to_load = model_path_to_load
 
         while os.path.isdir(model_path_to_load):
             controlnet = RBLNControlNetModel.from_pretrained(model_path_to_load, export=False, **kwargs)
             controlnets.append(controlnet)
-            rbln_config = RBLNConfig.load(model_path_to_load)
             idx += 1
-            model_path_to_load = model_id + f"_{idx}"
+            model_path_to_load = base_model_path_to_load + f"_{idx}"
 
         return cls(
             controlnets,
-            rbln_config=rbln_config,
-            **kwargs,
         )
 
     def save_pretrained(self, save_directory: Union[str, Path], **kwargs):
@@ -118,7 +98,7 @@ class RBLNMultiControlNetModel(RBLNModel):
         sample: torch.FloatTensor,
         timestep: Union[torch.Tensor, float, int],
         encoder_hidden_states: torch.Tensor,
-        controlnet_cond: List[torch.tensor],
+        controlnet_cond: List[torch.Tensor],
         conditioning_scale: List[float],
         class_labels: Optional[torch.Tensor] = None,
         timestep_cond: Optional[torch.Tensor] = None,

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

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Team. 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.
+
 # Copyright 2024 Rebellions Inc.
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,7 +34,6 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -33,8 +46,8 @@ 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 ....modeling_diffusers import RBLNDiffusionMixin
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -46,6 +59,7 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
     original_class = StableDiffusionControlNetPipeline
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet.py
     def check_inputs(
         self,
         prompt,
@@ -209,6 +223,7 @@ class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionC
                     f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
                 )
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet.py
     @torch.no_grad()
     @remove_compile_time_kwargs
     def __call__(

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

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Team. 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.
+
 # Copyright 2024 Rebellions Inc.
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,7 +34,6 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -32,8 +45,8 @@ from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
 
-from ....modeling_diffusers import RBLNDiffusionMixin
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -42,8 +55,10 @@ logger = logging.get_logger(__name__)
 
 
 class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionControlNetImg2ImgPipeline):
+    original_class = StableDiffusionControlNetImg2ImgPipeline
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_img2img.py
     def check_inputs(
         self,
         prompt,
@@ -201,6 +216,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
                     f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
                 )
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_img2img.py
     @torch.no_grad()
     @remove_compile_time_kwargs
     def __call__(

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

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Team. 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.
+
 # Copyright 2024 Rebellions Inc.
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,7 +34,6 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionXLPipeline class for inference of diffusion models on rbln devices."""
 
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -32,8 +45,8 @@ from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffus
 from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
 
-from ....modeling_diffusers import RBLNDiffusionMixin
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -42,8 +55,10 @@ logger = logging.get_logger(__name__)
 
 
 class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetPipeline):
+    original_class = StableDiffusionXLControlNetPipeline
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.py
     def check_inputs(
         self,
         prompt,
@@ -233,6 +248,7 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
                     f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
                 )
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.py
     @torch.no_grad()
     @remove_compile_time_kwargs
     def __call__(

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

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Team. 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.
+
 # Copyright 2024 Rebellions Inc.
 
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,7 +34,6 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionXLPipeline class for inference of diffusion models on rbln devices."""
 
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -32,8 +45,8 @@ from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffus
 from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
 
-from ....modeling_diffusers import RBLNDiffusionMixin
 from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...modeling_diffusers import RBLNDiffusionMixin
 from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -42,8 +55,10 @@ logger = logging.get_logger(__name__)
 
 
 class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLControlNetImg2ImgPipeline):
+    original_class = StableDiffusionXLControlNetImg2ImgPipeline
     _submodules = ["text_encoder", "text_encoder_2", "unet", "vae", "controlnet"]
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl_img2img.py
     def check_inputs(
         self,
         prompt,
@@ -245,6 +260,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableD
                     f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
                 )
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl_img2img.py
     @torch.no_grad()
     @remove_compile_time_kwargs
     def __call__(

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

@@ -23,3 +23,4 @@
 
 from .pipeline_stable_diffusion import RBLNStableDiffusionPipeline
 from .pipeline_stable_diffusion_img2img import RBLNStableDiffusionImg2ImgPipeline
+from .pipeline_stable_diffusion_inpaint import RBLNStableDiffusionInpaintPipeline

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

@@ -20,12 +20,12 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusionPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):
+    original_class = StableDiffusionPipeline
     _submodules = ["text_encoder", "unet", "vae"]

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

@@ -20,24 +20,12 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
-"""RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 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.).
-
-    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.
-    """
-
     original_class = StableDiffusionImg2ImgPipeline
     _submodules = ["text_encoder", "unet", "vae"]

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

@@ -0,0 +1,31 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from diffusers import StableDiffusionInpaintPipeline
+
+from ...modeling_diffusers import RBLNDiffusionMixin
+
+
+class RBLNStableDiffusionInpaintPipeline(RBLNDiffusionMixin, StableDiffusionInpaintPipeline):
+    original_class = StableDiffusionInpaintPipeline
+    _submodules = ["text_encoder", "unet", "vae"]

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

@@ -0,0 +1,26 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .pipeline_stable_diffusion_3 import RBLNStableDiffusion3Pipeline
+from .pipeline_stable_diffusion_3_img2img import RBLNStableDiffusion3Img2ImgPipeline
+from .pipeline_stable_diffusion_3_inpaint import RBLNStableDiffusion3InpaintPipeline

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

@@ -0,0 +1,31 @@
+# Copyright 2024 Rebellions Inc.
+
+# 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.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from diffusers import StableDiffusion3Pipeline
+
+from ...modeling_diffusers import RBLNDiffusionMixin
+
+
+class RBLNStableDiffusion3Pipeline(RBLNDiffusionMixin, StableDiffusion3Pipeline):
+    original_class = StableDiffusion3Pipeline
+    _submodules = ["transformer", "text_encoder_3", "text_encoder", "text_encoder_2", "vae"]