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

optimum/rbln/modeling_diffusers.py

@@ -20,14 +20,15 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
+import copy
 import importlib
 from os import PathLike
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import torch
 
-from .modeling_base import RBLNModel
-from .modeling_config import ContextRblnConfig, use_rbln_config
+from .modeling import RBLNModel
+from .modeling_config import RUNTIME_KEYWORDS, ContextRblnConfig, use_rbln_config
 from .utils.decorator_utils import remove_compile_time_kwargs
 
 
@@ -74,127 +75,40 @@ class RBLNDiffusionMixin:
 
     @classmethod
     @property
-    def use_encode(cls):
+    def img2img_pipeline(cls):
         return "Img2Img" in cls.__name__
 
     @classmethod
-    def _get_unet_batch_size(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> int:
-        # Calculates the batch size based on guidance scale
-        batch_size = rbln_config.get("batch_size", 1)
-        do_guidance = rbln_config.get("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        return batch_size * 2 if do_guidance else batch_size
-
-    @classmethod
-    def _get_vae_sample_size(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Union[int, Tuple[int, int]]:
-        image_size = (rbln_config.get("img_height"), rbln_config.get("img_width"))
-        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 cls.use_encode:
-                sample_size = model.vae.config.sample_size
-            else:
-                # In case of text2img, sample size of vae decoder is determined by unet.
-                unet_sample_size = model.unet.config.sample_size
-                if isinstance(unet_sample_size, int):
-                    sample_size = unet_sample_size * model.vae_scale_factor
-                else:
-                    sample_size = (
-                        unet_sample_size[0] * model.vae_scale_factor,
-                        unet_sample_size[1] * model.vae_scale_factor,
-                    )
-
-        else:
-            sample_size = (image_size[0], image_size[1])
-        return sample_size
-
-    @classmethod
-    def _get_unet_sample_size(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Union[int, Tuple[int, int]]:
-        image_size = (rbln_config.get("img_height"), rbln_config.get("img_width"))
-        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 cls.use_encode:
-                # In case of img2img, sample size of unet is determined by vae encoder.
-                vae_sample_size = model.vae.config.sample_size
-                if isinstance(vae_sample_size, int):
-                    sample_size = vae_sample_size // model.vae_scale_factor
-                else:
-                    sample_size = (
-                        vae_sample_size[0] // model.vae_scale_factor,
-                        vae_sample_size[1] // model.vae_scale_factor,
-                    )
-            else:
-                sample_size = model.unet.config.sample_size
-        else:
-            sample_size = (image_size[0] // model.vae_scale_factor, image_size[1] // model.vae_scale_factor)
-        return sample_size
-
-    @classmethod
-    def _get_default_config(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        # default configurations for each submodules
-        return {"img2img_pipeline": cls.use_encode}
+    @property
+    def inpaint_pipeline(cls):
+        return "Inpaint" in cls.__name__
 
     @classmethod
-    def get_default_rbln_config_text_encoder(
-        cls, model: torch.nn.Module, rbln_config: Dict[str, Any]
+    def get_submodule_rbln_config(
+        cls, model: torch.nn.Module, submodule_name: str, rbln_config: Dict[str, Any]
     ) -> Dict[str, Any]:
-        batch_size = rbln_config.get("batch_size", 1)
-        return {"batch_size": batch_size}
+        submodule = getattr(model, submodule_name)
+        submodule_class_name = submodule.__class__.__name__
 
-    @classmethod
-    def get_default_rbln_config_text_encoder_2(
-        cls, model: torch.nn.Module, rbln_config: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        batch_size = rbln_config.get("batch_size", 1)
-        return {"batch_size": batch_size}
+        if submodule_class_name == "MultiControlNetModel":
+            submodule_class_name = "ControlNetModel"
 
-    @classmethod
-    def get_default_rbln_config_unet(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        # configuration for unet
-        unet_batch_size = cls._get_unet_batch_size(model, rbln_config)
-        text_model_hidden_size = model.text_encoder_2.config.hidden_size if hasattr(model, "text_encoder_2") else None
-        return {
-            **cls._get_default_config(model, rbln_config),
-            "max_seq_len": model.text_encoder.config.max_position_embeddings,
-            "text_model_hidden_size": text_model_hidden_size,
-            "batch_size": unet_batch_size,
-            "sample_size": cls._get_unet_sample_size(model, rbln_config),
-            "is_controlnet": "controlnet" in model.config.keys(),
-        }
+        submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), f"RBLN{submodule_class_name}")
 
-    @classmethod
-    def get_default_rbln_config_vae(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        # configuration for vae
-        batch_size = rbln_config.get("batch_size", 1)
-        return {
-            **cls._get_default_config(model, rbln_config),
-            "sample_size": cls._get_vae_sample_size(model, rbln_config),
-            "batch_size": batch_size,
-        }
+        submodule_config = rbln_config.get(submodule_name, {})
+        submodule_config = copy.deepcopy(submodule_config)
 
-    @classmethod
-    def get_default_rbln_config_controlnet(cls, model: torch.nn.Module, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
-        # configuration for controlnet
-        unet_batch_size = cls._get_unet_batch_size(model, rbln_config)
-        text_model_hidden_size = model.text_encoder_2.config.hidden_size if hasattr(model, "text_encoder_2") else None
-        return {
-            **cls._get_default_config(model, rbln_config),
-            "max_seq_len": model.text_encoder.config.max_position_embeddings,
-            "vae_sample_size": cls._get_vae_sample_size(model, rbln_config),
-            "unet_sample_size": cls._get_unet_sample_size(model, rbln_config),
-            "batch_size": unet_batch_size,
-            "text_model_hidden_size": text_model_hidden_size,
-        }
+        pipe_global_config = {k: v for k, v in rbln_config.items() if k not in cls._submodules}
 
-    @classmethod
-    def get_default_rbln_config(
-        cls, model: torch.nn.Module, submodule_name: str, rbln_config: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        # Returns the default configuration based on submodule name
-        config_method = f"get_default_rbln_config_{submodule_name}"
-        if hasattr(cls, config_method):
-            return getattr(cls, config_method)(model, rbln_config)
-        raise ValueError(f"Unknown submodule: {submodule_name}")
+        submodule_config.update({k: v for k, v in pipe_global_config.items() if k not in submodule_config})
+        submodule_config.update(
+            {
+                "img2img_pipeline": cls.img2img_pipeline,
+                "inpaint_pipeline": cls.inpaint_pipeline,
+            }
+        )
+        submodule_config = submodule_cls.update_rbln_config_using_pipe(model, submodule_config)
+        return submodule_config
 
     @staticmethod
     def _maybe_apply_and_fuse_lora(
@@ -256,12 +170,32 @@ class RBLNDiffusionMixin:
 
         else:
             # raise error if any of submodules are torch module.
-            for name in cls._submodules:
-                if isinstance(kwargs.get(name), torch.nn.Module):
+            model_index_config = None
+            for submodule_name in cls._submodules:
+                if isinstance(kwargs.get(submodule_name), torch.nn.Module):
                     raise AssertionError(
-                        f"{name} is not compiled torch module. If you want to compile, set `export=True`."
+                        f"{submodule_name} is not compiled torch module. If you want to compile, set `export=True`."
                     )
 
+                # Load submodule outside if runtime kwargs(e.g. device) is specified.
+                if submodule_config := rbln_config.get(submodule_name):
+                    if any(kwd in submodule_config for kwd in RUNTIME_KEYWORDS):
+                        if model_index_config is None:
+                            model_index_config = cls.load_config(pretrained_model_name_or_path=model_id)
+
+                        module_name, class_name = model_index_config[submodule_name]
+                        if module_name != "optimum.rbln":
+                            raise ValueError(
+                                f"Invalid module_name '{module_name}' found in model_index.json for "
+                                f"submodule '{submodule_name}'. "
+                                "Expected 'optimum.rbln'. Please check the model_index.json configuration."
+                            )
+                        submodule_cls: RBLNModel = getattr(importlib.import_module("optimum.rbln"), class_name)
+                        submodule = submodule_cls.from_pretrained(
+                            model_id, export=False, subfolder=submodule_name, rbln_config=submodule_config
+                        )
+                        kwargs[submodule_name] = submodule
+
         with ContextRblnConfig(
             device=rbln_config.get("device"),
             device_map=rbln_config.get("device_map"),
@@ -291,16 +225,11 @@ class RBLNDiffusionMixin:
         model_save_dir: Optional[PathLike],
         rbln_config: Dict[str, Any],
     ) -> Dict[str, RBLNModel]:
-        # Compile submodules based on rbln_config
         compiled_submodules = {}
 
-        # FIXME : Currently, optimum-rbln for transformer does not use base rbln config.
-        base_rbln_config = {k: v for k, v in rbln_config.items() if k not in cls._submodules}
         for submodule_name in cls._submodules:
             submodule = passed_submodules.get(submodule_name) or getattr(model, submodule_name, None)
-            submodule_rbln_config = cls.get_default_rbln_config(model, submodule_name, rbln_config)
-            submodule_rbln_config.update(base_rbln_config)
-            submodule_rbln_config.update(rbln_config.get(submodule_name, {}))
+            submodule_rbln_config = cls.get_submodule_rbln_config(model, submodule_name, rbln_config)
 
             if submodule is None:
                 raise ValueError(f"submodule ({submodule_name}) cannot be accessed since it is not provided.")

optimum/rbln/transformers/__init__.py

@@ -28,7 +28,6 @@ from transformers.utils import _LazyModule
 
 _import_structure = {
     "cache_utils": ["RebelDynamicCache"],
-    "generation": ["BatchTextIteratorStreamer"],
     "models": [
         "RBLNAutoModel",
         "RBLNAutoModelForAudioClassification",
@@ -68,7 +67,6 @@ _import_structure = {
 
 if TYPE_CHECKING:
     from .cache_utils import RebelDynamicCache
-    from .generation import BatchTextIteratorStreamer
     from .models import (
         RBLNAutoModel,
         RBLNAutoModelForAudioClassification,

optimum/rbln/transformers/models/auto/auto_factory.py

@@ -22,8 +22,16 @@
 # from Rebellions Inc.
 
 import importlib
+import inspect
+import warnings
 
-from transformers import AutoConfig
+from transformers import AutoConfig, PretrainedConfig
+from transformers.dynamic_module_utils import get_class_from_dynamic_module
+from transformers.models.auto.auto_factory import _get_model_class
+
+from optimum.rbln.modeling_base import RBLNBaseModel
+from optimum.rbln.modeling_config import RBLNConfig
+from optimum.rbln.utils.model_utils import convert_hf_to_rbln_model_name, convert_rbln_to_hf_model_name
 
 
 class _BaseAutoModelClass:
@@ -33,46 +41,132 @@ class _BaseAutoModelClass:
     def __init__(self, *args, **kwargs):
         raise EnvironmentError(
             f"{self.__class__.__name__} is designed to be instantiated "
-            f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)` or "
-            f"`{self.__class__.__name__}.from_config(config)` methods."
+            f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)`"
         )
 
     @classmethod
     def get_rbln_cls(
         cls,
-        model_id,
+        pretrained_model_name_or_path,
         *args,
+        export=True,
         **kwargs,
     ):
-        # kwargs.update({"return_unused_kwargs": True})
-        config = AutoConfig.from_pretrained(model_id, return_unused_kwargs=True, **kwargs)[0]
-
-        if len(config.architectures) > 1:
-            raise ValueError(
-                f"Model with ID '{model_id}' has multiple architectures defined in the configuration: "
-                f"{config.architectures}. `_BaseAutoModelClass` require exactly one architecture. "
-            )
-
-        architecture_name = config.architectures[0]
-        if architecture_name not in cls._model_mapping.values():
-            raise ValueError(
-                f"The 'RBLN{architecture_name}' architecture is not supported by `{cls.__name__}.from_pretrained()`."
-                "Please use the appropriate class's `from_pretrained()` method to load this model."
-            )
-
-        rbln_class_name = "RBLN" + architecture_name
-        module = importlib.import_module("optimum.rbln")
+        """
+        Determine the appropriate RBLN model class based on the given model ID and configuration.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier or path to the pretrained model.
+            export (bool): Whether to infer the class based on Hugging Face (HF) architecture.
+            kwargs: Additional arguments for configuration and loading.
+
+        Returns:
+            RBLNBaseModel: The corresponding RBLN model class.
+        """
+        if export:
+            hf_model_class = cls.infer_hf_model_class(pretrained_model_name_or_path, **kwargs)
+            rbln_class_name = convert_hf_to_rbln_model_name(hf_model_class.__name__)
+        else:
+            rbln_class_name = cls.get_rbln_model_class_name(pretrained_model_name_or_path, **kwargs)
+
+            if convert_rbln_to_hf_model_name(rbln_class_name) not in cls._model_mapping_names.values():
+                raise ValueError(
+                    f"The architecture '{rbln_class_name}' is not supported by the `{cls.__name__}.from_pretrained()` method. "
+                    "Please use the `from_pretrained()` method of the appropriate class to load this model, "
+                    f"or directly use '{rbln_class_name}.from_pretrained()`."
+                )
 
         try:
+            module = importlib.import_module("optimum.rbln")
             rbln_cls = getattr(module, rbln_class_name)
         except AttributeError as e:
             raise AttributeError(
-                f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{model_id}'. "
+                f"Class '{rbln_class_name}' not found in 'optimum.rbln' module for model ID '{pretrained_model_name_or_path}'. "
                 "Ensure that the class name is correctly mapped and available in the 'optimum.rbln' module."
             ) from e
 
         return rbln_cls
 
+    @classmethod
+    def infer_hf_model_class(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        **kwargs,
+    ):
+        """
+        Infer the Hugging Face model class based on the configuration or model name.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier or path to the pretrained model.
+            kwargs: Additional arguments for configuration and loading.
+
+        Returns:
+            PretrainedModel: The inferred Hugging Face model class.
+        """
+
+        # Try to load configuration if provided or retrieve it from the model ID
+        config = kwargs.pop("config", None)
+        kwargs.update({"trust_remote_code": True})
+        kwargs["_from_auto"] = True
+
+        # Load configuration if not already provided
+        if not isinstance(config, PretrainedConfig):
+            config, kwargs = AutoConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                return_unused_kwargs=True,
+                **kwargs,
+            )
+
+        # Get hf_model_class from Config
+        has_remote_code = (
+            hasattr(config, "auto_map") and convert_rbln_to_hf_model_name(cls.__name__) in config.auto_map
+        )
+        if has_remote_code:
+            class_ref = config.auto_map[convert_rbln_to_hf_model_name(cls.__name__)]
+            model_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
+        elif type(config) in cls._model_mapping.keys():
+            model_class = _get_model_class(config, cls._model_mapping)
+        else:
+            raise ValueError(
+                f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
+                f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping.keys())}."
+            )
+
+        if model_class.__name__ != config.architectures[0]:
+            warnings.warn(
+                f"`{cls.__name__}.from_pretrained()` is invoking `{convert_hf_to_rbln_model_name(model_class.__name__)}.from_pretrained()`, which does not match the "
+                f"expected architecture `RBLN{config.architectures[0]}` from config. This mismatch could cause some operations to not be properly loaded "
+                f"from the checkpoint, leading to potential unintended behavior. If this is not intentional, consider calling the "
+                f"`from_pretrained()` method directly from the `RBLN{config.architectures[0]}` class instead.",
+                UserWarning,
+            )
+
+        return model_class
+
+    @classmethod
+    def get_rbln_model_class_name(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Retrieve the path to the compiled model directory for a given RBLN model.
+
+        Args:
+            pretrained_model_name_or_path (str): Identifier of the model.
+            kwargs: Additional arguments that match the parameters of `_load_compiled_model_dir`.
+
+        Returns:
+            str: Path to the compiled model directory.
+        """
+        sig = inspect.signature(RBLNBaseModel._load_compiled_model_dir)
+        valid_params = sig.parameters.keys()
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+
+        model_path_subfolder = RBLNBaseModel._load_compiled_model_dir(
+            model_id=pretrained_model_name_or_path, **filtered_kwargs
+        )
+        rbln_config = RBLNConfig.load(model_path_subfolder)
+
+        return rbln_config.meta["cls"]
+
     @classmethod
     def from_pretrained(
         cls,

optimum/rbln/transformers/models/auto/modeling_auto.py

@@ -21,18 +21,31 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
+
 from transformers.models.auto.modeling_auto import (
+    MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING,
     MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_CAUSAL_LM_MAPPING,
     MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,
+    MODEL_FOR_CTC_MAPPING,
     MODEL_FOR_CTC_MAPPING_NAMES,
+    MODEL_FOR_DEPTH_ESTIMATION_MAPPING,
     MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES,
+    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
     MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_MASKED_LM_MAPPING,
     MODEL_FOR_MASKED_LM_MAPPING_NAMES,
+    MODEL_FOR_QUESTION_ANSWERING_MAPPING,
     MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES,
+    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
     MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
+    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
     MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES,
+    MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING,
     MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES,
+    MODEL_FOR_VISION_2_SEQ_MAPPING,
     MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES,
+    MODEL_MAPPING,
     MODEL_MAPPING_NAMES,
 )
 
@@ -48,48 +61,60 @@ MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.update(
 
 
 class RBLNAutoModel(_BaseAutoModelClass):
-    _model_mapping = MODEL_MAPPING_NAMES
+    _model_mapping = MODEL_MAPPING
+    _model_mapping_names = MODEL_MAPPING_NAMES
 
 
 class RBLNAutoModelForCTC(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_CTC_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_CTC_MAPPING
+    _model_mapping_names = MODEL_FOR_CTC_MAPPING_NAMES
 
 
 class RBLNAutoModelForCausalLM(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_CAUSAL_LM_MAPPING
+    _model_mapping_names = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForSeq2SeqLM(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
+    _model_mapping_names = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForSpeechSeq2Seq(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING
+    _model_mapping_names = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
 
 
 class RBLNAutoModelForDepthEstimation(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_DEPTH_ESTIMATION_MAPPING
+    _model_mapping_names = MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForSequenceClassification(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
+    _model_mapping_names = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForVision2Seq(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_VISION_2_SEQ_MAPPING
+    _model_mapping_names = MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES
 
 
 class RBLNAutoModelForMaskedLM(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_MASKED_LM_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_MASKED_LM_MAPPING
+    _model_mapping_names = MODEL_FOR_MASKED_LM_MAPPING_NAMES
 
 
 class RBLNAutoModelForAudioClassification(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
+    _model_mapping_names = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForImageClassification(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
+    _model_mapping_names = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
 
 
 class RBLNAutoModelForQuestionAnswering(_BaseAutoModelClass):
-    _model_mapping = MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
+    _model_mapping = MODEL_FOR_QUESTION_ANSWERING_MAPPING
+    _model_mapping_names = MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES

optimum/rbln/transformers/models/bart/modeling_bart.py

@@ -24,9 +24,9 @@
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
-from transformers import BartConfig, BartForConditionalGeneration, BartModel, PretrainedConfig
+from transformers import BartForConditionalGeneration, PretrainedConfig
 
-from ....modeling_base import RBLNModel
+from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
 from ....utils.logging import get_logger
 from ...models.seq2seq import RBLNModelForSeq2SeqLM
@@ -41,9 +41,6 @@ if TYPE_CHECKING:
 
 
 class RBLNBartModel(RBLNModel):
-    original_model_class = BartModel
-    original_config_class = BartConfig
-
     @classmethod
     def _get_rbln_config(
         cls,
@@ -82,7 +79,7 @@ class RBLNBartModel(RBLNModel):
             if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
                 rbln_model_input_names = cls.rbln_model_input_names
             elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
-                input_names_order = inspect.signature(cls.original_model_class.forward).parameters.keys()
+                input_names_order = inspect.signature(cls.hf_class.forward).parameters.keys()
                 raise ValueError(
                     "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
                     f"and be sure to make the order of the inputs same as BartModel forward() arguments like ({list(input_names_order)})"

optimum/rbln/transformers/models/bert/modeling_bert.py

@@ -25,9 +25,9 @@ import inspect
 import logging
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
-from transformers import BertConfig, BertModel, PretrainedConfig
+from transformers import PretrainedConfig
 
-from ....modeling_base import RBLNModel
+from ....modeling import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
 
 
@@ -38,9 +38,6 @@ if TYPE_CHECKING:
 
 
 class RBLNBertModel(RBLNModel):
-    original_model_class = BertModel
-    original_config_class = BertConfig
-
     @classmethod
     def _get_rbln_config(
         cls,
@@ -75,7 +72,7 @@ class RBLNBertModel(RBLNModel):
             if rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names"):
                 rbln_model_input_names = cls.rbln_model_input_names
             elif rbln_model_input_names is None and hasattr(cls, "rbln_model_input_names") is False:
-                input_names_order = inspect.signature(cls.original_model_class.forward).parameters.keys()
+                input_names_order = inspect.signature(cls.hf_class.forward).parameters.keys()
                 raise ValueError(
                     "Specify the model input names obtained by the tokenizer via `rbln_model_input_names`, "
                     f"and be sure to make the order of the inputs same as BertModel forward() arguments like ({list(input_names_order)})"

optimum/rbln/transformers/models/clip/modeling_clip.py

@@ -28,16 +28,15 @@ import torch
 from transformers import (
     CLIPTextConfig,
     CLIPTextModel,
-    CLIPTextModelWithProjection,
     CLIPVisionConfig,
     CLIPVisionModel,
 )
 from transformers.modeling_outputs import BaseModelOutputWithPooling
 from transformers.models.clip.modeling_clip import CLIPTextModelOutput
 
-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
 
 
 logger = logging.getLogger(__name__)
@@ -57,20 +56,14 @@ class _TextEncoder(torch.nn.Module):
 
 
 class RBLNCLIPTextModel(RBLNModel):
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        with override_auto_classes(
-            config_func=CLIPTextConfig.from_pretrained,
-            model_func=CLIPTextModel.from_pretrained,
-            skip_taskmanager=False,
-        ):
-            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:
         return _TextEncoder(model).eval()
 
+    @classmethod
+    def update_rbln_config_using_pipe(cls, pipe: RBLNDiffusionMixin, rbln_config: Dict[str, Any]) -> Dict[str, Any]:
+        return rbln_config
+
     @classmethod
     def _get_rbln_config(
         cls,
@@ -114,7 +107,7 @@ class RBLNCLIPTextModel(RBLNModel):
 
 
 class RBLNCLIPTextModelWithProjection(RBLNCLIPTextModel):
-    original_model_class = CLIPTextModelWithProjection
+    pass
 
 
 class _VisionEncoder(torch.nn.Module):
@@ -128,16 +121,6 @@ class _VisionEncoder(torch.nn.Module):
 
 
 class RBLNCLIPVisionModel(RBLNModel):
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        with override_auto_classes(
-            config_func=CLIPVisionConfig.from_pretrained,
-            model_func=CLIPVisionModel.from_pretrained,
-            skip_taskmanager=False,
-        ):
-            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:
         return _VisionEncoder(model).eval()
@@ -146,7 +129,7 @@ class RBLNCLIPVisionModel(RBLNModel):
     def _get_rbln_config(
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
-        model_config: "CLIPTextConfig",
+        model_config: "CLIPVisionConfig",
         rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
         rbln_batch_size = rbln_kwargs.get("batch_size", 1)

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

@@ -22,9 +22,6 @@
 # from Rebellions Inc.
 
 from .decoderonly_architecture import (
-    DecoderOnlyAttention,
-    DecoderOnlyDecoderLayer,
-    DecoderOnlyModel,
     DecoderOnlyWrapper,
     RotaryEmbedding,
     apply_rotary_pos_emb,