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

optimum/rbln/__init__.py

@@ -25,6 +25,7 @@ from typing import TYPE_CHECKING
 
 from transformers.utils import _LazyModule
 
+from .__version__ import __version__
 from .utils import check_version_compats
 
 
@@ -54,13 +55,30 @@ _import_structure = {
     ],
     "transformers": [
         "BatchTextIteratorStreamer",
+        "RBLNAutoModel",
+        "RBLNAutoModelForAudioClassification",
+        "RBLNAutoModelForCausalLM",
+        "RBLNAutoModelForCTC",
+        "RBLNAutoModelForDepthEstimation",
+        "RBLNAutoModelForImageClassification",
+        "RBLNAutoModelForMaskedLM",
+        "RBLNAutoModelForQuestionAnswering",
+        "RBLNAutoModelForSeq2SeqLM",
+        "RBLNAutoModelForSequenceClassification",
+        "RBLNAutoModelForSpeechSeq2Seq",
+        "RBLNAutoModelForVision2Seq",
+        "RBLNBartModel",
+        "RBLNBertModel",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
+        "RBLNCLIPVisionModel",
         "RBLNDPTForDepthEstimation",
         "RBLNGemmaForCausalLM",
         "RBLNGPT2LMHeadModel",
         "RBLNWav2Vec2ForCTC",
         "RBLNLlamaForCausalLM",
+        "RBLNPhiForCausalLM",
+        "RBLNLlavaNextForConditionalGeneration",
         "RBLNMidmLMHeadModel",
         "RBLNMistralForCausalLM",
         "RBLNWhisperForConditionalGeneration",
@@ -80,7 +98,7 @@ _import_structure = {
         "RBLNStableDiffusionXLControlNetPipeline",
         "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
     ],
-    "modeling_config": ["RBLNRuntimeConfig", "RBLNConfig"],
+    "modeling_config": ["RBLNCompileConfig", "RBLNConfig"],
 }
 
 if TYPE_CHECKING:
@@ -117,18 +135,35 @@ if TYPE_CHECKING:
         RBLNModelForQuestionAnswering,
         RBLNModelForSequenceClassification,
     )
-    from .modeling_config import RBLNConfig, RBLNRuntimeConfig
+    from .modeling_config import RBLNCompileConfig, RBLNConfig
     from .modeling_seq2seq import RBLNModelForSeq2SeqLM
     from .transformers import (
         BatchTextIteratorStreamer,
+        RBLNAutoModel,
+        RBLNAutoModelForAudioClassification,
+        RBLNAutoModelForCausalLM,
+        RBLNAutoModelForCTC,
+        RBLNAutoModelForDepthEstimation,
+        RBLNAutoModelForImageClassification,
+        RBLNAutoModelForMaskedLM,
+        RBLNAutoModelForQuestionAnswering,
+        RBLNAutoModelForSeq2SeqLM,
+        RBLNAutoModelForSequenceClassification,
+        RBLNAutoModelForSpeechSeq2Seq,
+        RBLNAutoModelForVision2Seq,
+        RBLNBartModel,
+        RBLNBertModel,
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
+        RBLNCLIPVisionModel,
         RBLNDPTForDepthEstimation,
         RBLNGemmaForCausalLM,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
+        RBLNLlavaNextForConditionalGeneration,
         RBLNMidmLMHeadModel,
         RBLNMistralForCausalLM,
+        RBLNPhiForCausalLM,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
         RBLNXLMRobertaModel,

optimum/rbln/__version__.py

@@ -1 +1 @@
-__version__ = '0.1.9'
+__version__ = '0.1.11'

optimum/rbln/diffusers/models/autoencoder_kl.py

@@ -23,7 +23,7 @@
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Union
 
 import rebel
 import torch  # noqa: I001
@@ -34,7 +34,7 @@ from optimum.exporters import TasksManager
 from transformers import AutoConfig, AutoModel, PretrainedConfig
 
 from ...modeling_base import RBLNModel
-from ...modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNConfig, RBLNRuntimeConfig
+from ...modeling_config import DEFAULT_COMPILED_MODEL_NAME, RBLNCompileConfig, RBLNConfig
 from ...utils.runtime_utils import RBLNPytorchRuntime
 
 
@@ -63,10 +63,9 @@ class RBLNAutoencoderKL(RBLNModel):
     auto_model_class = AutoModel  # feature extraction
 
     def __post_init__(self, **kwargs):
-        self.dtype = torch.float32
-
-        self.rbln_use_encode = self.rbln_config.meta["rbln_use_encode"]
+        super().__post_init__(**kwargs)
 
+        self.rbln_use_encode = self.rbln_config.model_cfg["use_encode"]
         if self.rbln_use_encode:
             self.encoder = RBLNRuntimeVAEEncoder(runtime=self.model[0], main_input_name="x")
             self.decoder = RBLNRuntimeVAEDecoder(runtime=self.model[1], main_input_name="z")
@@ -81,20 +80,20 @@ class RBLNAutoencoderKL(RBLNModel):
             encoder_model.eval()
             decoder_model.eval()
 
-            enc_compiled_model = cls.compile(encoder_model, rbln_runtime_config=rbln_config["encoder"][0])
-            dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["decoder"][0])
+            enc_compiled_model = cls.compile(encoder_model, rbln_compile_config=rbln_config.compile_cfgs[0])
+            dec_compiled_model = cls.compile(decoder_model, rbln_compile_config=rbln_config.compile_cfgs[1])
 
