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

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

@@ -35,8 +35,9 @@ from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
 from transformers import CLIPTextModel
 
-from ....modeling_base import RBLNBaseModel
+from ....modeling_config import use_rbln_config
 from ....transformers import RBLNCLIPTextModel
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -46,6 +47,7 @@ logger = logging.get_logger(__name__)
 
 class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
     @classmethod
+    @use_rbln_config
     def from_pretrained(cls, model_id, **kwargs):
         """
         Pipeline for text-to-image generation using Stable Diffusion with ControlNet.
@@ -69,8 +71,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_config = {} if rbln_config is None else rbln_config
 
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
+        device = rbln_config.get("device", None)
+        device_map = rbln_config.get("device_map", None)
+        create_runtimes = rbln_config.get("create_runtimes", None)
+        optimize_host_memory = rbln_config.get("optimize_host_memory", None)
 
         kwargs_dict = {
             "pretrained_model_name_or_path": model_id,
@@ -98,13 +105,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_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         # compile model, create runtime
@@ -117,8 +130,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_config},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -127,11 +139,10 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_config},
             )
 
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
+        batch_size = rbln_config.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 +156,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_config},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -155,29 +165,25 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 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,
+                        RBLNControlNetModel.from_model(
+                            model=cid,
                             subfolder=subfolder_name,
-                            export=True,
                             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_config},
                         )
                     )
                 controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
                 controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
             else:
-                controlnet = RBLNControlNetModel.from_pretrained(
-                    model_id=controlnet.config._name_or_path,
+                controlnet = RBLNControlNetModel.from_model(
+                    model=controlnet,
                     subfolder="controlnet",
-                    export=True,
                     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_config},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -208,8 +214,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
             # 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:
+        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:

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

@@ -34,8 +34,9 @@ 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 ....modeling_config import use_rbln_config
 from ....transformers import RBLNCLIPTextModel
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -45,6 +46,7 @@ logger = logging.get_logger(__name__)
 
 class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2ImgPipeline):
     @classmethod
+    @use_rbln_config
     def from_pretrained(cls, model_id, **kwargs):
         """
         Pipeline for image-to-image generation using Stable Diffusion with ControlNet.
@@ -68,8 +70,13 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
         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_config = {} if rbln_config is None else rbln_config
 
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
+        device = rbln_config.get("device", None)
+        device_map = rbln_config.get("device_map", None)
+        create_runtimes = rbln_config.get("create_runtimes", None)
+        optimize_host_memory = rbln_config.get("optimize_host_memory", None)
 
         kwargs_dict = {
             "pretrained_model_name_or_path": model_id,
@@ -97,13 +104,19 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
             }
         )
 
-        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_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         # compile model, create runtime
@@ -116,8 +129,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                 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,
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -126,11 +138,10 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_config},
             )
 
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
+        batch_size = rbln_config.pop("batch_size", 1)
         unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
 
         if not isinstance(unet, RBLNUNet2DConditionModel):
@@ -144,8 +155,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                 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,
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -154,29 +164,25 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                 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,
+                        RBLNControlNetModel.from_model(
+                            model=cid,
                             subfolder=subfolder_name,
-                            export=True,
                             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_config},
                         )
                     )
                 controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
                 controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
             else:
