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

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

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

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

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

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

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 from diffusers import StableDiffusionPipeline
 
 from ...modeling_diffusers import RBLNDiffusionMixin
@@ -19,4 +20,4 @@ from ...modeling_diffusers import RBLNDiffusionMixin
 
 class RBLNStableDiffusionPipeline(RBLNDiffusionMixin, StableDiffusionPipeline):
     original_class = StableDiffusionPipeline
-    _submodules = ["text_encoder", "unet", "vae"]
+    _submodules = ["vae", "text_encoder", "unet"]

optimum/rbln/modeling.py

@@ -14,15 +14,16 @@
 
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 import rebel
 import torch
 from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
 from transformers import AutoConfig, PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutput
 
+from .configuration_utils import DEFAULT_COMPILED_MODEL_NAME, RBLNModelConfig
 from .modeling_base import RBLNBaseModel
-from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, use_rbln_config
 from .utils.logging import get_logger
 
 
@@ -48,6 +49,9 @@ class RBLNModel(RBLNBaseModel):
         ```
     """
 
+    output_class = None
+    output_key = "last_hidden_state"
+
     @classmethod
     def update_kwargs(cls, kwargs):
         """
@@ -56,12 +60,7 @@ class RBLNModel(RBLNBaseModel):
         For example, `torchscript`=True should be set because torch.jit
         does not support `transformers` output instances as module output;
         """
-        kwargs.update(
-            {
-                "torchscript": True,
-                "return_dict": False,
-            }
-        )
+        kwargs.update({"torchscript": True})
         return kwargs
 
     @classmethod
@@ -70,7 +69,7 @@ class RBLNModel(RBLNBaseModel):
         model: "PreTrainedModel",
         save_dir_path: Path,
         subfolder: str,
-        rbln_config: RBLNConfig,
+        rbln_config: RBLNModelConfig,
     ):
         """
         If you are unavoidably running on a CPU rather than an RBLN device,
@@ -78,30 +77,29 @@ class RBLNModel(RBLNBaseModel):
         """
 
     @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:
         # Wrap the model if needed.
         return model
 
     @classmethod
-    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNConfig):
+    def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNModelConfig):
         model = cls.wrap_model_if_needed(model, rbln_config)
         rbln_compile_config = rbln_config.compile_cfgs[0]
         compiled_model = cls.compile(model, rbln_compile_config=rbln_compile_config)
         return compiled_model
 
     @classmethod
-    @use_rbln_config
     def from_model(
         cls,
         model: "PreTrainedModel",
         config: Optional[PretrainedConfig] = None,
-        rbln_config: Dict[str, Any] = {},
+        rbln_config: Optional[RBLNModelConfig] = None,
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         subfolder: str = "",
         **kwargs,
     ):
         preprocessors = kwargs.pop("preprocessors", [])
-        rbln_kwargs = rbln_config
+        rbln_config, kwargs = cls.prepare_rbln_config(rbln_config=rbln_config, **kwargs)
 
         # Directory to save compile artifacts(.rbln) and original configs
         if model_save_dir is None:
@@ -141,14 +139,21 @@ class RBLNModel(RBLNBaseModel):
         for preprocessor in preprocessors:
             preprocessor.save_pretrained(save_dir_path / subfolder)
 
-        # ad-hoc
-        rbln_kwargs["n_model_params"] = sum(p.numel() for p in model.parameters())
+        # Load submodules
+        if len(cls._rbln_submodules) > 0:
+            rbln_submodules = cls._load_submodules(
+                model=model,
+                model_save_dir=save_dir,
+                rbln_config=rbln_config,
+                **kwargs,
+            )
+        else:
+            rbln_submodules = []
 
         # Get compilation arguments (e.g. input_info)
-        rbln_config: RBLNConfig = cls.get_rbln_config(
-            preprocessors=preprocessors, model_config=config, rbln_kwargs=rbln_kwargs
+        rbln_config: RBLNModelConfig = cls.update_rbln_config(
+            preprocessors=preprocessors, model=model, model_config=config, rbln_config=rbln_config
         )
-        # rbln_config.update_runtime_cfg(rbln_kwargs) # This is done in get_rbln_config
 
         compiled_model: Union[rebel.RBLNCompiledModel, Dict[str, rebel.RBLNCompiledModel]] = cls.get_compiled_model(
             model, rbln_config=rbln_config
@@ -167,17 +172,6 @@ class RBLNModel(RBLNBaseModel):
         # Save torch artifacts (e.g. embedding matrix if needed.)
         cls.save_torch_artifacts(model, save_dir_path=save_dir_path, subfolder=subfolder, rbln_config=rbln_config)
 
-        # Load submodules
-        if len(cls._rbln_submodules) > 0:
-            rbln_submodules = cls._load_submodules(
-                model=model,
-                model_save_dir=save_dir,
-                rbln_kwargs=rbln_kwargs,
-                **kwargs,
-            )
-        else:
-            rbln_submodules = []
-
         # Instantiate
         return cls._from_pretrained(
             model_id=save_dir_path,
@@ -201,8 +195,8 @@ class RBLNModel(RBLNBaseModel):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
-        # Some rbln-kwargs should be applied before loading torch module (i.e. quantized llm)
-        rbln_kwargs: Optional[Dict[str, Any]] = None,
+        # Some rbln-config should be applied before loading torch module (i.e. quantized llm)
+        rbln_config: Optional[RBLNModelConfig] = None,
         **kwargs,
     ) -> "PreTrainedModel":
         kwargs = cls.update_kwargs(kwargs)
@@ -222,18 +216,43 @@ class RBLNModel(RBLNBaseModel):
     def _create_runtimes(
         cls,
         compiled_models: List[rebel.RBLNCompiledModel],
-        rbln_device_map: Dict[str, int],
-        activate_profiler: Optional[bool] = None,
+        rbln_config: RBLNModelConfig,
     ) -> List[rebel.Runtime]:
-        if DEFAULT_COMPILED_MODEL_NAME not in rbln_device_map:
+        if DEFAULT_COMPILED_MODEL_NAME not in rbln_config.device_map:
             cls._raise_missing_compiled_file_error([DEFAULT_COMPILED_MODEL_NAME])
 
-        device = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
         return [
-            compiled_model.create_runtime(tensor_type="pt", device=device, activate_profiler=activate_profiler)
+            rebel.Runtime(
+                compiled_model,
+                tensor_type="pt",
+                device=rbln_config.device_map[DEFAULT_COMPILED_MODEL_NAME],
+                activate_profiler=rbln_config.activate_profiler,
+            )
             for compiled_model in compiled_models
         ]
 
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
+    def forward(self, *args, return_dict: Optional[bool] = None, **kwargs):
+        if self.hf_library_name == "transformers":
+            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        else:
+            return_dict = True if return_dict is None else return_dict
+
+        # Get output from the model
         output = self.model[0](*args, **kwargs)
-        return output
+
+        # Format output according to task requirements
+        return self._prepare_output(output, return_dict)
+
+    def _prepare_output(self, output, return_dict):
+        """
+        Prepare model output based on return_dict flag.
+        This method can be overridden by subclasses to provide task-specific output handling.
+        """
+        if not return_dict:
+            return (output,) if not isinstance(output, (tuple, list)) else output
+        else:
+            if self.output_class is None:
+                return BaseModelOutput(last_hidden_state=output)
+
+            # Create output with the appropriate class and key
+            return self.output_class(**{self.output_key: output})