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

optimum/rbln/__init__.py

@@ -30,24 +30,13 @@ from .utils import check_version_compats
 
 
 _import_structure = {
-    "modeling_alias": [
-        "RBLNASTForAudioClassification",
-        "RBLNBertForQuestionAnswering",
-        "RBLNDistilBertForQuestionAnswering",
-        "RBLNResNetForImageClassification",
-        "RBLNXLMRobertaForSequenceClassification",
-        "RBLNRobertaForSequenceClassification",
-        "RBLNRobertaForMaskedLM",
-        "RBLNViTForImageClassification",
-    ],
     "modeling": [
         "RBLNBaseModel",
         "RBLNModel",
-        "RBLNModelForQuestionAnswering",
-        "RBLNModelForAudioClassification",
-        "RBLNModelForImageClassification",
-        "RBLNModelForSequenceClassification",
-        "RBLNModelForMaskedLM",
+    ],
+    "modeling_config": [
+        "RBLNCompileConfig",
+        "RBLNConfig",
     ],
     "transformers": [
         "RBLNAutoModel",
@@ -83,6 +72,14 @@ _import_structure = {
         "RBLNMistralForCausalLM",
         "RBLNWhisperForConditionalGeneration",
         "RBLNXLMRobertaModel",
+        "RBLNASTForAudioClassification",
+        "RBLNBertForQuestionAnswering",
+        "RBLNDistilBertForQuestionAnswering",
+        "RBLNResNetForImageClassification",
+        "RBLNXLMRobertaForSequenceClassification",
+        "RBLNRobertaForSequenceClassification",
+        "RBLNRobertaForMaskedLM",
+        "RBLNViTForImageClassification",
     ],
     "diffusers": [
         "RBLNStableDiffusionPipeline",
@@ -103,15 +100,15 @@ _import_structure = {
         "RBLNStableDiffusion3Img2ImgPipeline",
         "RBLNStableDiffusion3InpaintPipeline",
         "RBLNStableDiffusion3Pipeline",
+        "RBLNDiffusionMixin",
     ],
-    "modeling_config": ["RBLNCompileConfig", "RBLNConfig"],
-    "modeling_diffusers": ["RBLNDiffusionMixin"],
 }
 
 if TYPE_CHECKING:
     from .diffusers import (
         RBLNAutoencoderKL,
         RBLNControlNetModel,
+        RBLNDiffusionMixin,
         RBLNMultiControlNetModel,
         RBLNSD3Transformer2DModel,
         RBLNStableDiffusion3Img2ImgPipeline,
@@ -132,25 +129,13 @@ if TYPE_CHECKING:
     from .modeling import (
         RBLNBaseModel,
         RBLNModel,
-        RBLNModelForAudioClassification,
-        RBLNModelForImageClassification,
-        RBLNModelForMaskedLM,
-        RBLNModelForQuestionAnswering,
-        RBLNModelForSequenceClassification,
     )
-    from .modeling_alias import (
-        RBLNASTForAudioClassification,
-        RBLNBertForQuestionAnswering,
-        RBLNResNetForImageClassification,
-        RBLNRobertaForMaskedLM,
-        RBLNRobertaForSequenceClassification,
-        RBLNT5ForConditionalGeneration,
-        RBLNViTForImageClassification,
-        RBLNXLMRobertaForSequenceClassification,
+    from .modeling_config import (
+        RBLNCompileConfig,
+        RBLNConfig,
     )
-    from .modeling_config import RBLNCompileConfig, RBLNConfig
-    from .modeling_diffusers import RBLNDiffusionMixin
     from .transformers import (
+        RBLNASTForAudioClassification,
         RBLNAutoModel,
         RBLNAutoModelForAudioClassification,
         RBLNAutoModelForCausalLM,
@@ -165,10 +150,12 @@ if TYPE_CHECKING:
         RBLNAutoModelForVision2Seq,
         RBLNBartForConditionalGeneration,
         RBLNBartModel,
+        RBLNBertForQuestionAnswering,
         RBLNBertModel,
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
         RBLNCLIPVisionModel,
+        RBLNDistilBertForQuestionAnswering,
         RBLNDPTForDepthEstimation,
         RBLNExaoneForCausalLM,
         RBLNGemmaForCausalLM,
@@ -179,12 +166,18 @@ if TYPE_CHECKING:
         RBLNMistralForCausalLM,
         RBLNPhiForCausalLM,
         RBLNQwen2ForCausalLM,
+        RBLNResNetForImageClassification,
+        RBLNRobertaForMaskedLM,
+        RBLNRobertaForSequenceClassification,
         RBLNT5EncoderModel,
         RBLNT5ForConditionalGeneration,
+        RBLNViTForImageClassification,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
+        RBLNXLMRobertaForSequenceClassification,
         RBLNXLMRobertaModel,
     )
+
 else:
     import sys
 

optimum/rbln/__version__.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.15'
-__version_tuple__ = version_tuple = (0, 1, 15)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

optimum/rbln/diffusers/__init__.py

@@ -54,9 +54,13 @@ _import_structure = {
         "RBLNControlNetModel",
         "RBLNSD3Transformer2DModel",
     ],
+    "modeling_diffusers": [
+        "RBLNDiffusionMixin",
+    ],
 }
 
 if TYPE_CHECKING:
+    from .modeling_diffusers import RBLNDiffusionMixin
     from .models import (
         RBLNAutoencoderKL,
         RBLNControlNetModel,

optimum/rbln/{modeling_diffusers.py → diffusers/modeling_diffusers.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.
+
 import copy
 import importlib
 from os import PathLike
@@ -27,10 +28,13 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import torch
 
-from .modeling import RBLNModel
-from .modeling_config import RUNTIME_KEYWORDS, ContextRblnConfig, use_rbln_config
-from .utils.decorator_utils import remove_compile_time_kwargs
+from ..modeling import RBLNModel
+from ..modeling_config import RUNTIME_KEYWORDS, ContextRblnConfig, use_rbln_config
+from ..utils.decorator_utils import remove_compile_time_kwargs
+from ..utils.logging import get_logger
+
 
+logger = get_logger(__name__)
 
 if TYPE_CHECKING:
     from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
@@ -177,30 +181,39 @@ class RBLNDiffusionMixin:
                         f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
                     )
 
-                # Load submodule outside if runtime kwargs(e.g. device) is specified.
-                if submodule_config := rbln_config.get(submodule_name):
-                    if any(kwd in submodule_config for kwd in RUNTIME_KEYWORDS):
-                        if model_index_config is None:
-                            model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
-
-                        module_name, class_name = model_index_config[submodule_name]
-                        if module_name != "optimum.rbln":
-                            raise ValueError(
-                                f"Invalid module_name '{module_name}' found in model_index.json for "
-                                f"submodule '{submodule_name}'. "
-                                "Expected 'optimum.rbln'. Please check the model_index.json configuration."
-                            )
-                        submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), class_name)
-                        submodule = submodule_cls.from_pretrained(
-                            model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
-                        )
-                        kwargs[submodule_name] = submodule
+                submodule_config = rbln_config.get(submodule_name, {})
+
+                for key, value in rbln_config.items():
+                    if key in RUNTIME_KEYWORDS and key not in submodule_config:
+                        submodule_config[key] = value
+
+                if not any(kwd in submodule_config for kwd in RUNTIME_KEYWORDS):
+                    continue
+
+                if model_index_config is None:
+                    model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
+
+                module_name, class_name = model_index_config[submodule_name]
+                if module_name != "optimum.rbln":
+                    raise ValueError(
+                        f"Invalid module_name '{module_name}' found in model_index.json for "
+                        f"submodule '{submodule_name}'. "
+                        "Expected 'optimum.rbln'. Please check the model_index.json configuration."
+                    )
+
+                submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), class_name)
+
+                submodule = submodule_cls.from_pretrained(
+                    model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
+                )
+                kwargs[submodule_name] = submodule
 
         with ContextRblnConfig(
             device=rbln_config.get("device"),
             device_map=rbln_config.get("device_map"),
             create_runtimes=rbln_config.get("create_runtimes"),
             optimize_host_mem=rbln_config.get("optimize_host_memory"),
+            activate_profiler=rbln_config.get("activate_profiler"),
         ):
             model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
 
