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

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

@@ -36,6 +36,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
 
@@ -68,8 +69,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_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,
@@ -97,13 +103,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_kwargs.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         # compile model, create runtime
@@ -116,8 +128,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_kwargs},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -126,11 +137,10 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                 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):
@@ -144,8 +154,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_kwargs},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -161,8 +170,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                             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)
@@ -175,8 +183,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                     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")
 
@@ -208,7 +215,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
             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

@@ -36,6 +36,7 @@ from transformers import CLIPTextModel
 
 from ....modeling_base import RBLNBaseModel
 from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -69,8 +70,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_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 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_kwargs.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         if not isinstance(vae, RBLNAutoencoderKL):
@@ -116,8 +128,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_kwargs},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -126,8 +137,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_kwargs},
             )
 
         if not isinstance(text_encoder_2, RBLNCLIPTextModel):
@@ -136,11 +146,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_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):
@@ -155,8 +164,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_kwargs},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -173,8 +181,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                             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_kwargs},
                         )
                     )
                 controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
@@ -188,8 +195,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                     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_kwargs},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -223,7 +229,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
             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],

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

@@ -36,6 +36,7 @@ from transformers import CLIPTextModel
 
 from ....modeling_base import RBLNBaseModel
 from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
@@ -69,8 +70,13 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
         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_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 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_kwargs.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         if not isinstance(vae, RBLNAutoencoderKL):
@@ -116,8 +128,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_kwargs},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -126,8 +137,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                **rbln_config_kwargs,
-                **rbln_constructor_kwargs,
+                rbln_config={**rbln_kwargs},
             )
 
         if not isinstance(text_encoder_2, RBLNCLIPTextModel):
@@ -136,11 +146,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_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):
@@ -155,8 +164,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_kwargs},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -173,8 +181,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                             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_kwargs},
                         )
                     )
                 controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
@@ -188,8 +195,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                     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_kwargs},
                 )
                 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)
@@ -223,7 +229,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
             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/stable_diffusion/pipeline_stable_diffusion.py

@@ -26,6 +26,7 @@ from diffusers import StableDiffusionPipeline
 
 from ....modeling_base import RBLNBaseModel
 from ....transformers import RBLNCLIPTextModel
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNUNet2DConditionModel
 
 
@@ -50,14 +51,27 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
         """
         export = kwargs.pop("export", None)
         model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
-        model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
+        rbln_config = kwargs.pop("rbln_config", None)
+        rbln_kwargs, _ = RBLNBaseModel.resolve_rbln_config(rbln_config, kwargs)
+
+        device = rbln_kwargs.get("device", None)
+        device_map = rbln_kwargs.get("device_map", None)
+        create_runtimes = rbln_kwargs.get("create_runtimes", None)
+        optimize_host_memory = rbln_kwargs.get("optimize_host_memory", None)
+
+        with ContextRblnConfig(
+            device=device,
+            device_map=device_map,
+            create_runtimes=create_runtimes,
+            optimze_host_mem=optimize_host_memory,
+        ):
+            model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
 
         if export is None or export is False:
             return model
 
         do_classifier_free_guidance = (
-            rbln_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
         )
 
         vae = RBLNAutoencoderKL.from_pretrained(
@@ -67,19 +81,17 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
             model_save_dir=model_save_dir,
             rbln_unet_sample_size=model.unet.config.sample_size,
             rbln_use_encode=False,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
+            rbln_config={**rbln_kwargs},
         )
         text_encoder = RBLNCLIPTextModel.from_pretrained(
             model_id=model_id,
             subfolder="text_encoder",
             export=True,
             model_save_dir=model_save_dir,
-            **rbln_config_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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -91,8 +103,7 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
             rbln_batch_size=unet_batch_size,
             rbln_use_encode=False,
             rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
+            rbln_config={**rbln_kwargs},
         )
 
         if model_save_dir is not None:
@@ -122,7 +133,7 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
 
         model.models = [vae.model[0], text_encoder.model[0], unet.model[0]]
 
-        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]]
 
         return model

optimum/rbln/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -26,6 +26,7 @@ from diffusers import StableDiffusionImg2ImgPipeline
 
 from ....modeling_base import RBLNBaseModel
 from ....transformers import RBLNCLIPTextModel
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNUNet2DConditionModel
 
 
@@ -50,14 +51,27 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
         """
         export = kwargs.pop("export", None)
         model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
