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

optimum/rbln/__init__.py

@@ -38,7 +38,7 @@ _import_structure = {
         "RBLNXLMRobertaForSequenceClassification",
         "RBLNRobertaForSequenceClassification",
         "RBLNRobertaForMaskedLM",
-        "RBLNViTForImageClassification"
+        "RBLNViTForImageClassification",
     ],
     "modeling_base": [
         "RBLNBaseModel",
@@ -76,6 +76,7 @@ _import_structure = {
         "RBLNQwen2ForCausalLM",
         "RBLNWav2Vec2ForCTC",
         "RBLNLlamaForCausalLM",
+        "RBLNT5EncoderModel",
         "RBLNT5ForConditionalGeneration",
         "RBLNPhiForCausalLM",
         "RBLNLlavaNextForConditionalGeneration",
@@ -99,6 +100,7 @@ _import_structure = {
         "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
     ],
     "modeling_config": ["RBLNCompileConfig", "RBLNConfig"],
+    "modeling_diffusers": ["RBLNDiffusionMixin"],
 }
 
 if TYPE_CHECKING:
@@ -136,6 +138,7 @@ if TYPE_CHECKING:
         RBLNModelForSequenceClassification,
     )
     from .modeling_config import RBLNCompileConfig, RBLNConfig
+    from .modeling_diffusers import RBLNDiffusionMixin
     from .transformers import (
         BatchTextIteratorStreamer,
         RBLNAutoModel,
@@ -166,6 +169,7 @@ if TYPE_CHECKING:
         RBLNMistralForCausalLM,
         RBLNPhiForCausalLM,
         RBLNQwen2ForCausalLM,
+        RBLNT5EncoderModel,
         RBLNT5ForConditionalGeneration,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,

optimum/rbln/__version__.py

@@ -1 +1 @@
-__version__ = '0.1.12'
+__version__ = '0.1.13'

optimum/rbln/diffusers/models/autoencoder_kl.py

@@ -22,7 +22,6 @@
 # from Rebellions Inc.
 
 import logging
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Union
 
 import rebel
@@ -30,11 +29,11 @@ import torch  # noqa: I001
 from diffusers import AutoencoderKL
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 from diffusers.models.modeling_outputs import AutoencoderKLOutput
-from optimum.exporters import TasksManager
-from transformers import AutoConfig, AutoModel, PretrainedConfig
+from transformers import PretrainedConfig
 
 from ...modeling_base import RBLNModel
 from ...modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
+from ...utils.context import override_auto_classes
 from ...utils.runtime_utils import RBLNPytorchRuntime
 
 
@@ -63,8 +62,7 @@ class RBLNAutoencoderKL(RBLNModel):
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
 
-        self.rbln_use_encode = self.rbln_config.model_cfg["use_encode"]
-        if self.rbln_use_encode:
+        if self.rbln_config.model_cfg.get("img2img_pipeline"):
             self.encoder = RBLNRuntimeVAEEncoder(runtime=self.model[0], main_input_name="x")
             self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[1], main_input_name="z")
         else:
@@ -91,38 +89,15 @@ class RBLNAutoencoderKL(RBLNModel):
 
             return dec_compiled_model
 
-        if rbln_config.model_cfg.get("use_encode", False):
+        if rbln_config.model_cfg.get("img2img_pipeline"):
             return compile_img2img()
         else:
             return compile_text2img()
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
-        def get_model_from_task(
-            task: str,
-            model_name_or_path: Union[str, Path],
-            **kwargs,
-        ):
-            return AutoencoderKL.from_pretrained(pretrained_model_name_or_path=model_name_or_path, **kwargs)
-
-        tasktmp = TasksManager.get_model_from_task
-        configtmp = AutoConfig.from_pretrained
-        modeltmp = AutoModel.from_pretrained
-        TasksManager.get_model_from_task = get_model_from_task
-
-        if kwargs.get("export", None):
-            # This is an ad-hoc to workaround save null values of the config.
-            # if export, pure optimum(not optimum-rbln) loads config using AutoConfig
-            # and diffusers model do not support loading by AutoConfig.
-            AutoConfig.from_pretrained = lambda *args, **kwargs: None
-        else:
-            AutoConfig.from_pretrained = AutoencoderKL.load_config
-
-        AutoModel.from_pretrained = AutoencoderKL.from_pretrained
-        rt = super().from_pretrained(*args, **kwargs)
-        AutoConfig.from_pretrained = configtmp
-        AutoModel.from_pretrained = modeltmp
-        TasksManager.get_model_from_task = tasktmp
+        with override_auto_classes(config_func=AutoencoderKL.load_config, model_func=AutoencoderKL.from_pretrained):
+            rt = super().from_pretrained(*args, **kwargs)
         return rt
 
     @classmethod