@@ -266,8 +279,8 @@ class RBLNDiffusionMixin:
         controlnet_rbln_config: Dict[str, Any],
     ):
         # Compile multiple ControlNet models for a MultiControlNet setup
-        from .diffusers.models.controlnet import RBLNControlNetModel
-        from .diffusers.pipelines.controlnet import RBLNMultiControlNetModel
+        from .models.controlnet import RBLNControlNetModel
+        from .pipelines.controlnet import RBLNMultiControlNetModel
 
         compiled_controlnets = [
             RBLNControlNetModel.from_model(
@@ -278,7 +291,7 @@ class RBLNDiffusionMixin:
             )
             for i, controlnet in enumerate(controlnets.nets)
         ]
-        return RBLNMultiControlNetModel(compiled_controlnets, config=controlnets.nets[0].config)
+        return RBLNMultiControlNetModel(compiled_controlnets)
 
     @classmethod
     def _construct_pipe(cls, model, submodules, model_save_dir, rbln_config):
@@ -324,6 +337,35 @@ class RBLNDiffusionMixin:
 
         return model
 
+    def get_compiled_image_size(self):
+        if hasattr(self, "vae"):
+            compiled_image_size = self.vae.image_size
+        else:
+            compiled_image_size = None
+        return compiled_image_size
+
+    def handle_additional_kwargs(self, **kwargs):
+        """
+        Function to handle additional compile-time parameters during inference.
+
+        If the additional variable is determined by another module, this method should be overrided.
+
+        Example:
+            ```python
+            if hasattr(self, "movq"):
+                compiled_image_size = self.movq.image_size
+                kwargs["height"] = compiled_image_size[0]
+                kwargs["width"] = compiled_image_size[1]
+
+            compiled_num_frames = self.unet.rbln_config.model_cfg.get("num_frames", None)
+            if compiled_num_frames is not None:
+                kwargs["num_frames"] = self.unet.rbln_config.model_cfg.get("num_frames")
+            return kwargs
+            ```
+        """
+        return kwargs
+
     @remove_compile_time_kwargs
     def __call__(self, *args, **kwargs):
