optimum-rbln 0.7.5rc0__tar.gz → 0.7.5rc2__tar.gz

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: optimum-rbln
-Version: 0.7.5rc0
+Version: 0.7.5rc2
 Summary: Optimum RBLN is the interface between the HuggingFace Transformers and Diffusers libraries and RBLN accelerators. It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.5rc0'
-__version_tuple__ = version_tuple = (0, 7, 5, 'rc0')
+__version__ = version = '0.7.5rc2'
+__version_tuple__ = version_tuple = (0, 7, 5, 'rc2')

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_controlnet.py

@@ -33,6 +33,8 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         img_height: Optional[int] = None,
         img_width: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
@@ -51,6 +53,8 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
+            height (Optional[int]): Height of the generated images.
+            width (Optional[int]): Width of the generated images.
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
@@ -65,11 +69,29 @@ class _RBLNStableDiffusionControlNetPipelineBaseConfig(RBLNModelConfig):
             accommodate classifier-free guidance.
         """
         super().__init__(**kwargs)
-        if image_size is not None and (img_height is not None or img_width is not None):
-            raise ValueError("image_size and img_height/img_width cannot both be provided")
 
-        if img_height is not None and img_width is not None:
+        # Initial check for image_size conflict remains as is
+        if image_size is not None and (
+            img_height is not None or img_width is not None or height is not None or width is not None
+        ):
+            raise ValueError("image_size cannot be provided alongside img_height/img_width or height/width")
+
+        # Prioritize height/width (HF-aligned)
+        if height is not None and width is not None:
+            if img_height is not None or img_width is not None:
+                # Raise error if both sets of arguments are provided
+                raise ValueError(
+                    "Cannot provide both 'height'/'width' and 'img_height'/'img_width' simultaneously. "
+                    "Please use one set of arguments for image dimensions, preferring 'height'/'width'."
+                )
+            image_size = (height, width)
+        elif (height is not None and width is None) or (height is None and width is not None):
+            raise ValueError("Both height and width must be provided together if used")
+        # Fallback to img_height/img_width for backward compatibility
+        elif img_height is not None and img_width is not None:
             image_size = (img_height, img_width)
+        elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
+            raise ValueError("Both img_height and img_width must be provided together if used")
 
         self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
         self.unet = self.init_submodule_config(
@@ -139,6 +161,8 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         img_height: Optional[int] = None,
         img_width: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
@@ -159,6 +183,8 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
+            height (Optional[int]): Height of the generated images.
+            width (Optional[int]): Width of the generated images.
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
@@ -173,11 +199,29 @@ class _RBLNStableDiffusionXLControlNetPipelineBaseConfig(RBLNModelConfig):
             accommodate classifier-free guidance.
         """
         super().__init__(**kwargs)
-        if image_size is not None and (img_height is not None or img_width is not None):
-            raise ValueError("image_size and img_height/img_width cannot both be provided")
 
-        if img_height is not None and img_width is not None:
+        # Initial check for image_size conflict remains as is
+        if image_size is not None and (
+            img_height is not None or img_width is not None or height is not None or width is not None
+        ):
+            raise ValueError("image_size cannot be provided alongside img_height/img_width or height/width")
+
+        # Prioritize height/width (HF-aligned)
+        if height is not None and width is not None:
+            if img_height is not None or img_width is not None:
+                # Raise error if both sets of arguments are provided
+                raise ValueError(
+                    "Cannot provide both 'height'/'width' and 'img_height'/'img_width' simultaneously. "
+                    "Please use one set of arguments for image dimensions, preferring 'height'/'width'."
+                )
+            image_size = (height, width)
+        elif (height is not None and width is None) or (height is None and width is not None):
+            raise ValueError("Both height and width must be provided together if used")
+        # Fallback to img_height/img_width for backward compatibility
+        elif img_height is not None and img_width is not None:
             image_size = (img_height, img_width)