-            return enc_compiled_model, dec_compiled_model
+            return {"encoder": enc_compiled_model, "decoder": dec_compiled_model}
 
         def compile_text2img():
             decoder_model = _VAEDecoder(model)
             decoder_model.eval()
 
-            dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["compiled_model"][0])
+            dec_compiled_model = cls.compile(decoder_model, rbln_compile_config=rbln_config.compile_cfgs[0])
 
             return dec_compiled_model
 
-        if rbln_config.meta.get("rbln_use_encode", False):
+        if rbln_config.model_cfg.get("use_encode", False):
             return compile_img2img()
         else:
             return compile_text2img()
@@ -133,23 +132,23 @@ class RBLNAutoencoderKL(RBLNModel):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
         model_config: "PretrainedConfig",
-        rbln_unet_sample_size: Optional[int] = None,
-        rbln_img_width: Optional[int] = None,
-        rbln_img_height: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        rbln_use_encode: Optional[bool] = None,
-        rbln_vae_scale_factor: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        meta = {}
+        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)
+
         if rbln_batch_size is None:
             rbln_batch_size = 1
 
-        meta["rbln_use_encode"] = rbln_use_encode
-        meta["rbln_batch_size"] = rbln_batch_size
+        model_cfg = {}
 
         if rbln_use_encode:
-            meta["rbln_img_width"] = rbln_img_width
-            meta["rbln_img_height"] = rbln_img_height
+            model_cfg["img_width"] = rbln_img_width
+            model_cfg["img_height"] = rbln_img_height
 
             vae_enc_input_info = [
                 ("x", [rbln_batch_size, model_config.in_channels, rbln_img_height, rbln_img_width], "float32")
@@ -167,20 +166,23 @@ class RBLNAutoencoderKL(RBLNModel):
                 )
             ]
 
-            enc_rbln_runtime_config = RBLNRuntimeConfig(compiled_model_name="encoder", input_info=vae_enc_input_info)
-            dec_rbln_runtime_config = RBLNRuntimeConfig(compiled_model_name="decoder", input_info=vae_dec_input_info)
+            enc_rbln_compile_config = RBLNCompileConfig(compiled_model_name="encoder", input_info=vae_enc_input_info)
+            dec_rbln_compile_config = RBLNCompileConfig(compiled_model_name="decoder", input_info=vae_dec_input_info)
 
