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

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py

@@ -16,131 +16,9 @@
 
 from diffusers import StableDiffusionXLImg2ImgPipeline
 
-from ....modeling_base import RBLNBaseModel
-from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
-from ....utils.runtime_utils import ContextRblnConfig
-from ...models import RBLNAutoencoderKL, RBLNUNet2DConditionModel
+from ....modeling_diffusers import RBLNDiffusionMixin
 
 
-class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
-    @classmethod
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for image-to-image generation using Stable Diffusion XL.
-
-        This model inherits from [`StableDiffusionXLPipeline`]. Check the superclass documentation for the generic methods the
-        library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
-
-        It implements the methods to convert a pre-trained StableDiffusionXL pipeline into a RBLNStableDiffusionXL pipeline by:
-        - transferring the checkpoint weights of the original into an optimized RBLN graph,
-        - compiling the resulting graph using the RBLN compiler.
-
-        Args:
-            model_id (`Union[str, Path]`):
-                Can be either:
-                - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
-                - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
-        """
-        export = kwargs.pop("export", None)
-        model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config = kwargs.pop("rbln_config", None)
-        rbln_kwargs, _ = RBLNBaseModel.resolve_rbln_config(rbln_config, kwargs)
-
-        device = rbln_kwargs.get("device", None)
-        device_map = rbln_kwargs.get("device_map", None)
-        create_runtimes = rbln_kwargs.get("create_runtimes", None)
-        optimize_host_memory = rbln_kwargs.get("optimize_host_memory", None)
-
-        with ContextRblnConfig(
-            device=device,
-            device_map=device_map,
-            create_runtimes=create_runtimes,
-            optimze_host_mem=optimize_host_memory,
-        ):
-            model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
-
-        if export is None or export is False:
-            return model
-
-        do_classifier_free_guidance = (
-            rbln_kwargs.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
-        )
-
-        vae = RBLNAutoencoderKL.from_pretrained(
-            model_id=model_id,
-            subfolder="vae",
-            export=True,
-            model_save_dir=model_save_dir,
-            rbln_unet_sample_size=model.unet.config.sample_size,
-            rbln_use_encode=True,
-            rbln_vae_scale_factor=model.vae_scale_factor,
-            rbln_config={**rbln_kwargs},
-        )
-        text_encoder = RBLNCLIPTextModel.from_pretrained(
-            model_id=model_id,
-            subfolder="text_encoder",
-            export=True,
-            model_save_dir=model_save_dir,
-            rbln_config={**rbln_kwargs},
-        )
-        text_encoder_2 = RBLNCLIPTextModelWithProjection.from_pretrained(
-            model_id=model_id,
-            subfolder="text_encoder_2",
-            export=True,
-            model_save_dir=model_save_dir,
-            rbln_config={**rbln_kwargs},
-        )
-
-        batch_size = rbln_kwargs.pop("batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        unet = RBLNUNet2DConditionModel.from_pretrained(
-            model_id=model_id,
-            subfolder="unet",
-            export=True,
-            model_save_dir=model_save_dir,
-            rbln_max_seq_len=model.text_encoder.config.max_position_embeddings,
-            rbln_text_model_hidden_size=model.text_encoder_2.config.hidden_size,
-            rbln_batch_size=unet_batch_size,
-            rbln_use_encode=True,
-            rbln_vae_scale_factor=model.vae_scale_factor,
-            rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
-            rbln_config={**rbln_kwargs},
-        )
-
-        if model_save_dir is not None:
-            # To skip saving original pytorch modules
-            del (model.vae, model.text_encoder, model.unet, model.text_encoder_2)
-
-            # Direct calling of `save_pretrained` causes config.unet = (None, None).
-            # So config must be saved again, later.
-            model.save_pretrained(model_save_dir)
-
-        model.vae = vae
-        model.text_encoder = text_encoder
-        model.unet = unet
-        model.text_encoder_2 = text_encoder_2
-        update_dict = {
-            "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
-            "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-            "text_encoder_2": ("optimum.rbln", "RBLNCLIPTextModelWithProjection"),
-        }
-        model.register_to_config(**update_dict)
-
-        if model_save_dir is not None:
-            # overwrite to replace incorrect config
-            model.save_config(model_save_dir)
-
-        model.models = [vae.model[0], vae.model[1], unet.model[0], text_encoder.model[0], text_encoder_2.model[0]]
-
-        if optimize_host_memory is False:
-            model.compiled_models = [
-                vae.compiled_models[0],
-                vae.compiled_models[1],
-                unet.compiled_models[0],
-                text_encoder.compiled_models[0],
-                text_encoder_2.compiled_models[0],
-            ]
-
-        return model
+class RBLNStableDiffusionXLImg2ImgPipeline(RBLNDiffusionMixin, StableDiffusionXLImg2ImgPipeline):
+    original_class = StableDiffusionXLImg2ImgPipeline
+    _submodules = ["text_encoder", "text_encoder_2", "unet", "vae"]

optimum/rbln/modeling_alias.py

@@ -28,7 +28,6 @@ from .modeling_base import (
     RBLNModelForQuestionAnswering,
     RBLNModelForSequenceClassification,
 )
-from .modeling_seq2seq import RBLNModelForSeq2SeqLM
 
 
 class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
@@ -47,14 +46,6 @@ class RBLNResNetForImageClassification(RBLNModelForImageClassification):
     pass
 
 
-class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    pass
-
-
-class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    pass
-
-
 class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
@@ -65,3 +56,7 @@ class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
 
 class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
     rbln_model_input_names = ["input_ids", "attention_mask"]
+
+
+class RBLNViTForImageClassification(RBLNModelForImageClassification):
+    pass

optimum/rbln/modeling_base.py

@@ -21,7 +21,6 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-import copy
 import importlib
 import inspect
 import logging
@@ -50,9 +49,9 @@ from transformers import (
     PretrainedConfig,
 )
 
-from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from .modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig, use_rbln_config
 from .utils.runtime_utils import UnavailableRuntime
-from .utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
+from .utils.save_utils import maybe_load_preprocessors
 
 
 if TYPE_CHECKING:
@@ -86,13 +85,11 @@ class SubModulesMixin:
             setattr(self, submodule_meta["name"], submodule)
 
     @classmethod
-    def _from_model(
+    def _export_submodules_from_model(
         cls,
         model: "PreTrainedModel",
         model_save_dir: str,
-        rbln_sub_configs_dict: Dict[str, Any],
         rbln_kwargs: Dict[str, Any],
-        subfolder=None,  # warning: will be ignored
         **kwargs,
     ) -> List["RBLNBaseModel"]:
         rbln_submodules = []
@@ -100,18 +97,15 @@ class SubModulesMixin:
             submodule_name = submodule["name"]
             torch_submodule: "PreTrainedModel" = getattr(model, submodule["name"])
             cls_name = torch_submodule.__class__.__name__
-            submodule_cls: "RBLNBaseModel" = getattr(importlib.import_module("optimum.rbln"), f"RBLN{cls_name}")
+            submodule_cls: "RBLNModel" = getattr(importlib.import_module("optimum.rbln"), f"RBLN{cls_name}")
 
-            if submodule_name in rbln_sub_configs_dict:
-                kwargs["rbln_config"] = rbln_sub_configs_dict[submodule_name]
+            if submodule_name in rbln_kwargs:
+                kwargs["rbln_config"] = rbln_kwargs[submodule_name]
 
-            rbln_submodule = submodule_cls._export(
-                model_id=None,
-                config=torch_submodule.config,
+            rbln_submodule = submodule_cls.from_model(
+                model=torch_submodule,
                 subfolder=submodule_name,
                 model_save_dir=model_save_dir,
-                model=torch_submodule,
-                **rbln_kwargs,
                 **kwargs,
             )
 
@@ -120,26 +114,29 @@ class SubModulesMixin:
         return rbln_submodules
 
     @classmethod
-    def _submodule_from_compiled_model(
-        cls, model_save_dir: str, rbln_sub_configs_dict: Dict[str, Any], rbln_kwargs: Dict[str, Any], **kwargs
+    def _load_submodules_from_compiled_models(
+        cls,
+        model_save_dir: str,
+        rbln_kwargs: Dict[str, Any],
+        **kwargs,
     ):
         rbln_submodules = []
         for submodule in cls._rbln_submodules:
             submodule_name = submodule["name"]
-            rbln_submodule_config_dict = rbln_sub_configs_dict.get(submodule_name, None)
+
+            if submodule_name in rbln_kwargs:
+                kwargs["rbln_config"] = rbln_kwargs[submodule_name]
 
             # Get cls name for call the constructor of the rbln class
             submodule_rbln_config = RBLNConfig.load(Path(model_save_dir) / submodule_name)
             submodule_cls_name = submodule_rbln_config.meta["cls"]
             submodule_cls: "RBLNBaseModel" = getattr(importlib.import_module("optimum.rbln"), submodule_cls_name)
 
-            config = OptimizedModel._load_config(Path(model_save_dir) / submodule_name, **kwargs)
+            config = OptimizedModel._load_config(Path(model_save_dir) / submodule_name)
             rbln_submodule = submodule_cls._from_pretrained(
                 model_id=model_save_dir,
                 config=config,
                 subfolder=submodule_name,
-                rbln_config=rbln_submodule_config_dict,
-                **rbln_kwargs,
                 **kwargs,
             )
             rbln_submodules.append(rbln_submodule)
@@ -149,27 +146,24 @@ class SubModulesMixin:
     def _load_submodules(
         cls,
         model_save_dir,
-        rbln_sub_configs_dict,
         rbln_kwargs,
         model=None,
         **kwargs,
     ):
-        # Two way :
+        # Two ways :
         # 1. Compile from pytorch object
         # 2. Load from compiled file
         if model is not None:
-            return cls._from_model(
+            return cls._export_submodules_from_model(
                 model=model,
                 model_save_dir=model_save_dir,
-                rbln_sub_configs_dict=rbln_sub_configs_dict,
                 rbln_kwargs=rbln_kwargs,
                 **kwargs,
             )
 
         else:
-            return cls._submodule_from_compiled_model(
+            return cls._load_submodules_from_compiled_models(
                 model_save_dir=model_save_dir,
-                rbln_sub_configs_dict=rbln_sub_configs_dict,
                 rbln_kwargs=rbln_kwargs,
                 **kwargs,
             )
@@ -370,45 +364,7 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
         return rbln_compiled_models
 
     @classmethod
-    def _split_submodule_config(cls, rbln_config_dict: Dict[str, Any] = {}) -> Dict[str, Any]:
-        # {"language_model" : {"rbln_tensor_parallel_size":4}}
-        rbln_sub_configs_dict: Dict[str, Dict[str, Any]] = {}
-
-        # Remove submodule-configs from rbln_config
-        if len(cls._rbln_submodules) > 0:
-            keys = list(rbln_config_dict.keys())
-            submodule_names = [m["name"] for m in cls._rbln_submodules]
-            for key in keys:
-                if key in submodule_names:
-                    rbln_sub_configs_dict[key] = rbln_config_dict.pop(key)
-
-        return rbln_sub_configs_dict
-
-    @classmethod
-    def resolve_rbln_config(cls, rbln_config: Union[RBLNConfig, Dict[str, Any]], kwargs):
-        if isinstance(rbln_config, RBLNConfig):
-            # Already resolved
-            return rbln_config, None
-
-        else:
-            if rbln_config is None:
-                rbln_config_dict = {}
-            else:
-                rbln_config_dict = rbln_config
-
-            rbln_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
-            rbln_sub_configs_dict = cls._split_submodule_config(rbln_config_dict)
-
-            for key in rbln_config_dict:
-                if key in rbln_kwargs:
-                    raise KeyError(f"duplicate key in both `rbln_config` and {key}")
-
-            merged_rbln_kwargs = copy.deepcopy(rbln_kwargs)
-            merged_rbln_kwargs.update(rbln_config_dict)
-
-            return (merged_rbln_kwargs, rbln_sub_configs_dict)
-
-    @classmethod
+    @use_rbln_config
     def _from_pretrained(
         cls,
         model_id: Union[str, Path],
@@ -430,7 +386,7 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
 
         if not from_export_method:
             # from compiled dir
-            rbln_kwargs, rbln_sub_configs_dict = cls.resolve_rbln_config(rbln_config, kwargs)
+            rbln_kwargs = rbln_config or {}
 
             model_path_subfolder = cls._load_compiled_model_dir(
                 model_id=model_id,
@@ -450,7 +406,6 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
             if len(cls._rbln_submodules) > 0:
                 rbln_submodules = cls._load_submodules(
                     model_save_dir=model_id,
-                    rbln_sub_configs_dict=rbln_sub_configs_dict,
                     rbln_kwargs=rbln_kwargs,
                     **kwargs,
                 )
@@ -477,9 +432,9 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
         cls,
         rbln_compiled_models: Dict[str, rebel.RBLNCompiledModel],
         rbln_config: RBLNConfig,
-        config,
-        model_save_dir: str,
-        subfolder: str,
+        config: "PretrainedConfig",
+        model_save_dir: Union[Path, str],
+        subfolder: Union[Path, str],
         rbln_submodules: List["RBLNBaseModel"] = [],
         **kwargs,
     ):
@@ -492,11 +447,18 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
         rbln_compiled_models = [rbln_compiled_models[cm_name] for cm_name in compiled_model_names]
 
         # create runtimes only if `rbln_create_runtimes` is enabled
-        models = (
-            cls._create_runtimes(rbln_compiled_models, rbln_config.device_map)
-            if rbln_config.create_runtimes
-            else UnavailableRuntime()
-        )
+        try:
+            models = (
+                cls._create_runtimes(rbln_compiled_models, rbln_config.device_map)
+                if rbln_config.create_runtimes
+                else UnavailableRuntime()
+            )
+
+        except rebel.core.exception.RBLNRuntimeError as e:
+            logger.warning(
+                f"Failed to create the runtime for the model due to a runtime error: {e.__class__.__name__} - {e}"
+            )
+            models = UnavailableRuntime()
 
         return cls(
             models,
@@ -538,25 +500,10 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
         rbln_config = cls._get_rbln_config(**others, rbln_kwargs=rbln_kwargs)
         return rbln_config
 
-    @staticmethod
-    def pop_rbln_kwargs_from_kwargs(kwargs: Dict[str, Any], runtime_only=False):
-        keys = list(kwargs.keys())
-        rbln_kwargs = {key[5:]: kwargs.pop(key) for key in keys if key.startswith("rbln_")}
-
-        if runtime_only:
-            rbln_kwargs = {
-                key: value
-                for key, value in rbln_kwargs.items()
-                if key in {"create_runtimes", "optimize_host_memory", "device", "device_map"}
-            }
-
-        return rbln_kwargs
-
     def can_generate(self):
         return False
 
     def to(self, *args, **kwargs):
-        # Do nothing
         return self
 
     def __call__(self, *args, **kwargs):
@@ -599,12 +546,49 @@ class RBLNBaseModel(OptimizedModel, ABC, SubModulesMixin):
 
     @classmethod
     @abstractmethod
-    def _export(cls, *args, **kwargs):
-        """
-        Exports a vanilla Transformers model into a rbln-compiled Module.
-        """
+    def get_pytorch_model(cls, *args, **kwargs):
+        pass
+
+    @classmethod
+    @abstractmethod
+    @use_rbln_config
+    def from_model(
+        cls,
+        model: "PreTrainedModel",
+        rbln_config: Dict[str, Any] = {},
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
+        subfolder: str = "",
+        **kwargs,
+    ):
         pass
 
+    @classmethod
+    @use_rbln_config
+    def _export(
+        cls,
+        model_id: Union[str, Path],
+        config: "PretrainedConfig",  # FIXME : optimum passes config, but we ignore it.
+        rbln_config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> "RBLNModel":
+        subfolder = kwargs.get("subfolder", "")
+        model_save_dir = kwargs.pop("model_save_dir", None)
+
+        rbln_kwargs = rbln_config
+        model: "PreTrainedModel" = cls.get_pytorch_model(
+            model_id=model_id,
+            rbln_kwargs=rbln_kwargs,
+            **kwargs,
+        )
+        preprocessors = maybe_load_preprocessors(model_id, subfolder=subfolder)
+        return cls.from_model(
+            model,
+            rbln_config=rbln_config,
+            preprocessors=preprocessors,
+            model_save_dir=model_save_dir,
+            **kwargs,
+        )
+
 
 class RBLNModel(RBLNBaseModel):
     """
