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

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

@@ -24,13 +24,15 @@
 
 from diffusers import StableDiffusionPipeline
 
-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, RBLNUNet2DConditionModel
 
 
 class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
     @classmethod
+    @use_rbln_config
     def from_pretrained(cls, model_id, **kwargs):
         """
         Pipeline for text-to-image generation using Stable Diffusion.
@@ -50,14 +52,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_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)
+
+        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_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         vae = RBLNAutoencoderKL.from_pretrained(
@@ -67,19 +82,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_config},
         )
         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_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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -91,8 +104,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_config},
         )
 
         if model_save_dir is not None:
@@ -120,9 +132,7 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
             # overwrite to replace incorrect config
             model.save_config(model_save_dir)
 
-        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

@@ -24,13 +24,15 @@
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
-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, RBLNUNet2DConditionModel
 
 
 class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
     @classmethod
+    @use_rbln_config
     def from_pretrained(cls, model_id, **kwargs):
         """
         Pipeline for image-to-image generation using Stable Diffusion.
@@ -50,14 +52,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_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)
+
+        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_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         # compile model, create runtime
@@ -69,19 +84,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_config},
         )
         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_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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -94,8 +107,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_config},
         )
 
         if model_save_dir is not None:
@@ -123,10 +135,7 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
             # overwrite to replace incorrect config
             model.save_config(model_save_dir)
 
-        # 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

@@ -16,13 +16,15 @@
 
 from diffusers import StableDiffusionXLPipeline
 
-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, RBLNUNet2DConditionModel
 
 
 class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
     @classmethod
+    @use_rbln_config
     def from_pretrained(cls, model_id, **kwargs):
         """
         Pipeline for text-to-image generation using Stable Diffusion XL.
@@ -42,14 +44,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_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)
+
+        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_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         vae = RBLNAutoencoderKL.from_pretrained(
@@ -59,27 +74,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_config},
         )
         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_config},
         )
         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_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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -92,8 +104,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_config},
         )
 
         if model_save_dir is not None:
@@ -112,7 +123,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)
 
@@ -120,9 +131,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             # overwrite to replace incorrect config
             model.save_config(model_save_dir)
 
-        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],

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

@@ -16,13 +16,15 @@
 
 from diffusers import StableDiffusionXLImg2ImgPipeline
 
-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, RBLNUNet2DConditionModel
 
 
 class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
     @classmethod
+    @use_rbln_config
     def from_pretrained(cls, model_id, **kwargs):
         """
         Pipeline for image-to-image generation using Stable Diffusion XL.
@@ -42,14 +44,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_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)
+
+        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_config.pop("guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
         vae = RBLNAutoencoderKL.from_pretrained(
@@ -60,27 +75,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_config},
         )
         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_config},
         )
         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_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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -94,8 +106,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_config},
         )
 
         if model_save_dir is not None:
@@ -114,7 +125,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)
 
@@ -122,9 +133,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             # overwrite to replace incorrect config
             model.save_config(model_save_dir)
 
-        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

@@ -28,7 +28,6 @@ from .modeling_base import (
     RBLNModelForQuestionAnswering,
     RBLNModelForSequenceClassification,
 )
-from .modeling_seq2seq import RBLNModelForSeq2SeqLM
 
 
 class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
@@ -47,21 +46,17 @@ class RBLNResNetForImageClassification(RBLNModelForImageClassification):
     pass
 
 
-class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    pass
-
-
-class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    pass
-
-
 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):
+    rbln_model_input_names = ["input_ids", "attention_mask"]
+
+
+class RBLNViTForImageClassification(RBLNModelForImageClassification):
     pass