-            rbln_config = RBLNConfig.from_rbln_runtime_configs(
-                [enc_rbln_runtime_config, dec_rbln_runtime_config],
-                _rbln_meta=meta,
+            compile_cfgs = [enc_rbln_compile_config, dec_rbln_compile_config]
+            rbln_config = RBLNConfig(
+                rbln_cls=cls.__name__,
+                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
 
-        meta["rbln_unet_sample_size"] = rbln_unet_sample_size
-        vae_config = RBLNRuntimeConfig(
+        model_cfg["unet_sample_size"] = rbln_unet_sample_size
+        vae_config = RBLNCompileConfig(
             input_info=[
                 (
                     "z",
@@ -194,7 +196,12 @@ class RBLNAutoencoderKL(RBLNModel):
                 )
             ],
         )
-        rbln_config = RBLNConfig.from_rbln_runtime_configs([vae_config], _rbln_meta=meta)
+        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

@@ -23,7 +23,7 @@
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import torch
 from diffusers import ControlNetModel
@@ -31,7 +31,7 @@ from optimum.exporters import TasksManager
 from transformers import AutoConfig, AutoModel, PretrainedConfig
 
 from ...modeling_base import RBLNModel
-from ...modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ...modeling_config import RBLNCompileConfig, RBLNConfig
 
 
 if TYPE_CHECKING:
@@ -109,21 +109,21 @@ class RBLNControlNetModel(RBLNModel):
     auto_model_class = AutoModel  # feature extraction
 
     def __post_init__(self, **kwargs):
-        self.dtype = torch.float32
+        super().__post_init__(**kwargs)
         self.use_encoder_hidden_states = any(
-            item[0] == "encoder_hidden_states" for item in self.rbln_config["compiled_model"][0].input_info
+            item[0] == "encoder_hidden_states" for item in self.rbln_config.compile_cfgs[0].input_info
         )
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
+        if "subfolder" in kwargs:
+            del kwargs["subfolder"]
+
         def get_model_from_task(
             task: str,
             model_name_or_path: Union[str, Path],
             **kwargs,
         ):
-            if "subfolder" in kwargs:
-                del kwargs["subfolder"]
-
             return ControlNetModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path, **kwargs)
 
         tasktmp = TasksManager.get_model_from_task
@@ -155,14 +155,14 @@ class RBLNControlNetModel(RBLNModel):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
         model_config: "PretrainedConfig",
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_text_model_hidden_size: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        rbln_img_width: Optional[int] = None,
-        rbln_img_height: Optional[int] = None,
-        rbln_vae_scale_factor: Optional[int] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        meta = {"type": "controlnet"}
+        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)
 
         if rbln_batch_size is None:
             rbln_batch_size = 1
@@ -170,28 +170,29 @@ class RBLNControlNetModel(RBLNModel):
         if rbln_max_seq_len is None:
             rbln_max_seq_len = 77
 
+        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")
+
         input_width = rbln_img_width // rbln_vae_scale_factor
         input_height = rbln_img_height // rbln_vae_scale_factor
 
-        rbln_runtime_config = RBLNRuntimeConfig(
-            input_info=[
-                (
-                    "sample",
-                    [
-                        rbln_batch_size,
-                        model_config.in_channels,
-                        input_height,
-                        input_width,
-                    ],
-                    "float32",
-                ),
-                ("timestep", [], "float32"),
-            ],
-            batch_size=rbln_batch_size,
-        )
+        input_info = [
+            (
+                "sample",
+                [
+                    rbln_batch_size,
+                    model_config.in_channels,
+                    input_height,
+                    input_width,
+                ],
+                "float32",
+            ),
+            ("timestep", [], "float32"),
+        ]
+
         use_encoder_hidden_states = any(element != "DownBlock2D" for element in model_config.down_block_types)
         if use_encoder_hidden_states:
-            rbln_runtime_config.input_info.append(
+            input_info.append(
                 (
                     "encoder_hidden_states",
                     [
@@ -202,19 +203,34 @@ class RBLNControlNetModel(RBLNModel):
                     "float32",
                 )
             )
-        rbln_runtime_config.input_info.append(
-            ("controlnet_cond", [rbln_batch_size, 3, rbln_img_height, rbln_img_width], "float32")
-        )
-        rbln_runtime_config.input_info.append(("conditioning_scale", [], "float32"))
+
+        input_info.append(("controlnet_cond", [rbln_batch_size, 3, rbln_img_height, rbln_img_width], "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
-            rbln_runtime_config.input_info.append(
-                ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")
-            )
-            rbln_runtime_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+            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_compile_config = RBLNCompileConfig(input_info=input_info)
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            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,
+            }
+        )
 
-        rbln_config = RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
         return rbln_config
 
     def forward(

optimum/rbln/diffusers/models/unet_2d_condition.py

@@ -32,7 +32,7 @@ from optimum.exporters import TasksManager
 from transformers import AutoConfig, AutoModel, PretrainedConfig
 
 from ...modeling_base import RBLNModel
-from ...modeling_config import RBLNConfig, RBLNRuntimeConfig
+from ...modeling_config import RBLNCompileConfig, RBLNConfig
 
 
 if TYPE_CHECKING:
@@ -130,8 +130,8 @@ class RBLNUNet2DConditionModel(RBLNModel):
     auto_model_class = AutoModel  # feature extraction
 
     def __post_init__(self, **kwargs):
-        self.dtype = torch.float32
-        self.in_features = self.rbln_config.meta.get("in_features", None)
+        super().__post_init__(**kwargs)
+        self.in_features = self.rbln_config.model_cfg.get("in_features", None)
         if self.in_features is not None:
 
             @dataclass
@@ -183,31 +183,31 @@ class RBLNUNet2DConditionModel(RBLNModel):
         cls,
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
         model_config: "PretrainedConfig",
-        rbln_max_seq_len: Optional[int] = None,
-        rbln_text_model_hidden_size: Optional[int] = None,
-        rbln_batch_size: Optional[int] = None,
-        rbln_in_features: Optional[int] = None,
-        rbln_use_encode: Optional[bool] = None,
-        rbln_img_width: Optional[int] = None,
-        rbln_img_height: Optional[int] = None,
-        rbln_vae_scale_factor: Optional[int] = None,
-        rbln_is_controlnet: Optional[bool] = None,
+        rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        meta = {"type": "unet"}
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
+        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
-
-        meta["rbln_use_encode"] = rbln_use_encode
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
 
         if rbln_use_encode:
-            # FIXME :: robust img shape getter
+            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:
-            # FIXME :: model_config.sample_size can be tuple or list
             input_width, input_height = model_config.sample_size, model_config.sample_size
 
         input_info = [
@@ -232,6 +232,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
                 "float32",
             ),
         ]
+
         if rbln_is_controlnet:
             if len(model_config.block_out_channels) > 0:
                 input_info.extend(
@@ -304,24 +305,35 @@ class RBLNUNet2DConditionModel(RBLNModel):
                 )
             )
 
-        rbln_runtime_config = RBLNRuntimeConfig(
-            input_info=input_info,
-            batch_size=rbln_batch_size,
-        )
+        rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
         if hasattr(model_config, "addition_embed_type") and model_config.addition_embed_type == "text_time":
-            # In case of sdxl
             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
-            meta["in_features"] = rbln_in_features
-            rbln_runtime_config.input_info.append(
+            rbln_compile_config.input_info.append(
                 ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")
             )
-            rbln_runtime_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+            rbln_compile_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+
+        rbln_config = RBLNConfig(
+            rbln_cls=cls.__name__,
+            compile_cfgs=[rbln_compile_config],
+            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
 
-        rbln_config = RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
         return rbln_config
 
     def forward(

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

@@ -37,6 +37,7 @@ from transformers import CLIPTextModel
 
 from ....modeling_base import RBLNBaseModel
 from ....transformers import RBLNCLIPTextModel
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -69,8 +70,13 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
         text_encoder = kwargs.pop("text_encoder", None)
         controlnet = kwargs.pop("controlnet", 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)
 
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(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)
 
         kwargs_dict = {
             "pretrained_model_name_or_path": model_id,
@@ -98,13 +104,19 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
             }
         )
 
-        model = super().from_pretrained(**{k: v for k, v in kwargs_dict.items() if v is not None})
+        with ContextRblnConfig(
+            device=device,
+            device_map=device_map,
+            create_runtimes=create_runtimes,
+            optimze_host_mem=optimize_host_memory,
+        ):
+            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
 
         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
+            rbln_kwargs.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         # compile model, create runtime
@@ -117,8 +129,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 rbln_unet_sample_size=model.unet.config.sample_size,
                 rbln_use_encode=False,
                 rbln_vae_scale_factor=model.vae_scale_factor,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_kwargs},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -127,11 +138,10 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_kwargs},
             )
 
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
+        batch_size = rbln_kwargs.pop("batch_size", 1)
         unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
 
         if not isinstance(unet, RBLNUNet2DConditionModel):
@@ -145,8 +155,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 rbln_use_encode=False,
                 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,
+                rbln_config={**rbln_kwargs},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -162,8 +171,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                             model_save_dir=model_save_dir,
                             rbln_batch_size=unet_batch_size,
                             rbln_vae_scale_factor=model.vae_scale_factor,
-                            **rbln_config_kwargs,
-                            **rbln_constructor_kwargs,
+                            rbln_config={**rbln_kwargs},
                         )
                     )
                 controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
@@ -176,8 +184,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                     model_save_dir=model_save_dir,
                     rbln_batch_size=unet_batch_size,
                     rbln_vae_scale_factor=model.vae_scale_factor,
-                    **rbln_config_kwargs,
-                    **rbln_constructor_kwargs,
+                    rbln_config={**rbln_kwargs},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -209,7 +216,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
             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:
+        if optimize_host_memory is False:
             model.compiled_models = [vae.compiled_models[0], text_encoder.compiled_models[0], unet.compiled_models[0]]
             if isinstance(controlnet, RBLNMultiControlNetModel):
                 for c_model in controlnet.nets: