optimum-rbln 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

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

@@ -22,17 +22,17 @@
 # from Rebellions Inc.
 """RBLNStableDiffusionXLPipeline class for inference of diffusion models on rbln devices."""
 
-from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
-from diffusers import StableDiffusionXLControlNetImg2ImgPipeline
+from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline
 from diffusers.image_processor import PipelineImageInput
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
 from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module
+from transformers import CLIPTextModel
 
 from ....modeling_base import RBLNBaseModel
 from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
@@ -63,103 +63,152 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
         """
         export = kwargs.pop("export", None)
-        text_encoder = kwargs.pop("text_encoder", None)
-        controlnets = kwargs.pop("controlnet", None)
         vae = kwargs.pop("vae", None)
+        unet = kwargs.pop("unet", None)
+        text_encoder = kwargs.pop("text_encoder", None)
+        text_encoder_2 = kwargs.pop("text_encoder_2", None)
+        controlnet = kwargs.pop("controlnet", None)
+        model_save_dir = kwargs.pop("model_save_dir", None)
 
         rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
+
         kwargs_dict = {
             "pretrained_model_name_or_path": model_id,
-            "vae": vae,
-            "controlnet": controlnets,
-            "text_encoder": text_encoder,
             **kwargs,
         }
 
+        kwargs_dict.update(
+            {
+                **({"vae": vae} if vae is not None and isinstance(vae, AutoencoderKL) else {}),
+                **({"unet": unet} if unet is not None and isinstance(unet, UNet2DConditionModel) else {}),
+                **(
+                    {"text_encoder": text_encoder}
+                    if text_encoder is not None and isinstance(text_encoder, CLIPTextModel)
+                    else {}
+                ),
+                **(
+                    {"controlnet": controlnet}
+                    if controlnet is not None
+                    and (
+                        isinstance(controlnet, ControlNetModel)
+                        or all(isinstance(c, ControlNetModel) for c in controlnet)
+                    )
+                    else {}
+                ),
+            }
+        )
+
         model = super().from_pretrained(**{k: v for k, v in kwargs_dict.items() if v is not None})
 
         if export is None or export is False:
             return model
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
-
-        model.save_pretrained(save_directory=save_dir_path, **kwargs)
-
         do_classifier_free_guidance = (
             rbln_config_kwargs.pop("rbln_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,
-            rbln_unet_sample_size=model.unet.config.sample_size,
-            rbln_use_encode=True,
-            rbln_vae_scale_factor=model.vae_scale_factor,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
-        )
-        text_encoder = RBLNCLIPTextModel.from_pretrained(
-            model_id=model_id,
-            subfolder="text_encoder",
-            export=True,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
-        )
-        text_encoder_2 = RBLNCLIPTextModelWithProjection.from_pretrained(
-            model_id=model_id,
-            subfolder="text_encoder_2",
-            export=True,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
-        )
-
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
+        if not isinstance(vae, RBLNAutoencoderKL):
+            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_kwargs,
+                **rbln_constructor_kwargs,
+            )
 
-        unet = RBLNUNet2DConditionModel.from_pretrained(
-            model_id=model_id,
-            subfolder="unet",
-            export=True,
-            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_kwargs,
-            **rbln_constructor_kwargs,
-        )
+        if not isinstance(text_encoder, RBLNCLIPTextModel):
+            text_encoder = RBLNCLIPTextModel.from_pretrained(
+                model_id=model_id,
+                subfolder="text_encoder",
+                export=True,
+                model_save_dir=model_save_dir,
+                **rbln_config_kwargs,
+                **rbln_constructor_kwargs,
+            )
 
-        if isinstance(controlnets, (list, tuple)):
-            controlnet = RBLNMultiControlNetModel.from_pretrained(
-                model_id=str(save_dir_path / "controlnet"),
+        if not isinstance(text_encoder_2, RBLNCLIPTextModel):
+            text_encoder_2 = RBLNCLIPTextModelWithProjection.from_pretrained(
+                model_id=model_id,
+                subfolder="text_encoder_2",
                 export=True,
-                rbln_batch_size=unet_batch_size,
-                rbln_vae_scale_factor=model.vae_scale_factor,
+                model_save_dir=model_save_dir,
                 **rbln_config_kwargs,
                 **rbln_constructor_kwargs,
             )
-            controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
-        else:
-            controlnet = RBLNControlNetModel.from_pretrained(
-                model_id=save_dir_path / "controlnet",
+
+        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
+        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
+
+        if not isinstance(unet, RBLNUNet2DConditionModel):
+            unet = RBLNUNet2DConditionModel.from_pretrained(
+                model_id=model_id,
+                subfolder="unet",
                 export=True,
-                rbln_batch_size=unet_batch_size,
+                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_kwargs,
                 **rbln_constructor_kwargs,
             )
-            controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
+        if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
+            if isinstance(controlnet, (list, tuple)):
+                multicontrolnet = []
+                for i, cid in enumerate(controlnet):
+                    subfolder_name = "controlnet" if i == 0 else f"controlnet_{i}"
+                    multicontrolnet.append(
+                        RBLNControlNetModel.from_pretrained(
+                            model_id=cid.config._name_or_path,
+                            subfolder=subfolder_name,
+                            export=True,
+                            model_save_dir=model_save_dir,
+                            rbln_batch_size=unet_batch_size,
+                            rbln_text_model_hidden_size=model.text_encoder_2.config.hidden_size,
+                            rbln_vae_scale_factor=model.vae_scale_factor,
+                            **rbln_config_kwargs,
+                            **rbln_constructor_kwargs,
+                        )
+                    )
+                controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
+                controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
+            else:
+                controlnet = RBLNControlNetModel.from_pretrained(
+                    model_id=controlnet.config._name_or_path,
+                    subfolder="controlnet",
+                    export=True,
+                    model_save_dir=model_save_dir,
+                    rbln_batch_size=unet_batch_size,
+                    rbln_text_model_hidden_size=model.text_encoder_2.config.hidden_size,
+                    rbln_vae_scale_factor=model.vae_scale_factor,
+                    **rbln_config_kwargs,
+                    **rbln_constructor_kwargs,
+                )
+                controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
+
+        if model_save_dir is not None:
+            # To skip saving original pytorch modules
+            del (model.vae, model.text_encoder, model.unet, model.controlnet)
+
+            # Direct calling of `save_pretrained` causes config.unet = (None, None).
+            # So config must be saved again, later.
+            model.save_pretrained(model_save_dir)
+
+        # replace modules
         model.vae = vae
         model.text_encoder = text_encoder
         model.unet = unet
         model.text_encoder_2 = text_encoder_2
         model.controlnet = controlnet
 
+        # update config to be able to load from file
         update_dict = {
             "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
             "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
@@ -169,14 +218,24 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
         }
         model.register_to_config(**update_dict)
 
-        model.models = [
-            vae.model[0],
-            vae.model[1],
-            unet.model[0],
-            text_encoder.model[0],
-            text_encoder_2.model[0],
-            controlnet.model[0],
-        ]
+        if model_save_dir is not None:
+            # overwrite to replace incorrect config
+            model.save_config(model_save_dir)
+
+        # use for CI to access each compiled model
+        if rbln_constructor_kwargs.pop("rbln_optimize_host_memory", None) is False:
+            model.compiled_models = [
+                vae.compiled_models[0],
+                vae.compiled_models[1],
+                text_encoder.compiled_models[0],
+                text_encoder_2.compiled_models[0],
+                unet.compiled_models[0],
+            ]
+            if isinstance(controlnet, RBLNMultiControlNetModel):
+                for c_model in controlnet.nets:
+                    model.compiled_models.append(c_model.compiled_models[0])
+            else:
+                model.compiled_models.append(controlnet.compiled_models[0])
 
         return model
 

optimum/rbln/modeling_alias.py

@@ -24,7 +24,9 @@
 from .modeling_base import (
     RBLNModelForAudioClassification,
     RBLNModelForImageClassification,
+    RBLNModelForMaskedLM,
     RBLNModelForQuestionAnswering,
+    RBLNModelForSequenceClassification,
 )
 from .modeling_seq2seq import RBLNModelForSeq2SeqLM
 
@@ -34,7 +36,11 @@ class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
 
 
 class RBLNBertForQuestionAnswering(RBLNModelForQuestionAnswering):
-    pass
+    rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
+
+
+class RBLNDistilBertForQuestionAnswering(RBLNModelForQuestionAnswering):
+    rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
 class RBLNResNetForImageClassification(RBLNModelForImageClassification):
@@ -47,3 +53,15 @@ class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
 
 class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
     pass
+
+
+class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
+    pass
+
+
+class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
+    pass
+
+
+class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
+    pass

optimum/rbln/modeling_base.py

@@ -39,7 +39,9 @@ from transformers import (
     AutoModel,
     AutoModelForAudioClassification,
     AutoModelForImageClassification,
+    AutoModelForMaskedLM,
     AutoModelForQuestionAnswering,
+    AutoModelForSequenceClassification,
     GenerationConfig,
     PretrainedConfig,
 )
@@ -49,10 +51,15 @@ from .utils.runtime_utils import UnavailableRuntime
 from .utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
 
 
-logger = logging.getLogger(__name__)
-
 if TYPE_CHECKING:
-    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer, PreTrainedModel
+    from transformers import (
+        AutoFeatureExtractor,
+        AutoProcessor,
+        AutoTokenizer,
+        PreTrainedModel,
+    )
+
+logger = logging.getLogger(__name__)
 
 
 class RBLNBaseModel(OptimizedModel, ABC):
@@ -154,13 +161,23 @@ class RBLNBaseModel(OptimizedModel, ABC):
                 Directory where to save the model file.
         """
         real_save_dir = self.model_save_dir / self.subfolder
