optimum-rbln 0.7.4a4__py3-none-any.whl → 0.7.4a6__py3-none-any.whl

optimum/rbln/diffusers/models/controlnet.py

@@ -13,20 +13,22 @@
 # limitations under the License.
 
 import importlib
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Dict, Optional, Union
 
 import torch
 from diffusers import ControlNetModel
+from diffusers.models.controlnet import ControlNetOutput
 from transformers import PretrainedConfig
 
+from ...configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ...modeling import RBLNModel
-from ...modeling_config import RBLNCompileConfig, RBLNConfig
 from ...utils.logging import get_logger
-from ..modeling_diffusers import RBLNDiffusionMixin
+from ..configurations import RBLNControlNetModelConfig
+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__)
@@ -98,6 +100,7 @@ class _ControlNetModel_Cross_Attention(torch.nn.Module):
 class RBLNControlNetModel(RBLNModel):
     hf_library_name = "diffusers"
     auto_model_class = ControlNetModel
+    output_class = ControlNetOutput
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
@@ -106,7 +109,7 @@ class RBLNControlNetModel(RBLNModel):
         )
 
     @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: RBLNModelConfig) -> torch.nn.Module:
         use_encoder_hidden_states = False
         for down_block in model.down_blocks:
             if use_encoder_hidden_states := getattr(down_block, "has_cross_attention", False):
@@ -118,73 +121,50 @@ class RBLNControlNetModel(RBLNModel):
             return _ControlNetModel(model).eval()
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+    def update_rbln_config_using_pipe(
+        cls,
+        pipe: RBLNDiffusionMixin,
+        rbln_config: "RBLNDiffusionMixinConfig",
+        submodule_name: str,
+    ) -> "RBLNDiffusionMixinConfig":
         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,
-            }
+        rbln_config.controlnet.max_seq_len = pipe.text_encoder.config.max_position_embeddings
+        text_model_hidden_size = pipe.text_encoder_2.config.hidden_size if hasattr(pipe, "text_encoder_2") else None
+        rbln_config.controlnet.text_model_hidden_size = text_model_hidden_size
+        rbln_config.controlnet.vae_sample_size = rbln_vae_cls.get_vae_sample_size(pipe, rbln_config.vae)
+        rbln_config.controlnet.unet_sample_size = rbln_unet_cls.get_unet_sample_size(
+            pipe, rbln_config.unet, image_size=rbln_config.image_size
         )
 
         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")
-        unet_sample_size = rbln_kwargs.get("unet_sample_size")
-        vae_sample_size = rbln_kwargs.get("vae_sample_size")
-
-        if batch_size is None:
-            batch_size = 1
-
-        if unet_sample_size is None:
-            raise ValueError(
-                "`rbln_unet_sample_size` (latent height, widht) must be specified (ex. unet's sample_size)"
-            )
+        rbln_config: RBLNControlNetModelConfig,
+    ) -> RBLNModelConfig:
+        if rbln_config.unet_sample_size is None:
+            raise ValueError("`unet_sample_size` (latent height, width) must be specified (ex. unet's sample_size)")
 
-        if vae_sample_size is None:
-            raise ValueError(
-                "`rbln_vae_sample_size` (input image height, width) must be specified (ex. vae's sample_size)"
-            )
+        if rbln_config.vae_sample_size is None:
+            raise ValueError("`vae_sample_size` (input image height, width) must be specified (ex. vae's sample_size)")
 