@@ -695,25 +679,19 @@ class RBLNModel(RBLNBaseModel):
         return compiled_model
 
     @classmethod
-    @torch.no_grad()
-    def _export(
+    @use_rbln_config
+    def from_model(
         cls,
-        model_id: Union[str, Path],
-        config: "PretrainedConfig",
-        use_auth_token: Optional[Union[bool, str]] = None,
-        revision: Optional[str] = None,
-        force_download: bool = False,
-        cache_dir: Optional[str] = None,
-        subfolder: str = "",
-        local_files_only: bool = False,
-        trust_remote_code: bool = False,
+        model: "PreTrainedModel",
+        rbln_config: Dict[str, Any] = {},
         model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
-        model: "PreTrainedModel" = None,
-        rbln_config: Optional[Dict[str, Any]] = None,
+        subfolder: str = "",
         **kwargs,
-    ) -> "RBLNModel":
-        rbln_kwargs, rbln_sub_configs_dict = cls.resolve_rbln_config(rbln_config, kwargs)
+    ):
+        preprocessors = kwargs.pop("preprocessors", [])
+        rbln_kwargs = rbln_config
 
+        # Directory to save compile artifacts(.rbln) and original configs
         if model_save_dir is None:
             save_dir = TemporaryDirectory()
             save_dir_path = Path(save_dir.name)