+        kwargs = self.handle_additional_kwargs(**kwargs)
         return super().__call__(*args, **kwargs)

optimum/rbln/diffusers/models/__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
@@ -35,6 +36,7 @@ _import_structure = {
     "controlnet": ["RBLNControlNetModel"],
     "transformers": ["RBLNSD3Transformer2DModel"],
 }
+
 if TYPE_CHECKING:
     from .autoencoders import (
         RBLNAutoencoderKL,

optimum/rbln/diffusers/models/autoencoders/autoencoder_kl.py

@@ -22,7 +22,7 @@
 # from Rebellions Inc.
 
 import logging
-from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import rebel
 import torch  # noqa: I001
@@ -32,7 +32,7 @@ from transformers import PretrainedConfig
 
 from ....modeling import RBLNModel
 from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 from .vae import RBLNRuntimeVAEDecoder, RBLNRuntimeVAEEncoder, _VAEDecoder, _VAEEncoder
 
 
@@ -88,22 +88,35 @@ class RBLNAutoencoderKL(RBLNModel):
     @classmethod
     def get_vae_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"))
+        noise_module = getattr(pipe, "unet", None) or getattr(pipe, "transformer", None)
+        vae_scale_factor = (
+            pipe.vae_scale_factor
+            if hasattr(pipe, "vae_scale_factor")
+            else 2 ** (len(pipe.vae.config.block_out_channels) - 1)
+        )
+
+        if noise_module is None:
+            raise AttributeError(
+                "Cannot find noise processing or predicting module attributes. ex. U-Net, Transformer, ..."
+            )
+
         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"]:
-                sample_size = pipe.vae.config.sample_size
+                sample_size = noise_module.config.sample_size
             elif rbln_config["inpaint_pipeline"]:
-                sample_size = pipe.unet.config.sample_size * pipe.vae_scale_factor
+                sample_size = noise_module.config.sample_size * vae_scale_factor
             else:
                 # In case of text2img, sample size of vae decoder is determined by unet.
-                unet_sample_size = pipe.unet.config.sample_size
-                if isinstance(unet_sample_size, int):
-                    sample_size = unet_sample_size * pipe.vae_scale_factor
+                noise_module_sample_size = noise_module.config.sample_size
+                if isinstance(noise_module_sample_size, int):
+                    sample_size = noise_module_sample_size * vae_scale_factor
                 else:
                     sample_size = (
-                        unet_sample_size[0] * pipe.vae_scale_factor,
-                        unet_sample_size[1] * pipe.vae_scale_factor,
+                        noise_module_sample_size[0] * vae_scale_factor,
+                        noise_module_sample_size[1] * vae_scale_factor,
                     )
         else:
             sample_size = (image_size[0], image_size[1])
@@ -192,15 +205,28 @@ class RBLNAutoencoderKL(RBLNModel):
 
     @classmethod
     def _create_runtimes(
-        cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_device_map: Dict[str, int],
+        activate_profiler: Optional[bool] = None,
     ) -> List[rebel.Runtime]:
         if len(compiled_models) == 1:
+            if DEFAULT_COMPILED_MODEL_NAME not in rbln_device_map:
+                cls._raise_missing_compiled_file_error([DEFAULT_COMPILED_MODEL_NAME])
+
             device_val = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
-            return [compiled_models[0].create_runtime(tensor_type="pt", device=device_val)]
+            return [
+                compiled_models[0].create_runtime(
+                    tensor_type="pt", device=device_val, activate_profiler=activate_profiler
+                )
+            ]
+
+        if any(model_name not in rbln_device_map for model_name in ["encoder", "decoder"]):
+            cls._raise_missing_compiled_file_error(["encoder", "decoder"])
 
         device_vals = [rbln_device_map["encoder"], rbln_device_map["decoder"]]
         return [
-            compiled_model.create_runtime(tensor_type="pt", device=device_val)
+            compiled_model.create_runtime(tensor_type="pt", device=device_val, activate_profiler=activate_profiler)
             for compiled_model, device_val in zip(compiled_models, device_vals)
         ]
 

optimum/rbln/diffusers/models/autoencoders/vae.py

@@ -21,7 +21,6 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-
 import logging
 from typing import TYPE_CHECKING
 

optimum/rbln/diffusers/models/controlnet.py

@@ -31,7 +31,7 @@ from transformers import PretrainedConfig
 
 from ...modeling import RBLNModel
 from ...modeling_config import RBLNCompileConfig, RBLNConfig
-from ...modeling_diffusers import RBLNDiffusionMixin
+from ..modeling_diffusers import RBLNDiffusionMixin
 
 
 if TYPE_CHECKING:

optimum/rbln/diffusers/models/transformers/transformer_sd3.py

@@ -31,7 +31,7 @@ from transformers import PretrainedConfig
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 if TYPE_CHECKING:

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

@@ -31,7 +31,7 @@ from transformers import PretrainedConfig
 
 from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
-from ....modeling_diffusers import RBLNDiffusionMixin
+from ...modeling_diffusers import RBLNDiffusionMixin
 
 
 if TYPE_CHECKING:
@@ -265,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__,

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

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

@@ -30,7 +30,6 @@ import torch
 from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
 
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNConfig
 from ...models.controlnet import RBLNControlNetModel
 
 
@@ -47,7 +46,6 @@ class RBLNMultiControlNetModel(RBLNModel):
     def __init__(
         self,
         models: List[RBLNControlNetModel],
-        **kwargs,
     ):
         self.nets = models
         self.dtype = torch.float32
@@ -67,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):

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
 
@@ -45,6 +58,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(RBLNDiffusionMixin, StableDif
     original_class = StableDiffusionControlNetImg2ImgPipeline
     _submodules = ["text_encoder", "unet", "vae", "controlnet"]
 
+    # Almost copied from diffusers.pipelines.controlnet.pipeline_controlnet_img2img.py
     def check_inputs(
         self,
         prompt,
@@ -202,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
 
@@ -45,6 +58,7 @@ class RBLNStableDiffusionXLControlNetPipeline(RBLNDiffusionMixin, StableDiffusio
     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,
@@ -234,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__(