-        if max_seq_len is None:
-            raise ValueError("`rbln_max_seq_len` (ex. text_encoder's max_position_embeddings )must be specified")
+        if rbln_config.max_seq_len is None:
+            raise ValueError("`max_seq_len` (ex. text_encoder's max_position_embeddings) must be specified")
 
         input_info = [
             (
                 "sample",
                 [
-                    batch_size,
+                    rbln_config.batch_size,
                     model_config.in_channels,
-                    unet_sample_size[0],
-                    unet_sample_size[1],
+                    rbln_config.unet_sample_size[0],
+                    rbln_config.unet_sample_size[1],
                 ],
                 "float32",
             ),
@@ -196,7 +176,7 @@ class RBLNControlNetModel(RBLNModel):
             input_info.append(
                 (
                     "encoder_hidden_states",
-                    [batch_size, max_seq_len, model_config.cross_attention_dim],
+                    [rbln_config.batch_size, rbln_config.max_seq_len, model_config.cross_attention_dim],
                     "float32",
                 )
             )
@@ -204,25 +184,18 @@ class RBLNControlNetModel(RBLNModel):
         input_info.append(
             (
                 "controlnet_cond",
-                [batch_size, 3, vae_sample_size[0], vae_sample_size[1]],
+                [rbln_config.batch_size, 3, rbln_config.vae_sample_size[0], rbln_config.vae_sample_size[1]],
                 "float32",
             )
         )
         input_info.append(("conditioning_scale", [], "float32"))
 
         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"]
-            input_info.append(("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32"))
-            input_info.append(("time_ids", [batch_size, 6], "float32"))
+            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"))
 
         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.set_compile_cfgs([rbln_compile_config])
         return rbln_config
 
     @property
@@ -237,6 +210,7 @@ class RBLNControlNetModel(RBLNModel):
         controlnet_cond: torch.FloatTensor,
         conditioning_scale: torch.Tensor = 1.0,
         added_cond_kwargs: Dict[str, torch.Tensor] = {},
+        return_dict: bool = True,
         **kwargs,
     ):
         sample_batch_size = sample.size()[0]
@@ -246,14 +220,14 @@ class RBLNControlNetModel(RBLNModel):
         ):
             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"
+                "This may be caused by the 'guidance_scale' parameter, which doubles the runtime batch size of ControlNet in Stable Diffusion. "
+                "Adjust the batch size of ControlNet during compilation to match the runtime 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(
+            output = self.model[0](
                 sample.contiguous(),
                 timestep.float(),
                 encoder_hidden_states,
@@ -262,14 +236,25 @@ class RBLNControlNetModel(RBLNModel):
                 **added_cond_kwargs,
             )
         else:
-            output = super().forward(
+            output = self.model[0](
                 sample.contiguous(),
                 timestep.float(),
                 controlnet_cond,
                 torch.tensor(conditioning_scale),
                 **added_cond_kwargs,
             )
+
         down_block_res_samples = output[:-1]
         mid_block_res_sample = output[-1]
+        output = (down_block_res_samples, mid_block_res_sample)
+        output = self._prepare_output(output, return_dict)
+        return output
 
-        return down_block_res_samples, mid_block_res_sample
+    def _prepare_output(self, output, return_dict):
+        if not return_dict:
+            return (output,) if not isinstance(output, (tuple, list)) else output
+        else:
+            return ControlNetOutput(
+                down_block_res_samples=output[:-1],
+                mid_block_res_sample=output[-1],
+            )

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

@@ -13,37 +13,22 @@
 # limitations under the License.
 
 from pathlib import Path
-from typing import Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import torch
 from diffusers.models.transformers.prior_transformer import PriorTransformer, PriorTransformerOutput
-from transformers import PretrainedConfig, PreTrainedModel
 
+from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
-from ....utils.runtime_utils import RBLNPytorchRuntime
-from ...modeling_diffusers import RBLNDiffusionMixin
+from ...configurations.models import RBLNPriorTransformerConfig
+from ...modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
 
-logger = get_logger(__name__)
-
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
 
-class RBLNRuntimePriorTransformer(RBLNPytorchRuntime):
-    def forward(
-        self, hidden_states, timestep, proj_embedding, encoder_hidden_states, attention_mask, return_dict: bool = True
-    ):
-        predicted_image_embedding = super().forward(
-            hidden_states,
-            timestep,
-            proj_embedding,
-            encoder_hidden_states,
-            attention_mask,
-        )
-        if return_dict:
-            return PriorTransformerOutput(predicted_image_embedding=predicted_image_embedding)
-        else:
-            return (predicted_image_embedding,)
+logger = get_logger(__name__)
 
 
 class _PriorTransformer(torch.nn.Module):
@@ -73,51 +58,28 @@ class _PriorTransformer(torch.nn.Module):
 class RBLNPriorTransformer(RBLNModel):
     hf_library_name = "diffusers"
     auto_model_class = PriorTransformer
+    output_class = PriorTransformerOutput
+    output_key = "predicted_image_embedding"
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
-        self.runtime = RBLNRuntimePriorTransformer(runtime=self.model[0])
         artifacts = torch.load(self.model_save_dir / self.subfolder / "torch_artifacts.pth", weights_only=False)
         self.clip_mean = artifacts["clip_mean"]
         self.clip_std = artifacts["clip_std"]
 
     @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: RBLNModelConfig) -> torch.nn.Module:
         return _PriorTransformer(model).eval()
 
     @classmethod