-                controlnet = RBLNControlNetModel.from_pretrained(
-                    model_id=controlnet.config._name_or_path,
+                controlnet = RBLNControlNetModel.from_model(
+                    model=controlnet,
                     subfolder="controlnet",
-                    export=True,
                     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_config},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -207,8 +213,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
             # 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:
+        if optimize_host_memory is False:
             model.compiled_models = [
                 vae.compiled_models[0],
                 vae.compiled_models[1],

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

@@ -34,8 +34,9 @@ from diffusers.utils import deprecate, logging
 from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
 from transformers import CLIPTextModel
 
-from ....modeling_base import RBLNBaseModel
+from ....modeling_config import use_rbln_config
 from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -45,6 +46,7 @@ logger = logging.get_logger(__name__)
 
 class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipeline):
     @classmethod
+    @use_rbln_config
     def from_pretrained(cls, model_id, **kwargs):
         """
         Pipeline for text-to-image generation using Stable Diffusion XL with ControlNet.
@@ -69,8 +71,13 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
         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.pop("rbln_config", None)
+        rbln_config = {} if rbln_config is None else rbln_config
 
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
+        device = rbln_config.get("device", None)
+        device_map = rbln_config.get("device_map", None)
+        create_runtimes = rbln_config.get("create_runtimes", None)
+        optimize_host_memory = rbln_config.get("optimize_host_memory", None)
 
         kwargs_dict = {
             "pretrained_model_name_or_path": model_id,
@@ -98,13 +105,19 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
             }
         )
 
-        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_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         if not isinstance(vae, RBLNAutoencoderKL):
@@ -116,8 +129,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 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_config},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -126,8 +138,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder_2, RBLNCLIPTextModel):
@@ -136,11 +147,10 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 subfolder="text_encoder_2",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_config},
             )
 
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
+        batch_size = rbln_config.pop("batch_size", 1)
         unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
 
         if not isinstance(unet, RBLNUNet2DConditionModel):
@@ -155,8 +165,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 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_config},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -165,31 +174,27 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 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,
+                        RBLNControlNetModel.from_model(
+                            model=cid,
                             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,
+                            rbln_config={**rbln_config},
                         )
                     )
                 controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
                 controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
             else:
-                controlnet = RBLNControlNetModel.from_pretrained(
-                    model_id=controlnet.config._name_or_path,
+                controlnet = RBLNControlNetModel.from_model(
+                    model=controlnet,
                     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,
+                    rbln_config={**rbln_config},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -222,20 +227,6 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
             # 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],
-                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
 
     def check_inputs(

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

@@ -34,8 +34,9 @@ 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 ....modeling_config import use_rbln_config
 from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -45,6 +46,7 @@ logger = logging.get_logger(__name__)
 
 class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNetImg2ImgPipeline):
     @classmethod
+    @use_rbln_config
     def from_pretrained(cls, model_id, **kwargs):
         """
         Pipeline for image-to-image generation using Stable Diffusion XL with ControlNet.
@@ -70,7 +72,13 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
         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)
+        rbln_config = kwargs.pop("rbln_config", None)
+        rbln_config = {} if rbln_config is None else rbln_config
+
+        device = rbln_config.get("device", None)
+        device_map = rbln_config.get("device_map", None)
+        create_runtimes = rbln_config.get("create_runtimes", None)
+        optimize_host_memory = rbln_config.get("optimize_host_memory", None)
 
         kwargs_dict = {
             "pretrained_model_name_or_path": model_id,
@@ -98,13 +106,19 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
             }
         )
 
-        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_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         if not isinstance(vae, RBLNAutoencoderKL):
@@ -116,8 +130,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 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,
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -126,8 +139,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder_2, RBLNCLIPTextModel):
@@ -136,11 +148,10 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 subfolder="text_encoder_2",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_config},
             )
 
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
+        batch_size = rbln_config.pop("batch_size", 1)
         unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
 
         if not isinstance(unet, RBLNUNet2DConditionModel):
@@ -155,8 +166,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 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,
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -165,31 +175,27 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 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,
+                        RBLNControlNetModel.from_model(
+                            model=cid,
                             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,
+                            rbln_config={**rbln_config},
                         )
                     )
                 controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
                 controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
             else:
-                controlnet = RBLNControlNetModel.from_pretrained(
-                    model_id=controlnet.config._name_or_path,
+                controlnet = RBLNControlNetModel.from_model(
+                    model=controlnet,
                     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,
+                    rbln_config={**rbln_config},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -213,7 +219,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
             "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
             "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
             "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-            "text_encoder_2": ("optimum.rbln", "RBLNCLIPTextModel"),
+            "text_encoder_2": ("optimum.rbln", "RBLNCLIPTextModelWithProjection"),
             "controlnet": controlnet_dict,
         }
         model.register_to_config(**update_dict)
@@ -222,21 +228,6 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
             # 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
 
     def check_inputs(