@@ -725,43 +703,31 @@ class RBLNModel(RBLNBaseModel):
                 save_dir_path = Path(model_save_dir)
                 save_dir_path.mkdir(exist_ok=True)
 
-        # Load pytorch model if needed.
-        if model is None:
-            model: "PreTrainedModel" = cls.get_pytorch_model(
-                model_id=model_id,
-                subfolder=subfolder,
-                revision=revision,
-                cache_dir=cache_dir,
-                use_auth_token=use_auth_token,
-                local_files_only=local_files_only,
-                force_download=force_download,
-                trust_remote_code=trust_remote_code,
-                rbln_kwargs=rbln_kwargs,
-                **kwargs,
-            )
-            preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
-        else:
-            preprocessors = []
+        # (Optional) Save preprocessors (tokenizer, image preprocessors, etc)
+        for preprocessor in preprocessors:
+            preprocessor.save_pretrained(save_dir_path)
 
-        # FIXME :: optimum passes AutoConfig.
+        # Save configs
+        # FIXME :: optimum passes AutoConfig. But here we ignore it.
         config = model.config
         if hasattr(model, "can_generate") and model.can_generate():
             generation_config = model.generation_config
             generation_config.save_pretrained(save_dir_path / subfolder)
-
         if not isinstance(config, PretrainedConfig):  # diffusers config
             config = PretrainedConfig(**config)
         config.save_pretrained(save_dir_path / subfolder)
 