+        save_directory_path = Path(save_directory)
         if os.path.exists(real_save_dir) and os.path.isdir(real_save_dir):
+            if save_directory_path.absolute() == real_save_dir.absolute():
+                raise FileExistsError(
+                    f"Cannot save model to '{save_directory}'. "
+                    f"This directory already exists and contains the model files."
+                )
             shutil.copytree(real_save_dir, save_directory, dirs_exist_ok=True)
             self.config.save_pretrained(save_directory)
             if self.generation_config is not None:
                 self.generation_config.save_pretrained(save_directory)
         else:
-            raise FileNotFoundError(f"Saving compiled model failed.({real_save_dir}).")
+            raise FileNotFoundError(
+                f"Unable to save the model. The model directory '{real_save_dir}' does not exist or is not accessible. "
+                f"Cannot save to the specified destination '{save_directory}'. "
+                f"Please ensure the model directory exists and you have the necessary permissions to access it."
+            )
 
     @classmethod
     def _from_pretrained(
@@ -194,7 +211,12 @@ class RBLNBaseModel(OptimizedModel, ABC):
                 token = HfFolder().get_token()
             else:
                 token = use_auth_token
-            repo_files = list(map(Path, HfApi().list_repo_files(model_id, revision=revision, token=token)))
+            repo_files = list(
+                map(
+                    Path,
+                    HfApi().list_repo_files(model_id, revision=revision, token=token),
+                )
+            )
 
             pattern = "*.rbln" if subfolder == "" else f"{subfolder}/*.rbln"
             rbln_files = [p for p in repo_files if p.match(pattern)]
@@ -285,7 +307,7 @@ class RBLNBaseModel(OptimizedModel, ABC):
             preprocessors,
             model_save_dir=model_save_dir,
             subfolder=subfolder,
-            rbln_compiled_models=None if rbln_optimize_host_memory else rbln_compiled_models,
+            rbln_compiled_models=(None if rbln_optimize_host_memory else rbln_compiled_models),
             **kwargs,
         )
 