@@ -132,34 +107,39 @@ class RBLNAutoencoderKL(RBLNModel):
         model_config: "PretrainedConfig",
         rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        rbln_unet_sample_size = rbln_kwargs.get("unet_sample_size", None)
-        rbln_img_width = rbln_kwargs.get("img_width", None)
-        rbln_img_height = rbln_kwargs.get("img_height", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_use_encode = rbln_kwargs.get("use_encode", None)
-        rbln_vae_scale_factor = rbln_kwargs.get("vae_scale_factor", None)
+        rbln_batch_size = rbln_kwargs.get("batch_size")
+        sample_size = rbln_kwargs.get("sample_size")
 
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
-        model_cfg = {}
+        if sample_size is None:
+            sample_size = model_config.sample_size
+
+        if isinstance(sample_size, int):
+            sample_size = (sample_size, sample_size)
+
+        if hasattr(model_config, "block_out_channels"):
+            vae_scale_factor = 2 ** (len(model_config.block_out_channels) - 1)
+        else:
+            # vae image processor default value 8 (int)
+            vae_scale_factor = 8
 
-        if rbln_use_encode:
-            model_cfg["img_width"] = rbln_img_width
-            model_cfg["img_height"] = rbln_img_height
+        dec_shape = (sample_size[0] // vae_scale_factor, sample_size[1] // vae_scale_factor)
+        enc_shape = (sample_size[0], sample_size[1])
 
+        if rbln_kwargs["img2img_pipeline"]:
             vae_enc_input_info = [
-                ("x", [rbln_batch_size, model_config.in_channels, rbln_img_height, rbln_img_width], "float32")
+                (
+                    "x",
+                    [rbln_batch_size, model_config.in_channels, enc_shape[0], enc_shape[1]],
+                    "float32",
+                )
             ]
             vae_dec_input_info = [
                 (
                     "z",
-                    [
-                        rbln_batch_size,
-                        model_config.latent_channels,
-                        rbln_img_height // rbln_vae_scale_factor,
-                        rbln_img_width // rbln_vae_scale_factor,
-                    ],
+                    [rbln_batch_size, model_config.latent_channels, dec_shape[0], dec_shape[1]],
                     "float32",
                 )
             ]
@@ -173,33 +153,22 @@ class RBLNAutoencoderKL(RBLNModel):
                 compile_cfgs=compile_cfgs,
                 rbln_kwargs=rbln_kwargs,
             )
-            rbln_config.model_cfg.update(model_cfg)
             return rbln_config
 
-        if rbln_unet_sample_size is None:
-            rbln_unet_sample_size = 64
-
-        model_cfg["unet_sample_size"] = rbln_unet_sample_size
         vae_config = RBLNCompileConfig(
             input_info=[
                 (
                     "z",
-                    [
-                        rbln_batch_size,
-                        model_config.latent_channels,
-                        rbln_unet_sample_size,
-                        rbln_unet_sample_size,
-                    ],
+                    [rbln_batch_size, model_config.latent_channels, dec_shape[0], dec_shape[1]],
                     "float32",
                 )
-            ],
+            ]
         )
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
             compile_cfgs=[vae_config],
             rbln_kwargs=rbln_kwargs,
         )
-        rbln_config.model_cfg.update(model_cfg)
         return rbln_config
 
     @classmethod

optimum/rbln/diffusers/models/controlnet.py

@@ -22,16 +22,15 @@
 # from Rebellions Inc.
 
 import logging
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import torch
 from diffusers import ControlNetModel
-from optimum.exporters import TasksManager
-from transformers import AutoConfig, AutoModel, PretrainedConfig
+from transformers import PretrainedConfig
 
 from ...modeling_base import RBLNModel
 from ...modeling_config import RBLNCompileConfig, RBLNConfig
+from ...utils.context import override_auto_classes
 
 
 if TYPE_CHECKING:
