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

optimum/rbln/__init__.py

@@ -35,11 +35,10 @@ _import_structure = {
         "RBLNBertForQuestionAnswering",
         "RBLNDistilBertForQuestionAnswering",
         "RBLNResNetForImageClassification",
-        "RBLNT5ForConditionalGeneration",
-        "RBLNBartForConditionalGeneration",
         "RBLNXLMRobertaForSequenceClassification",
         "RBLNRobertaForSequenceClassification",
         "RBLNRobertaForMaskedLM",
+        "RBLNViTForImageClassification"
     ],
     "modeling_base": [
         "RBLNBaseModel",
@@ -50,9 +49,6 @@ _import_structure = {
         "RBLNModelForSequenceClassification",
         "RBLNModelForMaskedLM",
     ],
-    "modeling_seq2seq": [
-        "RBLNModelForSeq2SeqLM",
-    ],
     "transformers": [
         "BatchTextIteratorStreamer",
         "RBLNAutoModel",
@@ -67,16 +63,20 @@ _import_structure = {
         "RBLNAutoModelForSequenceClassification",
         "RBLNAutoModelForSpeechSeq2Seq",
         "RBLNAutoModelForVision2Seq",
+        "RBLNBartForConditionalGeneration",
         "RBLNBartModel",
         "RBLNBertModel",
         "RBLNCLIPTextModel",
         "RBLNCLIPTextModelWithProjection",
         "RBLNCLIPVisionModel",
         "RBLNDPTForDepthEstimation",
+        "RBLNExaoneForCausalLM",
         "RBLNGemmaForCausalLM",
         "RBLNGPT2LMHeadModel",
+        "RBLNQwen2ForCausalLM",
         "RBLNWav2Vec2ForCTC",
         "RBLNLlamaForCausalLM",
+        "RBLNT5ForConditionalGeneration",
         "RBLNPhiForCausalLM",
         "RBLNLlavaNextForConditionalGeneration",
         "RBLNMidmLMHeadModel",
@@ -118,12 +118,12 @@ if TYPE_CHECKING:
     )
     from .modeling_alias import (
         RBLNASTForAudioClassification,
-        RBLNBartForConditionalGeneration,
         RBLNBertForQuestionAnswering,
         RBLNResNetForImageClassification,
         RBLNRobertaForMaskedLM,
         RBLNRobertaForSequenceClassification,
         RBLNT5ForConditionalGeneration,
+        RBLNViTForImageClassification,
         RBLNXLMRobertaForSequenceClassification,
     )
     from .modeling_base import (
@@ -136,7 +136,6 @@ if TYPE_CHECKING:
         RBLNModelForSequenceClassification,
     )
     from .modeling_config import RBLNCompileConfig, RBLNConfig
-    from .modeling_seq2seq import RBLNModelForSeq2SeqLM
     from .transformers import (
         BatchTextIteratorStreamer,
         RBLNAutoModel,
@@ -151,12 +150,14 @@ if TYPE_CHECKING:
         RBLNAutoModelForSequenceClassification,
         RBLNAutoModelForSpeechSeq2Seq,
         RBLNAutoModelForVision2Seq,
+        RBLNBartForConditionalGeneration,
         RBLNBartModel,
         RBLNBertModel,
         RBLNCLIPTextModel,
         RBLNCLIPTextModelWithProjection,
         RBLNCLIPVisionModel,
         RBLNDPTForDepthEstimation,
+        RBLNExaoneForCausalLM,
         RBLNGemmaForCausalLM,
         RBLNGPT2LMHeadModel,
         RBLNLlamaForCausalLM,
@@ -164,6 +165,8 @@ if TYPE_CHECKING:
         RBLNMidmLMHeadModel,
         RBLNMistralForCausalLM,
         RBLNPhiForCausalLM,
+        RBLNQwen2ForCausalLM,
+        RBLNT5ForConditionalGeneration,
         RBLNWav2Vec2ForCTC,
         RBLNWhisperForConditionalGeneration,
         RBLNXLMRobertaModel,

optimum/rbln/__version__.py

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

optimum/rbln/diffusers/models/autoencoder_kl.py

@@ -58,9 +58,7 @@ class RBLNRuntimeVAEDecoder(RBLNPytorchRuntime):
 
 
 class RBLNAutoencoderKL(RBLNModel):
-    model_type = "rbln_model"
     config_name = "config.json"
-    auto_model_class = AutoModel  # feature extraction
 
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)

optimum/rbln/diffusers/models/controlnet.py

@@ -105,9 +105,6 @@ class _ControlNetModel_Cross_Attention(torch.nn.Module):
 
 
 class RBLNControlNetModel(RBLNModel):
-    model_type = "rbln_model"
-    auto_model_class = AutoModel  # feature extraction
-
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
         self.use_encoder_hidden_states = any(
@@ -116,9 +113,6 @@ class RBLNControlNetModel(RBLNModel):
 
     @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],

optimum/rbln/diffusers/models/unet_2d_condition.py

@@ -126,9 +126,6 @@ class _UNet_SDXL(torch.nn.Module):
 
 
 class RBLNUNet2DConditionModel(RBLNModel):
-    model_type = "rbln_model"
-    auto_model_class = AutoModel  # feature extraction
-
     def __post_init__(self, **kwargs):
         super().__post_init__(**kwargs)
         self.in_features = self.rbln_config.model_cfg.get("in_features", None)

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

@@ -102,6 +102,10 @@ class RBLNMultiControlNetModel(RBLNModel):
             real_save_path = save_directory + suffix
             model.save_pretrained(real_save_path)
 
+    @classmethod
+    def _get_rbln_config(cls, **rbln_config_kwargs):
+        pass
+
     def forward(
         self,
         sample: torch.FloatTensor,

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

@@ -35,7 +35,7 @@ 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
@@ -47,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.
@@ -71,12 +72,12 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
         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 = {} if rbln_config is None else rbln_config
 
-        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)
+        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,
@@ -116,7 +117,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
             return model
 
         do_classifier_free_guidance = (
-            rbln_kwargs.pop("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
@@ -129,7 +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={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -138,10 +139,10 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                rbln_config={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
-        batch_size = rbln_kwargs.pop("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,7 +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={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -164,27 +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={**rbln_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={**rbln_kwargs},
+                    rbln_config={**rbln_config},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -215,7 +214,6 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
             # overwrite to replace incorrect config
             model.save_config(model_save_dir)
 
-        # use for CI to access each compiled model
         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):

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

@@ -34,7 +34,7 @@ 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
@@ -46,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.
@@ -70,12 +71,12 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
         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 = {} if rbln_config is None else rbln_config
 
-        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)
+        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,
@@ -115,7 +116,7 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
             return model
 
         do_classifier_free_guidance = (
-            rbln_kwargs.pop("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
@@ -128,7 +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={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -137,10 +138,10 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                rbln_config={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
-        batch_size = rbln_kwargs.pop("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):
@@ -154,7 +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={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -163,27 +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={**rbln_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={**rbln_kwargs},
+                    rbln_config={**rbln_config},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -214,7 +213,6 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
             # overwrite to replace incorrect config
             model.save_config(model_save_dir)
 
-        # use for CI to access each compiled model
         if optimize_host_memory is False:
             model.compiled_models = [
                 vae.compiled_models[0],

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

@@ -34,7 +34,7 @@ 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
@@ -46,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.
@@ -71,12 +72,12 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
         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 = {} if rbln_config is None else rbln_config
 
-        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)
+        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,
@@ -116,7 +117,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
             return model
 
         do_classifier_free_guidance = (
-            rbln_kwargs.pop("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):
@@ -128,7 +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={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -137,7 +138,7 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                rbln_config={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder_2, RBLNCLIPTextModel):
@@ -146,10 +147,10 @@ class RBLNStableDiffusionXLControlNetPipeline(StableDiffusionXLControlNetPipelin
                 subfolder="text_encoder_2",
                 export=True,
                 model_save_dir=model_save_dir,
-                rbln_config={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
-        batch_size = rbln_kwargs.pop("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):
@@ -164,7 +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={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -173,29 +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={**rbln_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={**rbln_kwargs},
+                    rbln_config={**rbln_config},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -228,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 optimize_host_memory 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,7 +34,7 @@ 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
@@ -46,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,13 +71,14 @@ 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 = {} if rbln_config is None else rbln_config
 
-        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)
+        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,
@@ -116,7 +118,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
             return model
 
         do_classifier_free_guidance = (
-            rbln_kwargs.pop("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):
@@ -128,7 +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={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder, RBLNCLIPTextModel):
@@ -137,7 +139,7 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 subfolder="text_encoder",
                 export=True,
                 model_save_dir=model_save_dir,
-                rbln_config={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(text_encoder_2, RBLNCLIPTextModel):
@@ -146,10 +148,10 @@ class RBLNStableDiffusionXLControlNetImg2ImgPipeline(StableDiffusionXLControlNet
                 subfolder="text_encoder_2",
                 export=True,
                 model_save_dir=model_save_dir,
-                rbln_config={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
-        batch_size = rbln_kwargs.pop("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):
@@ -164,7 +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={**rbln_kwargs},
+                rbln_config={**rbln_config},
             )
 
         if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
@@ -173,29 +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={**rbln_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={**rbln_kwargs},
+                    rbln_config={**rbln_config},
                 )
                 controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
 
@@ -228,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 optimize_host_memory 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(