-        # Get compilation arguments
+        # 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.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
         )
 
-        # Save compiled models
+        # Save compiled models (.rbln)
         (save_dir_path / subfolder).mkdir(exist_ok=True)
         if not isinstance(compiled_model, dict):
             compiled_models = {DEFAULT_COMPILED_MODEL_NAME: compiled_model}
@@ -771,6 +737,7 @@ class RBLNModel(RBLNBaseModel):
             cm.save(save_dir_path / subfolder / f"{compiled_model_name}.rbln")
         rbln_config.save(save_dir_path / subfolder)
 
+        # 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
@@ -778,7 +745,6 @@ class RBLNModel(RBLNBaseModel):
             rbln_submodules = cls._load_submodules(
                 model=model,
                 model_save_dir=save_dir,
-                rbln_sub_configs_dict=rbln_sub_configs_dict,
                 rbln_kwargs=rbln_kwargs,
                 **kwargs,
             )
@@ -790,12 +756,7 @@ class RBLNModel(RBLNBaseModel):
             model_id=save_dir_path,
             config=config,
             model_save_dir=save_dir,
-            use_auth_token=use_auth_token,
-            revision=revision,
-            force_download=force_download,
-            cache_dir=cache_dir,
             subfolder=subfolder,
-            local_files_only=local_files_only,
             rbln_config=rbln_config,
             rbln_compiled_models=compiled_models,
             rbln_submodules=rbln_submodules,