@@ -375,7 +397,7 @@ class RBLNBaseModel(OptimizedModel, ABC):
         return self.forward(*args, **kwargs)
 
     @classmethod
-    def wrap_model_if_needed(cls, model: torch.nn.Module) -> torch.nn.Module:
+    def wrap_model_if_needed(cls, model: torch.nn.Module, rbln_config: RBLNConfig) -> torch.nn.Module:
         # Wrap the model if needed.
         return model
 
@@ -398,7 +420,9 @@ class RBLNBaseModel(OptimizedModel, ABC):
     @classmethod
     @abstractmethod
     def _create_runtimes(
-        cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_device_map: Dict[str, int],
     ) -> List[rebel.Runtime]:
         # compiled_models -> runtimes
         pass
@@ -495,7 +519,7 @@ class RBLNModel(RBLNBaseModel):
 
     @classmethod
     def get_compiled_model(cls, model: "PreTrainedModel", rbln_config: RBLNConfig):
-        model = cls.wrap_model_if_needed(model)
+        model = cls.wrap_model_if_needed(model, rbln_config)
         rbln_runtime_configs = list(rbln_config.values())
         if len(rbln_runtime_configs) != 1:
             raise ValueError
@@ -596,7 +620,9 @@ class RBLNModel(RBLNBaseModel):
 
     @classmethod
     def _create_runtimes(
-        cls, compiled_models: List[rebel.RBLNCompiledModel], rbln_device_map: Dict[str, int]
+        cls,
+        compiled_models: List[rebel.RBLNCompiledModel],
+        rbln_device_map: Dict[str, int],
     ) -> List[rebel.Runtime]:
         device = rbln_device_map[DEFAULT_COMPILED_MODEL_NAME]
         return [compiled_model.create_runtime(tensor_type="pt", device=device) for compiled_model in compiled_models]