@@ -113,23 +112,11 @@ class RBLNControlNetModel(RBLNModel):
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
-        def get_model_from_task(
-            task: str,
-            model_name_or_path: Union[str, Path],
-            **kwargs,
+        with override_auto_classes(
+            config_func=ControlNetModel.load_config,
+            model_func=ControlNetModel.from_pretrained,
         ):
-            return ControlNetModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path, **kwargs)
-
-        tasktmp = TasksManager.get_model_from_task
-        configtmp = AutoConfig.from_pretrained
-        modeltmp = AutoModel.from_pretrained
-        TasksManager.get_model_from_task = get_model_from_task
-        AutoConfig.from_pretrained = ControlNetModel.load_config
-        AutoModel.from_pretrained = ControlNetModel.from_pretrained
-        rt = super().from_pretrained(*args, **kwargs)
-        AutoConfig.from_pretrained = configtmp
-        AutoModel.from_pretrained = modeltmp
-        TasksManager.get_model_from_task = tasktmp
+            rt = super().from_pretrained(*args, **kwargs)
         return rt
 
     @classmethod
@@ -151,33 +138,35 @@ class RBLNControlNetModel(RBLNModel):
         model_config: "PretrainedConfig",
         rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_text_model_hidden_size = rbln_kwargs.get("text_model_hidden_size", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_img_width = rbln_kwargs.get("img_width", None)
-        rbln_img_height = rbln_kwargs.get("img_height", None)
-        rbln_vae_scale_factor = rbln_kwargs.get("vae_scale_factor", None)
+        batch_size = rbln_kwargs.get("batch_size")
+        max_seq_len = rbln_kwargs.get("max_seq_len")
+        unet_sample_size = rbln_kwargs.get("unet_sample_size")
+        vae_sample_size = rbln_kwargs.get("vae_sample_size")
 
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
+        if batch_size is None:
+            batch_size = 1
 
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = 77
+        if unet_sample_size is None:
+            raise ValueError(
+                "`rbln_unet_sample_size` (latent height, widht) must be specified (ex. unet's sample_size)"
+            )
 
-        if rbln_img_width is None or rbln_img_height is None or rbln_vae_scale_factor is None:
-            raise ValueError("rbln_img_width, rbln_img_height, and rbln_vae_scale_factor must be provided")
+        if vae_sample_size is None:
+            raise ValueError(
+                "`rbln_vae_sample_size` (input image height, width) must be specified (ex. vae's sample_size)"
+            )
 
-        input_width = rbln_img_width // rbln_vae_scale_factor
-        input_height = rbln_img_height // rbln_vae_scale_factor
+        if max_seq_len is None:
+            raise ValueError("`rbln_max_seq_len` (ex. text_encoder's max_position_embeddings )must be specified")
 
         input_info = [
             (
                 "sample",
                 [
-                    rbln_batch_size,
+                    batch_size,
                     model_config.in_channels,
-                    input_height,
-                    input_width,
+                    unet_sample_size[0],
+                    unet_sample_size[1],
                 ],
                 "float32",
             ),
@@ -189,23 +178,24 @@ class RBLNControlNetModel(RBLNModel):
             input_info.append(
                 (
                     "encoder_hidden_states",
-                    [
-                        rbln_batch_size,
-                        rbln_max_seq_len,
-                        model_config.cross_attention_dim,
-                    ],
+                    [batch_size, max_seq_len, model_config.cross_attention_dim],
                     "float32",
                 )
             )
 
-        input_info.append(("controlnet_cond", [rbln_batch_size, 3, rbln_img_height, rbln_img_width], "float32"))
+        input_info.append(
+            (
+                "controlnet_cond",
+                [batch_size, 3, vae_sample_size[0], vae_sample_size[1]],
+                "float32",
+            )
+        )
         input_info.append(("conditioning_scale", [], "float32"))
 
         if hasattr(model_config, "addition_embed_type") and model_config.addition_embed_type == "text_time":
-            if rbln_text_model_hidden_size is None:
-                rbln_text_model_hidden_size = 768
-            input_info.append(("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32"))
-            input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+            rbln_text_model_hidden_size = rbln_kwargs["text_model_hidden_size"]
+            input_info.append(("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32"))
+            input_info.append(("time_ids", [batch_size, 6], "float32"))
 
         rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
@@ -215,16 +205,6 @@ class RBLNControlNetModel(RBLNModel):
             rbln_kwargs=rbln_kwargs,
         )
 
-        rbln_config.model_cfg.update(
-            {
-                "max_seq_len": rbln_max_seq_len,
-                "batch_size": rbln_batch_size,
-                "img_width": rbln_img_width,
-                "img_height": rbln_img_height,
-                "vae_scale_factor": rbln_vae_scale_factor,
-            }
-        )
-
         return rbln_config
 
     def forward(

optimum/rbln/diffusers/models/unet_2d_condition.py

@@ -23,16 +23,15 @@
 
 import logging
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
 from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
-from optimum.exporters import TasksManager
-from transformers import AutoConfig, AutoModel, PretrainedConfig
+from transformers import PretrainedConfig
 
 from ...modeling_base import RBLNModel
 from ...modeling_config import RBLNCompileConfig, RBLNConfig
+from ...utils.context import override_auto_classes
 
 
 if TYPE_CHECKING:
@@ -143,29 +142,11 @@ class RBLNUNet2DConditionModel(RBLNModel):
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
-        def get_model_from_task(
-            task: str,
-            model_name_or_path: Union[str, Path],
-            **kwargs,
+        with override_auto_classes(
+            config_func=UNet2DConditionModel.load_config,
+            model_func=UNet2DConditionModel.from_pretrained,
         ):
-            return UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path, **kwargs)
-
-        tasktmp = TasksManager.get_model_from_task
-        configtmp = AutoConfig.from_pretrained
-        modeltmp = AutoModel.from_pretrained
-        TasksManager.get_model_from_task = get_model_from_task
-        if kwargs.get("export", None):
-            # This is an ad-hoc to workaround save null values of the config.
-            # if export, pure optimum(not optimum-rbln) loads config using AutoConfig
-            # and diffusers model do not support loading by AutoConfig.
-            AutoConfig.from_pretrained = lambda *args, **kwargs: None
-        else:
-            AutoConfig.from_pretrained = UNet2DConditionModel.load_config
-        AutoModel.from_pretrained = UNet2DConditionModel.from_pretrained
-        rt = super().from_pretrained(*args, **kwargs)
-        AutoConfig.from_pretrained = configtmp
-        AutoModel.from_pretrained = modeltmp
-        TasksManager.get_model_from_task = tasktmp
+            rt = super().from_pretrained(*args, **kwargs)
         return rt
 
     @classmethod
@@ -182,137 +163,68 @@ class RBLNUNet2DConditionModel(RBLNModel):
         model_config: "PretrainedConfig",
         rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_text_model_hidden_size = rbln_kwargs.get("text_model_hidden_size", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_in_features = rbln_kwargs.get("in_features", None)
-        rbln_use_encode = rbln_kwargs.get("use_encode", None)
-        rbln_img_width = rbln_kwargs.get("img_width", None)
-        rbln_img_height = rbln_kwargs.get("img_height", None)
-        rbln_vae_scale_factor = rbln_kwargs.get("vae_scale_factor", None)
-        rbln_is_controlnet = rbln_kwargs.get("is_controlnet", None)
-
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = 77
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        if rbln_use_encode:
-            if rbln_img_width is None or rbln_img_height is None or rbln_vae_scale_factor is None:
-                raise ValueError(
-                    "rbln_img_width, rbln_img_height, and rbln_vae_scale_factor must be provided when rbln_use_encode is True"
-                )
-            input_width = rbln_img_width // rbln_vae_scale_factor
-            input_height = rbln_img_height // rbln_vae_scale_factor
-        else:
-            input_width, input_height = model_config.sample_size, model_config.sample_size
+        batch_size = rbln_kwargs.get("batch_size")
+        max_seq_len = rbln_kwargs.get("max_seq_len")
+        sample_size = rbln_kwargs.get("sample_size")
+        is_controlnet = rbln_kwargs.get("is_controlnet")
+        rbln_in_features = None
+
+        if batch_size is None:
+            batch_size = 1
+
+        if sample_size is None:
+            sample_size = model_config.sample_size
+
+        if isinstance(sample_size, int):
+            sample_size = (sample_size, sample_size)
+
+        if max_seq_len is None:
+            raise ValueError("`rbln_max_seq_len` (ex. text_encoder's max_position_embeddings )must be specified")
 
         input_info = [
-            (
-                "sample",
-                [
-                    rbln_batch_size,
-                    model_config.in_channels,
-                    input_height,
-                    input_width,
-                ],
-                "float32",
-            ),
+            ("sample", [batch_size, model_config.in_channels, sample_size[0], sample_size[1]], "float32"),
             ("timestep", [], "float32"),
-            (
-                "encoder_hidden_states",
-                [
-                    rbln_batch_size,
-                    rbln_max_seq_len,
-                    model_config.cross_attention_dim,
-                ],
-                "float32",
-            ),
+            ("encoder_hidden_states", [batch_size, max_seq_len, model_config.cross_attention_dim], "float32"),
         ]
 
-        if rbln_is_controlnet:
-            if len(model_config.block_out_channels) > 0:
-                input_info.extend(
-                    [
-                        (
-                            f"down_block_additional_residuals_{i}",
-                            [rbln_batch_size, model_config.block_out_channels[0], input_height, input_width],
-                            "float32",
-                        )
-                        for i in range(3)
-                    ]
-                )
-            if len(model_config.block_out_channels) > 1:
-                input_info.append(
-                    (
-                        "down_block_additional_residuals_3",
-                        [rbln_batch_size, model_config.block_out_channels[0], input_height // 2, input_width // 2],
-                        "float32",
-                    )
-                )
-                input_info.extend(
-                    [
-                        (
-                            f"down_block_additional_residuals_{i}",
-                            [rbln_batch_size, model_config.block_out_channels[1], input_height // 2, input_width // 2],
-                            "float32",
-                        )
-                        for i in range(4, 6)
-                    ]
-                )
-            if len(model_config.block_out_channels) > 2:
-                input_info.append(
-                    (
-                        f"down_block_additional_residuals_{6}",
-                        [rbln_batch_size, model_config.block_out_channels[1], input_height // 4, input_width // 4],
-                        "float32",
-                    )
-                )
-                input_info.extend(
-                    [
-                        (
-                            f"down_block_additional_residuals_{i}",
-                            [rbln_batch_size, model_config.block_out_channels[2], input_height // 4, input_width // 4],
-                            "float32",
-                        )
-                        for i in range(7, 9)
-                    ]
-                )
-            if len(model_config.block_out_channels) > 3:
-                input_info.extend(
-                    [
-                        (
-                            f"down_block_additional_residuals_{i}",
-                            [rbln_batch_size, model_config.block_out_channels[3], input_height // 8, input_width // 8],
-                            "float32",
-                        )
-                        for i in range(9, 12)
-                    ]
-                )
-            input_info.append(
-                (
-                    "mid_block_additional_residual",
-                    [
-                        rbln_batch_size,
-                        model_config.block_out_channels[-1],
-                        input_height // 2 ** (len(model_config.block_out_channels) - 1),
-                        input_width // 2 ** (len(model_config.block_out_channels) - 1),
-                    ],
-                    "float32",
-                )
-            )
+        if is_controlnet:
+            # down block addtional residuals
+            first_shape = [batch_size, model_config.block_out_channels[0], sample_size[0], sample_size[1]]
+            height, width = sample_size[0], sample_size[1]
+            input_info.append(("down_block_additional_residuals_0", first_shape, "float32"))
+            name_idx = 1
+            for idx, _ in enumerate(model_config.down_block_types):
+                shape = [batch_size, model_config.block_out_channels[idx], height, width]
+                for _ in range(model_config.layers_per_block):
+                    input_info.append((f"down_block_additional_residuals_{name_idx}", shape, "float32"))
+                    name_idx += 1
+                if idx != len(model_config.down_block_types) - 1:
+                    height = height // 2
+                    width = width // 2
+                    shape = [batch_size, model_config.block_out_channels[idx], height, width]
+                    input_info.append((f"down_block_additional_residuals_{name_idx}", shape, "float32"))
+                    name_idx += 1
+
+            # mid block addtional residual
+            num_cross_attn_blocks = model_config.down_block_types.count("CrossAttnDownBlock2D")
+            out_channels = model_config.block_out_channels[-1]
+            shape = [
+                batch_size,
+                out_channels,
+                sample_size[0] // 2**num_cross_attn_blocks,
+                sample_size[1] // 2**num_cross_attn_blocks,
+            ]
+            input_info.append(("mid_block_additional_residual", shape, "float32"))
 
         rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
         if hasattr(model_config, "addition_embed_type") and model_config.addition_embed_type == "text_time":
-            if rbln_text_model_hidden_size is None:
-                rbln_text_model_hidden_size = 768
-            if rbln_in_features is None:
-                rbln_in_features = model_config.projection_class_embeddings_input_dim
+            rbln_text_model_hidden_size = rbln_kwargs["text_model_hidden_size"]
+            rbln_in_features = model_config.projection_class_embeddings_input_dim
             rbln_compile_config.input_info.append(
-                ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")
+                ("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32")
             )
-            rbln_compile_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+            rbln_compile_config.input_info.append(("time_ids", [batch_size, 6], "float32"))
 
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
@@ -320,14 +232,6 @@ class RBLNUNet2DConditionModel(RBLNModel):
             rbln_kwargs=rbln_kwargs,
         )
 
-        rbln_config.model_cfg.update(
-            {
-                "max_seq_len": rbln_max_seq_len,
-                "batch_size": rbln_batch_size,
-                "use_encode": rbln_use_encode,
-            }
-        )
-
         if rbln_in_features is not None:
             rbln_config.model_cfg["in_features"] = rbln_in_features