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

optimum/rbln/diffusers/models/{autoencoder_kl.py → autoencoders/autoencoder_kl.py}

@@ -22,19 +22,18 @@
 # from Rebellions Inc.
 
 import logging
-from typing import TYPE_CHECKING, Any, Dict, List, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import rebel
 import torch  # noqa: I001
 from diffusers import AutoencoderKL
-from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 from diffusers.models.modeling_outputs import AutoencoderKLOutput
 from transformers import PretrainedConfig
 
-from ...modeling_base import RBLNModel
-from ...modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
-from ...utils.context import override_auto_classes
-from ...utils.runtime_utils import RBLNPytorchRuntime
+from ....modeling import RBLNModel
+from ....modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from ...modeling_diffusers import RBLNDiffusionMixin
+from .vae import RBLNRuntimeVAEDecoder, RBLNRuntimeVAEEncoder, _VAEDecoder, _VAEEncoder
 
 
 if TYPE_CHECKING:
@@ -44,30 +43,22 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class RBLNRuntimeVAEEncoder(RBLNPytorchRuntime):
-    def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
-        moments = self.forward(x.contiguous())
-        posterior = DiagonalGaussianDistribution(moments)
-        return AutoencoderKLOutput(latent_dist=posterior)
-
-
-class RBLNRuntimeVAEDecoder(RBLNPytorchRuntime):
-    def decode(self, z: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
-        return (self.forward(z),)
-
-
 class RBLNAutoencoderKL(RBLNModel):
+    auto_model_class = AutoencoderKL
     config_name = "config.json"
+    hf_library_name = "diffusers"
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
 
-        if self.rbln_config.model_cfg.get("img2img_pipeline"):
+        if self.rbln_config.model_cfg.get("img2img_pipeline") or self.rbln_config.model_cfg.get("inpaint_pipeline"):
             self.encoder = RBLNRuntimeVAEEncoder(runtime=self.model[0], main_input_name="x")
             self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[1], main_input_name="z")
         else:
             self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[0], main_input_name="z")
 
+        self.image_size = self.rbln_config.model_cfg["sample_size"]
+
     @classmethod
     def get_compiled_model(cls, model, rbln_config: RBLNConfig):
         def compile_img2img():
@@ -89,16 +80,53 @@ class RBLNAutoencoderKL(RBLNModel):
 
             return dec_compiled_model
 
-        if rbln_config.model_cfg.get("img2img_pipeline"):
+        if rbln_config.model_cfg.get("img2img_pipeline") or rbln_config.model_cfg.get("inpaint_pipeline"):
             return compile_img2img()
         else:
             return compile_text2img()
 
     @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        with override_auto_classes(config_func=AutoencoderKL.load_config, model_func=AutoencoderKL.from_pretrained):
-            rt = super().from_pretrained(*args, **kwargs)
-        return rt
+    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 = noise_module.config.sample_size
+            elif rbln_config["inpaint_pipeline"]:
+                sample_size = noise_module.config.sample_size * vae_scale_factor
+            else:
+                # In case of text2img, sample size of vae decoder is determined by unet.
+                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 = (
+                        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])
+
+        return sample_size
+
+    @classmethod
+    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+        rbln_config.update({"sample_size": cls.get_vae_sample_size(pipe, rbln_config)})
+        return rbln_config
 
     @classmethod
     def _get_rbln_config(
@@ -109,6 +137,8 @@ class RBLNAutoencoderKL(RBLNModel):
     ) -> RBLNConfig:
         rbln_batch_size = rbln_kwargs.get("batch_size")
         sample_size = rbln_kwargs.get("sample_size")
+        is_img2img = rbln_kwargs.get("img2img_pipeline")
+        is_inpaint = rbln_kwargs.get("inpaint_pipeline")
 
         if rbln_batch_size is None:
             rbln_batch_size = 1
@@ -119,6 +149,8 @@ class RBLNAutoencoderKL(RBLNModel):
         if isinstance(sample_size, int):
             sample_size = (sample_size, sample_size)
 
+        rbln_kwargs["sample_size"] = sample_size
+
         if hasattr(model_config, "block_out_channels"):
             vae_scale_factor = 2 ** (len(model_config.block_out_channels) - 1)
         else:
