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

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

@@ -24,7 +24,7 @@
 
 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
@@ -32,6 +32,7 @@ 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.
@@ -52,12 +53,12 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
         export = kwargs.pop("export", 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)
 
         with ContextRblnConfig(
             device=device,
@@ -71,7 +72,7 @@ class RBLNStableDiffusionPipeline(StableDiffusionPipeline):
             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
         )
 
         vae = RBLNAutoencoderKL.from_pretrained(
@@ -81,17 +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={**rbln_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={**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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -103,7 +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={**rbln_kwargs},
+            rbln_config={**rbln_config},
         )
 
         if model_save_dir is not None:
@@ -131,8 +132,6 @@ 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 optimize_host_memory is False:
             model.compiled_models = [vae.compiled_models[0], text_encoder.compiled_models[0], unet.compiled_models[0]]
 

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

@@ -24,7 +24,7 @@
 
 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
@@ -32,6 +32,7 @@ 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.
@@ -52,12 +53,12 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
         export = kwargs.pop("export", 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)
 
         with ContextRblnConfig(
             device=device,
@@ -71,7 +72,7 @@ class RBLNStableDiffusionImg2ImgPipeline(StableDiffusionImg2ImgPipeline):
             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
@@ -83,17 +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={**rbln_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={**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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -106,7 +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={**rbln_kwargs},
+            rbln_config={**rbln_config},
         )
 
         if model_save_dir is not None:
@@ -134,9 +135,6 @@ 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 optimize_host_memory is False:
             model.compiled_models = [
                 vae.compiled_models[0],

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

@@ -16,7 +16,7 @@
 
 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
@@ -24,6 +24,7 @@ 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.
@@ -44,12 +45,12 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
         export = kwargs.pop("export", 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)
 
         with ContextRblnConfig(
             device=device,
@@ -63,7 +64,7 @@ class RBLNStableDiffusionXLPipeline(StableDiffusionXLPipeline):
             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
         )
 
         vae = RBLNAutoencoderKL.from_pretrained(
@@ -73,24 +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={**rbln_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={**rbln_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={**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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -103,7 +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={**rbln_kwargs},
+            rbln_config={**rbln_config},
         )
 
         if model_save_dir is not None:
@@ -130,8 +131,6 @@ 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 optimize_host_memory is False:
             model.compiled_models = [
                 vae.compiled_models[0],

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

@@ -16,7 +16,7 @@
 
 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
@@ -24,6 +24,7 @@ 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.
@@ -44,12 +45,12 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
         export = kwargs.pop("export", 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)
 
         with ContextRblnConfig(
             device=device,
@@ -63,7 +64,7 @@ class RBLNStableDiffusionXLImg2ImgPipeline(StableDiffusionXLImg2ImgPipeline):
             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
         )
 
         vae = RBLNAutoencoderKL.from_pretrained(
@@ -74,24 +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={**rbln_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={**rbln_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={**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
 
         unet = RBLNUNet2DConditionModel.from_pretrained(
@@ -105,7 +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={**rbln_kwargs},
+            rbln_config={**rbln_config},
         )
 
         if model_save_dir is not None:
@@ -132,8 +133,6 @@ 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 optimize_host_memory is False:
             model.compiled_models = [
                 vae.compiled_models[0],

optimum/rbln/modeling_alias.py

@@ -28,7 +28,6 @@ from .modeling_base import (
     RBLNModelForQuestionAnswering,
     RBLNModelForSequenceClassification,
 )
-from .modeling_seq2seq import RBLNModelForSeq2SeqLM
 
 
 class RBLNASTForAudioClassification(RBLNModelForAudioClassification):
@@ -47,14 +46,6 @@ class RBLNResNetForImageClassification(RBLNModelForImageClassification):
     pass
 
 
-class RBLNT5ForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    pass
-
-
-class RBLNBartForConditionalGeneration(RBLNModelForSeq2SeqLM):
-    pass
-
-
 class RBLNXLMRobertaForSequenceClassification(RBLNModelForSequenceClassification):
     rbln_model_input_names = ["input_ids", "attention_mask"]
 
@@ -65,3 +56,7 @@ class RBLNRobertaForSequenceClassification(RBLNModelForSequenceClassification):
 
 class RBLNRobertaForMaskedLM(RBLNModelForMaskedLM):
     rbln_model_input_names = ["input_ids", "attention_mask"]
+
+
+class RBLNViTForImageClassification(RBLNModelForImageClassification):
+    pass