@@ -616,8 +642,8 @@ class RBLNModelForQuestionAnswering(RBLNModel):
         preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
         model_config: Optional["PretrainedConfig"] = None,
         rbln_max_seq_len: Optional[int] = None,
-        rbln_model_input_names: Optional[List[str]] = None,
         rbln_batch_size: Optional[int] = None,
+        rbln_model_input_names: Optional[List[str]] = None,
     ) -> RBLNConfig:
         if rbln_max_seq_len is None:
             for tokenizer in preprocessors:
@@ -627,15 +653,15 @@ class RBLNModelForQuestionAnswering(RBLNModel):
             if rbln_max_seq_len is None:
                 raise ValueError("`rbln_max_seq_len` should be specified!")
 
-        if rbln_model_input_names is None:
-            # These are BERT's inputs
-            rbln_model_input_names = ["input_ids", "attention_mask", "token_type_ids"]
-
         if rbln_batch_size is None:
             rbln_batch_size = 1
+
+        if rbln_model_input_names is not None:
+            cls.rbln_model_input_names = rbln_model_input_names
+
         input_info = [
             (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
-            for model_input_name in rbln_model_input_names
+            for model_input_name in cls.rbln_model_input_names
         ]
 
         rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
@@ -672,7 +698,13 @@ class RBLNModelForImageClassification(RBLNModel):
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
-        input_info = [("pixel_values", [rbln_batch_size, 3, rbln_image_size, rbln_image_size], "float32")]
+        input_info = [
+            (
+                "pixel_values",
+                [rbln_batch_size, 3, rbln_image_size, rbln_image_size],
+                "float32",
+            )
+        ]
 
         rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
         rbln_runtime_config.batch_size = rbln_batch_size
@@ -737,7 +769,11 @@ class RBLNModelForAudioClassification(RBLNModel):
         meta["rbln_num_mel_bins"] = rbln_num_mel_bins
 
         model_input_info = [
-            ("input_values", [rbln_batch_size, rbln_max_length, rbln_num_mel_bins], "float32"),
+            (
+                "input_values",
+                [rbln_batch_size, rbln_max_length, rbln_num_mel_bins],
+                "float32",
+            ),
         ]
 
         rbln_runtime_config = RBLNRuntimeConfig(input_info=model_input_info, batch_size=rbln_batch_size)
@@ -748,3 +784,111 @@ class RBLNModelForAudioClassification(RBLNModel):
         )
 
         return rbln_config