@@ -128,7 +160,7 @@ class RBLNAutoencoderKL(RBLNModel):
         dec_shape = (sample_size[0] // vae_scale_factor, sample_size[1] // vae_scale_factor)
         enc_shape = (sample_size[0], sample_size[1])
 
-        if rbln_kwargs["img2img_pipeline"]:
+        if is_img2img or is_inpaint:
             vae_enc_input_info = [
                 (
                     "x",
@@ -173,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)
         ]
 
@@ -191,36 +236,3 @@ class RBLNAutoencoderKL(RBLNModel):
 
     def decode(self, z: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
         return self.decoder.decode(z)
-
-
-class _VAEDecoder(torch.nn.Module):
-    def __init__(self, vae: "AutoencoderKL"):
-        super().__init__()
-        self.vae = vae
-
-    def forward(self, z):
-        vae_out = self.vae.decode(z, return_dict=False)
-        return vae_out
-
-
-class _VAEEncoder(torch.nn.Module):
-    def __init__(self, vae: "AutoencoderKL"):
-        super().__init__()
-        self.vae = vae
-
-    def encode(self, x: torch.FloatTensor, return_dict: bool = True):
-        if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
-            return self.tiled_encode(x, return_dict=return_dict)
-
-        if self.use_slicing and x.shape[0] > 1:
-            encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
-            h = torch.cat(encoded_slices)
-        else:
-            h = self.encoder(x)
-
-        moments = self.quant_conv(h)
-        return moments
-
-    def forward(self, x):
-        vae_out = _VAEEncoder.encode(self.vae, x, return_dict=False)
-        return vae_out

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

@@ -0,0 +1,83 @@
+# 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.
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch  # noqa: I001
+from diffusers import AutoencoderKL
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+
+from ....utils.runtime_utils import RBLNPytorchRuntime
+
+
+if TYPE_CHECKING:
+    import torch
+
+logger = logging.getLogger(__name__)
+
+
+class RBLNRuntimeVAEEncoder(RBLNPytorchRuntime):
+    def encode(self, x: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        moments = self.forward(x.contiguous())
+        posterior = DiagonalGaussianDistribution(moments)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+
+class RBLNRuntimeVAEDecoder(RBLNPytorchRuntime):
+    def decode(self, z: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
+        return (self.forward(z),)
+
+
+class _VAEDecoder(torch.nn.Module):
+    def __init__(self, vae: "AutoencoderKL"):
+        super().__init__()
+        self.vae = vae
+
+    def forward(self, z):
+        vae_out = self.vae.decode(z, return_dict=False)
+        return vae_out
+
+
+class _VAEEncoder(torch.nn.Module):
+    def __init__(self, vae: "AutoencoderKL"):
+        super().__init__()
+        self.vae = vae
+
+    def encode(self, x: torch.FloatTensor, return_dict: bool = True):
+        if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
+            return self.tiled_encode(x, return_dict=return_dict)
+
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self.encoder(x)
+            if self.quant_conv is not None:
+                h = self.quant_conv(h)
+        return h
+
+    def forward(self, x):
+        vae_out = _VAEEncoder.encode(self.vae, x, return_dict=False)
+        return vae_out

optimum/rbln/diffusers/models/controlnet.py

@@ -21,6 +21,7 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
+import importlib
 import logging
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
@@ -28,9 +29,9 @@ import torch
 from diffusers import ControlNetModel
 from transformers import PretrainedConfig
 
-from ...modeling_base import RBLNModel
+from ...modeling import RBLNModel
 from ...modeling_config import RBLNCompileConfig, RBLNConfig
-from ...utils.context import override_auto_classes
+from ..modeling_diffusers import RBLNDiffusionMixin
 
 
 if TYPE_CHECKING:
@@ -104,21 +105,15 @@ class _ControlNetModel_Cross_Attention(torch.nn.Module):
 
 
 class RBLNControlNetModel(RBLNModel):
+    hf_library_name = "diffusers"
+    auto_model_class = ControlNetModel
+
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
         self.use_encoder_hidden_states = any(
             item[0] == "encoder_hidden_states" for item in self.rbln_config.compile_cfgs[0].input_info
         )
 
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        with override_auto_classes(
-            config_func=ControlNetModel.load_config,
-            model_func=ControlNetModel.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:
         use_encoder_hidden_states = False
@@ -131,6 +126,38 @@ class RBLNControlNetModel(RBLNModel):
         else:
             return _ControlNetModel(model).eval()
 
+    @classmethod
+    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+        rbln_vae_cls = getattr(importlib.import_module("optimum.rbln"), f"RBLN{pipe.vae.__class__.__name__}")
+        rbln_unet_cls = getattr(importlib.import_module("optimum.rbln"), f"RBLN{pipe.unet.__class__.__name__}")
+        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,
+                "vae_sample_size": rbln_vae_cls.get_vae_sample_size(pipe, rbln_config),
+                "unet_sample_size": rbln_unet_cls.get_unet_sample_size(pipe, rbln_config),
+                "batch_size": batch_size,
+            }
+        )
+
+        return rbln_config
+
     @classmethod
     def _get_rbln_config(
         cls,
@@ -207,6 +234,10 @@ class RBLNControlNetModel(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.FloatTensor,
@@ -217,9 +248,18 @@ class RBLNControlNetModel(RBLNModel):
         added_cond_kwargs: Dict[str, torch.Tensor] = {},
         **kwargs,
     ):
-        """
-        The [`ControlNetModel`] forward method.
-        """
+        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 ControlNet'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 self.use_encoder_hidden_states:
             output = super().forward(

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

@@ -0,0 +1,24 @@
+# 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 .transformer_sd3 import RBLNSD3Transformer2DModel

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

@@ -0,0 +1,203 @@
+# 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.
+
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import torch
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
+from transformers import PretrainedConfig
+
+from ....modeling import RBLNModel
+from ....modeling_config import RBLNCompileConfig, RBLNConfig
+from ...modeling_diffusers import RBLNDiffusionMixin
+
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class SD3Transformer2DModelWrapper(torch.nn.Module):
+    def __init__(self, model: "SD3Transformer2DModel") -> None:
+        super().__init__()
+        self.model = model
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        pooled_projections: torch.FloatTensor = None,
+        timestep: torch.LongTensor = None,
+        # need controlnet support?
+        block_controlnet_hidden_states: List = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+    ):
+        return self.model(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            pooled_projections=pooled_projections,
+            timestep=timestep,
+            return_dict=False,
+        )
+
+
+class RBLNSD3Transformer2DModel(RBLNModel):
+    hf_library_name = "diffusers"
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
+        return SD3Transformer2DModelWrapper(model).eval()
+
+    @classmethod
+    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+        sample_size = rbln_config.get("sample_size", pipe.default_sample_size)
+        img_width = rbln_config.get("img_width")
+        img_height = rbln_config.get("img_height")
+
+        if (img_width is None) ^ (img_height is None):
+            raise RuntimeError
+
+        elif img_width and img_height:
+            sample_size = img_height // pipe.vae_scale_factor, img_width // pipe.vae_scale_factor
+
+        prompt_max_length = rbln_config.get("max_sequence_length", 256)
+        prompt_embed_length = pipe.tokenizer_max_length + prompt_max_length
+
+        batch_size = rbln_config.get("batch_size")
+        if not batch_size:
+            do_classifier_free_guidance = rbln_config.get("guidance_scale", 5.0) > 1.0
+            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(
+            {
+                "batch_size": batch_size,
+                "prompt_embed_length": prompt_embed_length,
+                "sample_size": sample_size,
+            }
+        )
+
+        return rbln_config
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model_config: "PretrainedConfig",
+        rbln_kwargs: Dict[str, Any] = {},
+    ) -> RBLNConfig:
+        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+
+        sample_size = rbln_kwargs.get("sample_size", model_config.sample_size)
+        if isinstance(sample_size, int):
+            sample_size = (sample_size, sample_size)
+
+        rbln_prompt_embed_length = rbln_kwargs.get("prompt_embed_length")
+        if rbln_prompt_embed_length is None:
+            raise ValueError("rbln_prompt_embed_length should be specified.")
+
+        input_info = [
+            (
+                "hidden_states",
+                [
+                    rbln_batch_size,
+                    model_config.in_channels,
+                    sample_size[0],
+                    sample_size[1],
+                ],
+                "float32",
+            ),
+            (
+                "encoder_hidden_states",
+                [
+                    rbln_batch_size,
+                    rbln_prompt_embed_length,
+                    model_config.joint_attention_dim,
+                ],
+                "float32",
+            ),
+            (
+                "pooled_projections",
+                [
+                    rbln_batch_size,
+                    model_config.pooled_projection_dim,
+                ],
+                "float32",
+            ),
+            ("timestep", [rbln_batch_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,
+        )
+
+        rbln_config.model_cfg.update({"batch_size": rbln_batch_size})
+
+        return rbln_config
+
+    @property
+    def compiled_batch_size(self):
+        return self.rbln_config.compile_cfgs[0].input_info[0][1][0]
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        pooled_projections: torch.FloatTensor = None,
+        timestep: torch.LongTensor = None,
+        block_controlnet_hidden_states: List = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        sample_batch_size = hidden_states.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 Transformers' 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"
+            )
+
+        sample = super().forward(hidden_states, encoder_hidden_states, pooled_projections, timestep)
+        return Transformer2DModelOutput(sample=sample)

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

@@ -0,0 +1,24 @@
+# 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 .unet_2d_condition import RBLNUNet2DConditionModel