@@ -817,7 +778,6 @@ class RBLNModel(RBLNBaseModel):
 
 
 class RBLNModelForQuestionAnswering(RBLNModel):
-    model_type = "rbln_model"
     auto_model_class = AutoModelForQuestionAnswering
     rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
 
@@ -878,7 +838,6 @@ class RBLNModelForImageClassification(RBLNModel):
     This is a generic model class that will be instantiated as one of the model classes of the library (with a image classification head) when created with the from_pretrained() class method
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForImageClassification
 
     @classmethod
@@ -894,8 +853,17 @@ class RBLNModelForImageClassification(RBLNModel):
         if rbln_image_size is None:
             for processor in preprocessors:
                 if hasattr(processor, "size"):
-                    rbln_image_size = processor.size["shortest_edge"]
+                    if all(required_key in processor.size.keys() for required_key in ["height", "width"]):
+                        rbln_image_size = (processor.size["height"], processor.size["width"])
+                    elif "shortest_edge" in processor.size.keys():
+                        rbln_image_size = (processor.size["shortest_edge"], processor.size["shortest_edge"])
+                    elif "longest_edge" in processor.size.keys():
+                        rbln_image_size = (processor.size["longest_edge"], processor.size["longest_edge"])
                     break
+
+            if rbln_image_size is None:
+                rbln_image_size = model_config.image_size
+
             if rbln_image_size is None:
                 raise ValueError("`rbln_image_size` should be specified!")
 