+
+
+class RBLNModelForSequenceClassification(RBLNModel):
+    """
+    This is a generic model class that will be instantiated as one of the model classes of the library (with a sequence classification head) when created with the from_pretrained() class method
+    This model inherits from [`RBLNModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+
+    A class to convert and run pre-trained transformers based SequenceClassification models on RBLN devices.
+    It implements the methods to convert a pre-trained transformers SequenceClassification model into a RBLN transformer model by:
+    - transferring the checkpoint weights of the original into an optimized RBLN graph,
+    - compiling the resulting graph using the RBLN compiler.
+
+    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
+    def _get_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_max_seq_len: Optional[int] = None,
+        rbln_model_input_names: Optional[List[str]] = None,
+        rbln_batch_size: Optional[int] = None,
+    ) -> RBLNConfig:
+        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
+            model_config, "max_position_embeddings", None
+        )
+
+        if rbln_max_seq_len is None:
+            rbln_max_seq_len = max_position_embeddings
+            if rbln_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_max_seq_len is None:
+                    raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_model_input_names is None:
+            # These are BERT's inputs
+            rbln_model_input_names = ["input_ids", "attention_mask"]
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+        input_info = [
+            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
+            for model_input_name in rbln_model_input_names
+        ]
+
+        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
+        rbln_runtime_config.batch_size = rbln_batch_size
+        meta = {"rbln_max_seq_len": rbln_max_seq_len}
+
+        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+
+
+class RBLNModelForMaskedLM(RBLNModel):
+    model_type = "rbln_model"
+    auto_model_class = AutoModelForMaskedLM
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Optional[Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"]],
+        model_config: Optional["PretrainedConfig"] = None,
+        rbln_max_seq_len: Optional[int] = None,
+        rbln_model_input_names: Optional[List[str]] = None,
+        rbln_batch_size: Optional[int] = None,
+    ) -> RBLNConfig:
+        max_position_embeddings = getattr(model_config, "n_positions", None) or getattr(
+            model_config, "max_position_embeddings", None
+        )
+
+        if rbln_max_seq_len is None:
+            rbln_max_seq_len = max_position_embeddings
+            if rbln_max_seq_len is None:
+                for tokenizer in preprocessors:
+                    if hasattr(tokenizer, "model_max_length"):
+                        rbln_max_seq_len = tokenizer.model_max_length
+                        break
+                if rbln_max_seq_len is None:
+                    raise ValueError("`rbln_max_seq_len` should be specified!")
+
+        if max_position_embeddings is not None and rbln_max_seq_len > max_position_embeddings:
+            raise ValueError("`rbln_enc_max_seq_len` should be less or equal than max_position_embeddings!")
+
+        if rbln_model_input_names is None:
+            # These are BERT's inputs
+            rbln_model_input_names = ["input_ids", "attention_mask"]
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+        input_info = [
+            (model_input_name, [rbln_batch_size, rbln_max_seq_len], "int64")
+            for model_input_name in rbln_model_input_names
+        ]
+
+        rbln_runtime_config = RBLNRuntimeConfig(input_info=input_info)
+        rbln_runtime_config.batch_size = rbln_batch_size
+        meta = {"rbln_max_seq_len": rbln_max_seq_len}
+
+        return RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)

optimum/rbln/transformers/__init__.py

@@ -27,30 +27,38 @@ from transformers.utils import _LazyModule
 
 
 _import_structure = {
+    "cache_utils": ["RebelDynamicCache"],
     "generation": ["BatchTextIteratorStreamer"],
     "models": [
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
         "RBLNDPTForDepthEstimation",
+        "RBLNGemmaForCausalLM",
         "RBLNGPT2LMHeadModel",
         "RBLNWav2Vec2ForCTC",
         "RBLNWhisperForConditionalGeneration",
         "RBLNLlamaForCausalLM",
         "RBLNMidmLMHeadModel",
+        "RBLNMistralForCausalLM",
+        "RBLNXLMRobertaModel",
     ],
 }
 
 if TYPE_CHECKING:
+    from .cache_utils import RebelDynamicCache
     from .generation import BatchTextIteratorStreamer
     from .models import (
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
         RBLNDPTForDepthEstimation,
+        RBLNGemmaForCausalLM,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
         RBLNMidmLMHeadModel,
+        RBLNMistralForCausalLM,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
+        RBLNXLMRobertaModel,
     )
 else:
     import sys