-        model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
+        rbln_config = kwargs.pop("rbln_config", None)
+        rbln_kwargs, _ = RBLNBaseModel.resolve_rbln_config(rbln_config, kwargs)
+
+        device = rbln_kwargs.get("device", None)
+        device_map = rbln_kwargs.get("device_map", None)
+        create_runtimes = rbln_kwargs.get("create_runtimes", None)
+        optimize_host_memory = rbln_kwargs.get("optimize_host_memory", None)
+
+        with ContextRblnConfig(
+            device=device,
+            device_map=device_map,
+            create_runtimes=create_runtimes,
+            optimze_host_mem=optimize_host_memory,
+        ):
+            model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
 
         if export is None or export is False:
             return model
 
         do_classifier_free_guidance = (
-            rbln_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
@@ -69,19 +83,17 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
             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_kwargs},
         )
         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,
+            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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -94,8 +106,7 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
             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_kwargs},
         )
 
         if model_save_dir is not None:
@@ -126,7 +137,7 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
         # vae encoder, vae decoder, text_encoder, unet
         model.models = [vae.model[0], vae.model[1], text_encoder.model[0], unet.model[0]]
 
-        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/stable_diffusion_xl/pipeline_stable_diffusion_xl.py

@@ -18,6 +18,7 @@ from diffusers import StableDiffusionXLPipeline
 
 from ....modeling_base import RBLNBaseModel
 from ....transformers import RBLNCLIPTextModel, RBLNCLIPTextModelWithProjection
+from ....utils.runtime_utils import ContextRblnConfig
 from ...models import RBLNAutoencoderKL, RBLNUNet2DConditionModel
 
 
@@ -42,14 +43,27 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
         """
         export = kwargs.pop("export", None)
         model_save_dir = kwargs.pop("model_save_dir", None)
-        rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
-        model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
+        rbln_config = kwargs.pop("rbln_config", None)
+        rbln_kwargs, _ = RBLNBaseModel.resolve_rbln_config(rbln_config, kwargs)
+
+        device = rbln_kwargs.get("device", None)
+        device_map = rbln_kwargs.get("device_map", None)
+        create_runtimes = rbln_kwargs.get("create_runtimes", None)
+        optimize_host_memory = rbln_kwargs.get("optimize_host_memory", None)
+
+        with ContextRblnConfig(
+            device=device,
+            device_map=device_map,
+            create_runtimes=create_runtimes,
+            optimze_host_mem=optimize_host_memory,
+        ):
+            model = super().from_pretrained(pretrained_model_name_or_path=model_id, **kwargs)
 
         if export is None or export is False:
             return model
 
         do_classifier_free_guidance = (
-            rbln_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
         )
 
         vae = RBLNAutoencoderKL.from_pretrained(
@@ -59,27 +73,24 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             model_save_dir=model_save_dir,
             rbln_unet_sample_size=model.unet.config.sample_size,
             rbln_use_encode=False,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
+            rbln_config={**rbln_kwargs},
         )
         text_encoder = RBLNCLIPTextModel.from_pretrained(
             model_id=model_id,
             subfolder="text_encoder",
             export=True,
             model_save_dir=model_save_dir,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
+            rbln_config={**rbln_kwargs},
         )
         text_encoder_2 = RBLNCLIPTextModelWithProjection.from_pretrained(
             model_id=model_id,
             subfolder="text_encoder_2",
             export=True,
             model_save_dir=model_save_dir,
-            **rbln_config_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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -92,8 +103,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             rbln_batch_size=unet_batch_size,
             rbln_use_encode=False,
             rbln_is_controlnet=True if "controlnet" in model.config.keys() else False,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
+            rbln_config={**rbln_kwargs},
         )
 
         if model_save_dir is not None:
@@ -112,7 +122,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             "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"),
         }
         model.register_to_config(**update_dict)
 
@@ -122,7 +132,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
 
         model.models = [vae.model[0], unet.model[0], text_encoder.model[0], text_encoder_2.model[0]]
 
-        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],
                 unet.compiled_models[0],