@@ -903,12 +871,20 @@ class RBLNModelForImageClassification(RBLNModel):
             rbln_batch_size = 1
 
         if isinstance(rbln_image_size, int):
-            rbln_image_size = rbln_image_size, rbln_image_size
+            rbln_image_height, rbln_image_width = rbln_image_size, rbln_image_size
+        elif isinstance(rbln_image_size, (list, tuple)):
+            rbln_image_height, rbln_image_width = rbln_image_size[0], rbln_image_size[1]
+        elif isinstance(rbln_image_size, dict):
+            rbln_image_height, rbln_image_width = rbln_image_size["height"], rbln_image_size["width"]
+        else:
+            raise ValueError(
+                "`rbln_image_size` should be `int` (ex. 224), `tuple` (ex. 224, 224), `dict` (ex. {'height': 224, 'width': 224}) format"
+            )
 
         input_info = [
             (
                 "pixel_values",
-                [rbln_batch_size, 3, rbln_image_size[0], rbln_image_size[1]],
+                [rbln_batch_size, 3, rbln_image_height, rbln_image_width],
                 "float32",
             )
         ]
@@ -930,7 +906,6 @@ class RBLNModelForAudioClassification(RBLNModel):
     Currently, this model class only supports the 'AST' model from the transformers library. Future updates may include support for additional model types.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForAudioClassification
 
     @classmethod
@@ -1005,7 +980,6 @@ class RBLNModelForSequenceClassification(RBLNModel):
     Currently, this model class supports the 'XLMRoberta' and 'Roberta' model from the transformers library. Future updates may include support for additional model types.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForSequenceClassification
 
     @classmethod
@@ -1070,7 +1044,6 @@ class RBLNModelForSequenceClassification(RBLNModel):
 
 
 class RBLNModelForMaskedLM(RBLNModel):
-    model_type = "rbln_model"
     auto_model_class = AutoModelForMaskedLM
 
     @classmethod