-    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        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."
-                )
-        embedding_dim = rbln_config.get("embedding_dim", pipe.prior.config.embedding_dim)
-        num_embeddings = rbln_config.get("num_embeddings", pipe.prior.config.num_embeddings)
-
-        rbln_config.update(
-            {
-                "batch_size": batch_size,
-                "embedding_dim": embedding_dim,
-                "num_embeddings": num_embeddings,
-            }
-        )
-
+    def update_rbln_config_using_pipe(
+        cls, pipe: RBLNDiffusionMixin, rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
         return rbln_config
 
     @classmethod
     def save_torch_artifacts(
-        cls,
-        model: "PreTrainedModel",
-        save_dir_path: Path,
-        subfolder: str,
-        rbln_config: RBLNConfig,
+        cls, model: "PreTrainedModel", save_dir_path: Path, subfolder: str, rbln_config: RBLNModelConfig
     ):
         save_dict = {}
         save_dict["clip_mean"] = model.clip_mean
@@ -125,50 +87,51 @@ class RBLNPriorTransformer(RBLNModel):
         torch.save(save_dict, save_dir_path / subfolder / "torch_artifacts.pth")
 
     @classmethod
-    def _get_rbln_config(
+    def _update_rbln_config(
         cls,
-        preprocessors,
-        model_config: PretrainedConfig,
-        rbln_kwargs,
-    ) -> RBLNConfig:
-        batch_size = rbln_kwargs.get("batch_size") or 1
-        embedding_dim = rbln_kwargs.get("embedding_dim") or model_config.embedding_dim
-        num_embeddings = rbln_kwargs.get("num_embeddings") or model_config.num_embeddings
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model: "PreTrainedModel",
+        model_config: "PretrainedConfig",
+        rbln_config: RBLNPriorTransformerConfig,
+    ) -> RBLNPriorTransformerConfig:
+        rbln_config.embedding_dim = rbln_config.embedding_dim or model_config.embedding_dim
+        rbln_config.num_embeddings = rbln_config.num_embeddings or model_config.num_embeddings
 
         input_info = [
-            ("hidden_states", [batch_size, embedding_dim], "float32"),
+            ("hidden_states", [rbln_config.batch_size, rbln_config.embedding_dim], "float32"),
             ("timestep", [], "float32"),
-            ("proj_embedding", [batch_size, embedding_dim], "float32"),
-            ("encoder_hidden_states", [batch_size, num_embeddings, embedding_dim], "float32"),
-            ("attention_mask", [batch_size, num_embeddings], "float32"),
+            ("proj_embedding", [rbln_config.batch_size, rbln_config.embedding_dim], "float32"),
+            (
+                "encoder_hidden_states",
+                [rbln_config.batch_size, rbln_config.num_embeddings, rbln_config.embedding_dim],
+                "float32",
+            ),
+            ("attention_mask", [rbln_config.batch_size, rbln_config.num_embeddings], "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.set_compile_cfgs([rbln_compile_config])
         return rbln_config
 
+    def post_process_latents(self, prior_latents):
+        prior_latents = (prior_latents * self.clip_std) + self.clip_mean
+        return prior_latents
+
     def forward(
         self,
         hidden_states,
         timestep: Union[torch.Tensor, float, int],
         proj_embedding: torch.Tensor,
         encoder_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.BoolTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
         return_dict: bool = True,
     ):
-        return self.runtime.forward(
-            hidden_states.contiguous(),
+        # Convert timestep(long) and attention_mask(bool) to float
+        return super().forward(
+            hidden_states,
             timestep.float(),
             proj_embedding,
             encoder_hidden_states,
             attention_mask.float(),
-            return_dict,
+            return_dict=return_dict,
         )
-
-    def post_process_latents(self, prior_latents):
-        prior_latents = (prior_latents * self.clip_std) + self.clip_mean
-        return prior_latents

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

@@ -19,14 +19,16 @@ from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
 from transformers import PretrainedConfig
 
+from ....configuration_utils import RBLNCompileConfig, RBLNModelConfig
 from ....modeling import RBLNModel
-from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
-from ...modeling_diffusers import RBLNDiffusionMixin
+from ...configurations import RBLNSD3Transformer2DModelConfig
 
 
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+
+    from ...modeling_diffusers import RBLNDiffusionMixin, RBLNDiffusionMixinConfig
 
 logger = get_logger(__name__)
 
@@ -58,84 +60,64 @@ class SD3Transformer2DModelWrapper(torch.nn.Module):
 
 class RBLNSD3Transformer2DModel(RBLNModel):
     hf_library_name = "diffusers"
+    auto_model_class = SD3Transformer2DModel
+    output_class = Transformer2DModelOutput
+    output_key = "sample"
 
     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:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNModelConfig) -> 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."
+    def update_rbln_config_using_pipe(
+        cls, pipe: "RBLNDiffusionMixin", rbln_config: "RBLNDiffusionMixinConfig", submodule_name: str
+    ) -> "RBLNDiffusionMixinConfig":
+        if rbln_config.sample_size is None:
+            if rbln_config.image_size is not None:
+                rbln_config.transformer.sample_size = (
+                    rbln_config.image_size[0] // pipe.vae_scale_factor,
+                    rbln_config.image_size[1] // pipe.vae_scale_factor,
                 )
+            else:
+                rbln_config.transformer.sample_size = pipe.default_sample_size
 
-        rbln_config.update(
-            {
-                "batch_size": batch_size,
-                "prompt_embed_length": prompt_embed_length,
-                "sample_size": sample_size,
-            }
-        )
-
+        prompt_embed_length = pipe.tokenizer_max_length + rbln_config.max_seq_len
+        rbln_config.transformer.prompt_embed_length = prompt_embed_length
         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:
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
+        rbln_config: RBLNSD3Transformer2DModelConfig,
+    ) -> RBLNSD3Transformer2DModelConfig:
+        if rbln_config.sample_size is None:
+            rbln_config.sample_size = model_config.sample_size
 
-        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.")
+        if isinstance(rbln_config.sample_size, int):
+            rbln_config.sample_size = (rbln_config.sample_size, rbln_config.sample_size)
 
         input_info = [
             (
                 "hidden_states",
                 [
-                    rbln_batch_size,
+                    rbln_config.batch_size,
                     model_config.in_channels,
-                    sample_size[0],
-                    sample_size[1],
+                    rbln_config.sample_size[0],
+                    rbln_config.sample_size[1],
                 ],
                 "float32",
             ),
             (
                 "encoder_hidden_states",
                 [
-                    rbln_batch_size,
-                    rbln_prompt_embed_length,
+                    rbln_config.batch_size,
+                    rbln_config.prompt_embed_length,
                     model_config.joint_attention_dim,
                 ],
                 "float32",
@@ -143,24 +125,16 @@ class RBLNSD3Transformer2DModel(RBLNModel):
             (
                 "pooled_projections",
                 [
-                    rbln_batch_size,
+                    rbln_config.batch_size,
                     model_config.pooled_projection_dim,
                 ],
                 "float32",
             ),
-            ("timestep", [rbln_batch_size], "float32"),
+            ("timestep", [rbln_config.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})
-
+        compile_config = RBLNCompileConfig(input_info=input_info)
+        rbln_config.set_compile_cfgs([compile_config])
         return rbln_config
 
     @property
@@ -184,11 +158,12 @@ class RBLNSD3Transformer2DModel(RBLNModel):
             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}). "
+                f"Mismatch between transformer'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"
+                "Adjust the batch size of transformer during compilation.\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)
+        return super().forward(
+            hidden_states, encoder_hidden_states, pooled_projections, timestep, return_dict=return_dict
+        )