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

optimum/rbln/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py

@@ -18,6 +18,7 @@ from diffusers import StableDiffusionXLImg2ImgPipeline
 
 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 RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
         """
         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(
@@ -60,27 +74,24 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             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},
         )
         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(
@@ -94,8 +105,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             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:
@@ -114,7 +124,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             "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)
 
@@ -124,7 +134,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
 
         model.models = [vae.model[0], vae.model[1], 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],
                 vae.compiled_models[1],

optimum/rbln/modeling_alias.py

@@ -56,12 +56,12 @@ class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
 
 
 class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
-    pass
+    rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
 class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
-    pass
+    rbln_model_input_names = ["input_ids", "attention_mask"]
 
 
 class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
-    pass
+    rbln_model_input_names = ["input_ids", "attention_mask"]