+        elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
+            raise ValueError("Both img_height and img_width must be provided together if used")
 
         self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
         self.text_encoder_2 = self.init_submodule_config(

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_kandinsky2_2.py

@@ -35,6 +35,8 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
         image_size: Optional[Tuple[int, int]] = None,
         img_height: Optional[int] = None,
         img_width: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -50,6 +52,8 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
                 Cannot be used together with img_height/img_width.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
+            height (Optional[int]): Height of the generated images.
+            width (Optional[int]): Width of the generated images.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
         Raises:
@@ -60,11 +64,29 @@ class _RBLNKandinskyV22PipelineBaseConfig(RBLNModelConfig):
             accommodate classifier-free guidance.
         """
         super().__init__(**kwargs)
-        if image_size is not None and (img_height is not None or img_width is not None):
-            raise ValueError("image_size and img_height/img_width cannot both be provided")
 
-        if img_height is not None and img_width is not None:
+        # Initial check for image_size conflict remains as is
+        if image_size is not None and (
+            img_height is not None or img_width is not None or height is not None or width is not None
+        ):
+            raise ValueError("image_size cannot be provided alongside img_height/img_width or height/width")
+
+        # Prioritize height/width (HF-aligned)
+        if height is not None and width is not None:
+            if img_height is not None or img_width is not None:
+                # Raise error if both sets of arguments are provided
+                raise ValueError(
+                    "Cannot provide both 'height'/'width' and 'img_height'/'img_width' simultaneously. "
+                    "Please use one set of arguments for image dimensions, preferring 'height'/'width'."
+                )
+            image_size = (height, width)
+        elif (height is not None and width is None) or (height is None and width is not None):
+            raise ValueError("Both height and width must be provided together if used")
+        # Fallback to img_height/img_width for backward compatibility
+        elif img_height is not None and img_width is not None:
             image_size = (img_height, img_width)
+        elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
+            raise ValueError("Both img_height and img_width must be provided together if used")
 
         self.unet = self.init_submodule_config(RBLNUNet2DConditionModelConfig, unet, sample_size=sample_size)
         self.movq = self.init_submodule_config(
@@ -186,6 +208,8 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         img_height: Optional[int] = None,
         img_width: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
         guidance_scale: Optional[float] = None,
         prior_prior: Optional[RBLNPriorTransformerConfig] = None,
         prior_image_encoder: Optional[RBLNCLIPVisionModelWithProjectionConfig] = None,
@@ -212,6 +236,8 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
+            height (Optional[int]): Height of the generated images.
+            width (Optional[int]): Width of the generated images.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
             prior_prior (Optional[RBLNPriorTransformerConfig]): Direct configuration for the prior transformer.
                 Used if prior_pipe is not provided.
@@ -226,6 +252,30 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
         """
         super().__init__(**kwargs)
+
+        # Initial check for image_size conflict remains as is
+        if image_size is not None and (
+            img_height is not None or img_width is not None or height is not None or width is not None
+        ):
+            raise ValueError("image_size cannot be provided alongside img_height/img_width or height/width")
+
+        # Prioritize height/width (HF-aligned)
+        if height is not None and width is not None:
+            if img_height is not None or img_width is not None:
+                # Raise error if both sets of arguments are provided
+                raise ValueError(
+                    "Cannot provide both 'height'/'width' and 'img_height'/'img_width' simultaneously. "
+                    "Please use one set of arguments for image dimensions, preferring 'height'/'width'."
+                )
+            image_size = (height, width)
+        elif (height is not None and width is None) or (height is None and width is not None):
+            raise ValueError("Both height and width must be provided together if used")
+        # Fallback to img_height/img_width for backward compatibility
+        elif img_height is not None and img_width is not None:
+            image_size = (img_height, img_width)
+        elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
+            raise ValueError("Both img_height and img_width must be provided together if used")
+
         self.prior_pipe = self.init_submodule_config(
             RBLNKandinskyV22PriorPipelineConfig,
             prior_pipe,
@@ -243,8 +293,6 @@ class _RBLNKandinskyV22CombinedPipelineBaseConfig(RBLNModelConfig):
             batch_size=batch_size,
             sample_size=sample_size,
             image_size=image_size,
-            img_height=img_height,
-            img_width=img_width,
             guidance_scale=guidance_scale,
         )
 

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion.py

@@ -32,6 +32,8 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         img_height: Optional[int] = None,
         img_width: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
@@ -48,6 +50,8 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
+            height (Optional[int]): Height of the generated images.
+            width (Optional[int]): Width of the generated images.
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
@@ -62,11 +66,29 @@ class _RBLNStableDiffusionPipelineBaseConfig(RBLNModelConfig):
             accommodate classifier-free guidance.
         """
         super().__init__(**kwargs)
-        if image_size is not None and (img_height is not None or img_width is not None):
-            raise ValueError("image_size and img_height/img_width cannot both be provided")
 
-        if img_height is not None and img_width is not None:
+        # Initial check for image_size conflict remains as is
+        if image_size is not None and (
+            img_height is not None or img_width is not None or height is not None or width is not None
+        ):
+            raise ValueError("image_size cannot be provided alongside img_height/img_width or height/width")
+
+        # Prioritize height/width (HF-aligned)
+        if height is not None and width is not None:
+            if img_height is not None or img_width is not None:
+                # Raise error if both sets of arguments are provided
+                raise ValueError(
+                    "Cannot provide both 'height'/'width' and 'img_height'/'img_width' simultaneously. "
+                    "Please use one set of arguments for image dimensions, preferring 'height'/'width'."
+                )
+            image_size = (height, width)
+        elif (height is not None and width is None) or (height is None and width is not None):
+            raise ValueError("Both height and width must be provided together if used")
+        # Fallback to img_height/img_width for backward compatibility
+        elif img_height is not None and img_width is not None:
             image_size = (img_height, img_width)
+        elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
+            raise ValueError("Both img_height and img_width must be provided together if used")
 
         self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
         self.unet = self.init_submodule_config(

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_3.py

@@ -37,6 +37,8 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         img_height: Optional[int] = None,
         img_width: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
         guidance_scale: Optional[float] = None,
         **kwargs,
     ):
@@ -59,6 +61,8 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
+            height (Optional[int]): Height of the generated images.
+            width (Optional[int]): Width of the generated images.
             guidance_scale (Optional[float]): Scale for classifier-free guidance.
             **kwargs: Additional arguments passed to the parent RBLNModelConfig.
 
@@ -70,11 +74,29 @@ class _RBLNStableDiffusion3PipelineBaseConfig(RBLNModelConfig):
             accommodate classifier-free guidance.
         """
         super().__init__(**kwargs)
-        if image_size is not None and (img_height is not None or img_width is not None):
-            raise ValueError("image_size and img_height/img_width cannot both be provided")
 
-        if img_height is not None and img_width is not None:
+        # Initial check for image_size conflict remains as is
+        if image_size is not None and (
+            img_height is not None or img_width is not None or height is not None or width is not None
+        ):
+            raise ValueError("image_size cannot be provided alongside img_height/img_width or height/width")
+
+        # Prioritize height/width (HF-aligned)
+        if height is not None and width is not None:
+            if img_height is not None or img_width is not None:
+                # Raise error if both sets of arguments are provided
+                raise ValueError(
+                    "Cannot provide both 'height'/'width' and 'img_height'/'img_width' simultaneously. "
+                    "Please use one set of arguments for image dimensions, preferring 'height'/'width'."
+                )
+            image_size = (height, width)
+        elif (height is not None and width is None) or (height is None and width is not None):
+            raise ValueError("Both height and width must be provided together if used")
+        # Fallback to img_height/img_width for backward compatibility
+        elif img_height is not None and img_width is not None:
             image_size = (img_height, img_width)
+        elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
+            raise ValueError("Both img_height and img_width must be provided together if used")
 
         max_seq_len = max_seq_len or 256
 

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/diffusers/configurations/pipelines/configuration_stable_diffusion_xl.py

@@ -33,6 +33,8 @@ class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
         batch_size: Optional[int] = None,
         img_height: Optional[int] = None,
         img_width: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
         sample_size: Optional[Tuple[int, int]] = None,
         image_size: Optional[Tuple[int, int]] = None,
         guidance_scale: Optional[float] = None,
@@ -51,6 +53,8 @@ class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
             batch_size (Optional[int]): Batch size for inference, applied to all submodules.
             img_height (Optional[int]): Height of the generated images.
             img_width (Optional[int]): Width of the generated images.
+            height (Optional[int]): Height of the generated images.
+            width (Optional[int]): Width of the generated images.
             sample_size (Optional[Tuple[int, int]]): Spatial dimensions for the UNet model.
             image_size (Optional[Tuple[int, int]]): Alternative way to specify image dimensions.
                 Cannot be used together with img_height/img_width.
@@ -65,11 +69,29 @@ class _RBLNStableDiffusionXLPipelineBaseConfig(RBLNModelConfig):
             accommodate classifier-free guidance.
         """
         super().__init__(**kwargs)
-        if image_size is not None and (img_height is not None or img_width is not None):
-            raise ValueError("image_size and img_height/img_width cannot both be provided")
 
-        if img_height is not None and img_width is not None:
+        # Initial check for image_size conflict remains as is
+        if image_size is not None and (
+            img_height is not None or img_width is not None or height is not None or width is not None
+        ):
+            raise ValueError("image_size cannot be provided alongside img_height/img_width or height/width")
+
+        # Prioritize height/width (HF-aligned)
+        if height is not None and width is not None:
+            if img_height is not None or img_width is not None:
+                # Raise error if both sets of arguments are provided
+                raise ValueError(
+                    "Cannot provide both 'height'/'width' and 'img_height'/'img_width' simultaneously. "
+                    "Please use one set of arguments for image dimensions, preferring 'height'/'width'."
+                )
+            image_size = (height, width)
+        elif (height is not None and width is None) or (height is None and width is not None):
+            raise ValueError("Both height and width must be provided together if used")
+        # Fallback to img_height/img_width for backward compatibility
+        elif img_height is not None and img_width is not None:
             image_size = (img_height, img_width)
+        elif (img_height is not None and img_width is None) or (img_height is None and img_width is not None):
+            raise ValueError("Both img_height and img_width must be provided together if used")
 
         self.text_encoder = self.init_submodule_config(RBLNCLIPTextModelConfig, text_encoder, batch_size=batch_size)
         self.text_encoder_2 = self.init_submodule_config(

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/diffusers/models/transformers/prior_transformer.py

@@ -58,8 +58,7 @@ class _PriorTransformer(torch.nn.Module):
 class RBLNPriorTransformer(RBLNModel):
     hf_library_name = "diffusers"
     auto_model_class = PriorTransformer
-    output_class = PriorTransformerOutput
-    output_key = "predicted_image_embedding"
+    _output_class = PriorTransformerOutput
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/diffusers/models/transformers/transformer_sd3.py

@@ -61,8 +61,7 @@ class SD3Transformer2DModelWrapper(torch.nn.Module):
 class RBLNSD3Transformer2DModel(RBLNModel):
     hf_library_name = "diffusers"
     auto_model_class = SD3Transformer2DModel
-    output_class = Transformer2DModelOutput
-    output_key = "sample"
+    _output_class = Transformer2DModelOutput
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/diffusers/models/unets/unet_2d_condition.py

@@ -143,8 +143,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
     hf_library_name = "diffusers"
     auto_model_class = UNet2DConditionModel
     _rbln_config_class = RBLNUNet2DConditionModelConfig
-    output_class = UNet2DConditionOutput
-    output_key = "sample"
+    _output_class = UNet2DConditionOutput
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/modeling.py

@@ -14,7 +14,7 @@
 
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Union, get_args, get_origin, get_type_hints
 
 import rebel
 import torch
@@ -49,8 +49,7 @@ class RBLNModel(RBLNBaseModel):
         ```
     """
 
-    output_class = None
-    output_key = "last_hidden_state"
+    _output_class = None
 
     @classmethod
     def update_kwargs(cls, kwargs):
@@ -245,16 +244,61 @@ class RBLNModel(RBLNBaseModel):
         # Format output according to task requirements
         return self._prepare_output(output, return_dict)
 
+    @classmethod
+    def get_hf_output_class(cls):
+        """
+        Dynamically gets the output class from the corresponding HuggingFace model class.
+
+        Returns:
+            type: The appropriate output class from transformers or diffusers
+        """
+        if cls._output_class:
+            return cls._output_class
+
+        hf_class = cls.get_hf_class()
+        if hf_class is None:
+            raise ValueError(f"No HuggingFace model class found for {cls.__name__}")
+
+        hints = get_type_hints(hf_class.forward) if hasattr(hf_class, "forward") else {}
+        ret = hints.get("return")
+
+        if ret is not None:
+            candidates = get_args(ret) if get_origin(ret) is Union else (ret,)
+
+            for t in candidates:
+                if t is type(None):  # Skip NoneType in Union
+                    continue
+                mod = getattr(t, "__module__", "")
+                if "transformers" in mod or "diffusers" in mod:
+                    cls._output_class = t
+                    return t
+
+        # Fallback to BaseModelOutput
+        cls._output_class = BaseModelOutput
+        return BaseModelOutput
+
     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.
         """
+        tuple_output = (output,) if not isinstance(output, (tuple, list)) else tuple(output)
         if not return_dict:
-            return (output,) if not isinstance(output, (tuple, list)) else output
+            return tuple_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})
+            output_class = self.get_hf_output_class()
+            if hasattr(output_class, "loss"):
+                tuple_output = (None,) + tuple_output
+
+            # Truncate if we have too many outputs, otherwise use as is
+            if hasattr(output_class, "__annotations__"):
+                num_fields = len(output_class.__annotations__)
+                if len(tuple_output) > num_fields:
+                    tuple_output = tuple_output[:num_fields]
+                    logger.warning(
+                        f"Truncating output to {num_fields} fields for {output_class.__name__}. "
+                        f"Expected {num_fields} fields, but got {len(tuple_output)} fields."
+                        "This is unexpected. Please report this issue to the developers."
+                    )
+
+            return output_class(*tuple_output)

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/modeling_base.py

@@ -178,9 +178,27 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
         return str(model_path)
 
     @classmethod
-    def _load_compiled_models(cls, model_path: str):
+    def _load_compiled_models(cls, model_path: str, expected_compiled_model_names: List[str]):
         compiled_models = Path(model_path).glob("*.rbln")
-        rbln_compiled_models = {cm.stem: rebel.RBLNCompiledModel(cm) for cm in compiled_models}
+        expected_compiled_models = [
+            Path(model_path) / f"{compiled_model_name}.rbln" for compiled_model_name in expected_compiled_model_names
+        ]
+        unexpected_compiled_models = [cm for cm in compiled_models if cm not in expected_compiled_models]
+        if unexpected_compiled_models:
+            # TODO(jongho): fix after May release. raise error if unexpected compiled models are found
+            logger.warning(
+                f"Unexpected compiled models found: {[cm.name for cm in unexpected_compiled_models]}. "
+                f"Please check the model path: {model_path}"
+            )
+
+        rbln_compiled_models = {}
+        for compiled_model in expected_compiled_models:
+            if not compiled_model.exists():
+                raise FileNotFoundError(
+                    f"Expected RBLN compiled model '{compiled_model.name}' not found at '{model_path}'. "
+                    "Please ensure all models specified in `rbln_config` are present."
+                )
+            rbln_compiled_models[compiled_model.stem] = rebel.RBLNCompiledModel(compiled_model)
         return rbln_compiled_models
 
     @classmethod
@@ -271,7 +289,8 @@ class RBLNBaseModel(SubModulesMixin, PushToHubMixin, PreTrainedModel):
                     )
                     config = PretrainedConfig(**config)
 
-            rbln_compiled_models = cls._load_compiled_models(model_path_subfolder)
+            compiled_model_names = [cfg.compiled_model_name for cfg in rbln_config.compile_cfgs]
+            rbln_compiled_models = cls._load_compiled_models(model_path_subfolder, compiled_model_names)
 
             if subfolder != "":
                 model_save_dir = Path(model_path_subfolder).absolute().parent

{optimum_rbln-0.7.5rc0 → optimum_rbln-0.7.5rc2}/src/optimum/rbln/transformers/modeling_generic.py

@@ -36,11 +36,7 @@ from transformers import (
 )
 from transformers.modeling_outputs import (
     BaseModelOutput,
-    DepthEstimatorOutput,
-    ImageClassifierOutput,
-    MaskedLMOutput,
     QuestionAnsweringModelOutput,
-    SequenceClassifierOutput,
 )
 
 from ..configuration_utils import RBLNCompileConfig
@@ -63,8 +59,6 @@ class _RBLNTransformerEncoder(RBLNModel):
     auto_model_class = AutoModel
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
     rbln_dtype = "int64"
-    output_class = BaseModelOutput
-    output_key = "last_hidden_state"
 
     @classmethod
     def _update_rbln_config(
@@ -149,7 +143,6 @@ class _RBLNImageModel(RBLNModel):
     auto_model_class = AutoModel
     main_input_name = "pixel_values"
     output_class = BaseModelOutput
-    output_key = "last_hidden_state"
 
     @classmethod
     def _update_rbln_config(
@@ -223,15 +216,11 @@ class RBLNModelForQuestionAnswering(_RBLNTransformerEncoder):
 class RBLNModelForSequenceClassification(_RBLNTransformerEncoder):
     auto_model_class = AutoModelForSequenceClassification
     rbln_model_input_names = ["input_ids", "attention_mask"]
-    output_class = SequenceClassifierOutput
-    output_key = "logits"
 
 
 class RBLNModelForMaskedLM(_RBLNTransformerEncoder):
     auto_model_class = AutoModelForMaskedLM
     rbln_model_input_names = ["input_ids", "attention_mask"]
-    output_class = MaskedLMOutput
-    output_key = "logits"
 
 
 class RBLNModelForTextEncoding(_RBLNTransformerEncoder):
@@ -243,20 +232,14 @@ class RBLNTransformerEncoderForFeatureExtraction(_RBLNTransformerEncoder):
     # TODO: RBLNModel is also for feature extraction.
     auto_model_class = AutoModel
     rbln_model_input_names = ["input_ids", "attention_mask"]
-    output_class = BaseModelOutput
-    output_key = "last_hidden_state"
 
 
 class RBLNModelForImageClassification(_RBLNImageModel):
     auto_model_class = AutoModelForImageClassification
-    output_class = ImageClassifierOutput
-    output_key = "logits"
 
 
 class RBLNModelForDepthEstimation(_RBLNImageModel):
     auto_model_class = AutoModelForDepthEstimation
-    output_class = DepthEstimatorOutput
-    output_key = "predicted_depth"
 
 
 class RBLNModelForAudioClassification(RBLNModel):
@@ -273,8 +256,6 @@ class RBLNModelForAudioClassification(RBLNModel):
     """
 
     auto_model_class = AutoModelForAudioClassification
-    output_class = SequenceClassifierOutput
-    output_key = "logits"
 
     @classmethod
     def _update_rbln_config(