optimum-rbln 0.1.4__py3-none-any.whl → 0.1.8__py3-none-any.whl

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

@@ -22,18 +22,18 @@
 # from Rebellions Inc.
 """RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
-from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
 import torch.nn.functional as F
-from diffusers import StableDiffusionControlNetPipeline
+from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionControlNetPipeline
 from diffusers.image_processor import PipelineImageInput
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.pipelines.controlnet.pipeline_controlnet import retrieve_timesteps
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 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 ....transformers import RBLNCLIPTextModel
@@ -64,18 +64,40 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
         """
         export = kwargs.pop("export", None)
+        vae = kwargs.pop("vae", None)
+        unet = kwargs.pop("unet", None)
         text_encoder = kwargs.pop("text_encoder", None)
-        controlnets = kwargs.pop("controlnet", None)
+        controlnet = kwargs.pop("controlnet", None)
+        model_save_dir = kwargs.pop("model_save_dir", None)
 
         rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
 
         kwargs_dict = {
             "pretrained_model_name_or_path": model_id,
-            "text_encoder": text_encoder,
-            "controlnet": controlnets,
             **kwargs,
         }
 
+        kwargs_dict.update(
+            {
+                **({"vae": vae} if vae is not None and isinstance(vae, AutoencoderKL) else {}),
+                **({"unet": unet} if unet is not None and isinstance(unet, UNet2DConditionModel) else {}),
+                **(
+                    {"text_encoder": text_encoder}
+                    if text_encoder is not None and isinstance(text_encoder, CLIPTextModel)
+                    else {}
+                ),
+                **(
+                    {"controlnet": controlnet}
+                    if controlnet is not None
+                    and (
+                        isinstance(controlnet, ControlNetModel)
+                        or all(isinstance(c, ControlNetModel) for c in controlnet)
+                    )
+                    else {}
+                ),
+            }
+        )
+
         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:
@@ -85,64 +107,87 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
             rbln_config_kwargs.pop("rbln_guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
-
-        model.save_pretrained(save_directory=save_dir_path, **kwargs)
-
         # compile model, create runtime
-        vae = RBLNAutoencoderKL.from_pretrained(
-            model_id=save_dir_path / "vae",
-            export=True,
-            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,
-        )
-
-        text_encoder = RBLNCLIPTextModel.from_pretrained(
-            model_id=save_dir_path / "text_encoder",
-            export=True,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
-        )
-
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        unet = RBLNUNet2DConditionModel.from_pretrained(
-            model_id=save_dir_path / "unet",
-            export=True,
-            rbln_max_seq_len=text_encoder.config.max_position_embeddings,
-            rbln_batch_size=unet_batch_size,
-            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,
-        )
-
-        if isinstance(controlnets, (list, tuple)):
-            controlnet = RBLNMultiControlNetModel.from_pretrained(
-                model_id=str(save_dir_path / "controlnet"),
+        if not isinstance(vae, RBLNAutoencoderKL):
+            vae = RBLNAutoencoderKL.from_pretrained(
+                model_id=model_id,
+                subfolder="vae",
                 export=True,
-                rbln_batch_size=unet_batch_size,
+                model_save_dir=model_save_dir,
+                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,
             )
-            controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
-        else:
-            controlnet = RBLNControlNetModel.from_pretrained(
-                model_id=save_dir_path / "controlnet",
+
+        if not isinstance(text_encoder, RBLNCLIPTextModel):
+            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,
+            )
+
+        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
+        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
+
+        if not isinstance(unet, RBLNUNet2DConditionModel):
+            unet = RBLNUNet2DConditionModel.from_pretrained(
+                model_id=model_id,
+                subfolder="unet",
+                export=True,
+                model_save_dir=model_save_dir,
+                rbln_max_seq_len=text_encoder.config.max_position_embeddings,
                 rbln_batch_size=unet_batch_size,
+                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,
             )
-            controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
+
+        if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
+            if isinstance(controlnet, (list, tuple)):
+                multicontrolnet = []
+                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,
+                            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_kwargs,
+                            **rbln_constructor_kwargs,
+                        )
+                    )
+                controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
+                controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
+            else:
+                controlnet = RBLNControlNetModel.from_pretrained(
+                    model_id=controlnet.config._name_or_path,
+                    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_kwargs,
+                    **rbln_constructor_kwargs,
+                )
+                controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
+
+        if model_save_dir is not None:
+            # To skip saving original pytorch modules
+            del (model.vae, model.text_encoder, model.unet, model.controlnet)
+
+            # Direct calling of `save_pretrained` causes config.unet = (None, None).
+            # So config must be saved again, later.
+            model.save_pretrained(model_save_dir)
 
         # replace modules
         model.vae = vae
@@ -159,7 +204,18 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
         }
         model.register_to_config(**update_dict)
 
-        model.models = [vae.model[0], text_encoder.model[0], unet.model[0], controlnet.model[0]]
+        if model_save_dir is not None:
+            # overwrite to replace incorrect config
+            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:
+            model.compiled_models = [vae.compiled_models[0], text_encoder.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
 

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

@@ -22,17 +22,17 @@
 # from Rebellions Inc.
 """RBLNStableDiffusionPipeline class for inference of diffusion models on rbln devices."""
 
-from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
 import torch.nn.functional as F
-from diffusers import StableDiffusionControlNetImg2ImgPipeline
+from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionControlNetImg2ImgPipeline
 from diffusers.image_processor import PipelineImageInput
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 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 ....transformers import RBLNCLIPTextModel
@@ -63,18 +63,40 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
                 - A path to a *directory* containing a model saved using [`~OptimizedModel.save_pretrained`],
         """
         export = kwargs.pop("export", None)
+        vae = kwargs.pop("vae", None)
+        unet = kwargs.pop("unet", None)
         text_encoder = kwargs.pop("text_encoder", None)
-        controlnets = kwargs.pop("controlnet", None)
+        controlnet = kwargs.pop("controlnet", None)
+        model_save_dir = kwargs.pop("model_save_dir", None)
 
         rbln_config_kwargs, rbln_constructor_kwargs = RBLNBaseModel.pop_rbln_kwargs_from_kwargs(kwargs)
 
         kwargs_dict = {
             "pretrained_model_name_or_path": model_id,
-            "text_encoder": text_encoder,
-            "controlnet": controlnets,
             **kwargs,
         }
 
+        kwargs_dict.update(
+            {
+                **({"vae": vae} if vae is not None and isinstance(vae, AutoencoderKL) else {}),
+                **({"unet": unet} if unet is not None and isinstance(unet, UNet2DConditionModel) else {}),
+                **(
+                    {"text_encoder": text_encoder}
+                    if text_encoder is not None and isinstance(text_encoder, CLIPTextModel)
+                    else {}
+                ),
+                **(
+                    {"controlnet": controlnet}
+                    if controlnet is not None
+                    and (
+                        isinstance(controlnet, ControlNetModel)
+                        or all(isinstance(c, ControlNetModel) for c in controlnet)
+                    )
+                    else {}
+                ),
+            }
+        )
+
         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:
@@ -84,64 +106,87 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
             rbln_config_kwargs.pop("rbln_guidance_scale", 5.0) > 1.0 and model.unet.config.time_cond_proj_dim is None
         )
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
-
-        model.save_pretrained(save_directory=save_dir_path, **kwargs)
-
         # compile model, create runtime
-        vae = RBLNAutoencoderKL.from_pretrained(
-            model_id=save_dir_path / "vae",
-            export=True,
-            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,
-        )
-
-        text_encoder = RBLNCLIPTextModel.from_pretrained(
-            model_id=save_dir_path / "text_encoder",
-            export=True,
-            **rbln_config_kwargs,
-            **rbln_constructor_kwargs,
-        )
-
-        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
-        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
-
-        unet = RBLNUNet2DConditionModel.from_pretrained(
-            model_id=save_dir_path / "unet",
-            export=True,
-            rbln_max_seq_len=text_encoder.config.max_position_embeddings,
-            rbln_batch_size=unet_batch_size,
-            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,
-        )
-
-        if isinstance(controlnets, (list, tuple)):
-            controlnet = RBLNMultiControlNetModel.from_pretrained(
-                model_id=str(save_dir_path / "controlnet"),
+        if not isinstance(vae, RBLNAutoencoderKL):
+            vae = RBLNAutoencoderKL.from_pretrained(
+                model_id=model_id,
+                subfolder="vae",
                 export=True,
-                rbln_batch_size=unet_batch_size,
+                model_save_dir=model_save_dir,
+                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,
             )
-            controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
-        else:
-            controlnet = RBLNControlNetModel.from_pretrained(
-                model_id=save_dir_path / "controlnet",
+
+        if not isinstance(text_encoder, RBLNCLIPTextModel):
+            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,
+            )
+
+        batch_size = rbln_config_kwargs.pop("rbln_batch_size", 1)
+        unet_batch_size = batch_size * 2 if do_classifier_free_guidance else batch_size
+
+        if not isinstance(unet, RBLNUNet2DConditionModel):
+            unet = RBLNUNet2DConditionModel.from_pretrained(
+                model_id=model_id,
+                subfolder="unet",
                 export=True,
+                model_save_dir=model_save_dir,
+                rbln_max_seq_len=text_encoder.config.max_position_embeddings,
                 rbln_batch_size=unet_batch_size,
+                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,
             )
-            controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
+
+        if not isinstance(controlnet, (RBLNControlNetModel, RBLNMultiControlNetModel)):
+            if isinstance(controlnet, (list, tuple)):
+                multicontrolnet = []
+                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,
+                            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_kwargs,
+                            **rbln_constructor_kwargs,
+                        )
+                    )
+                controlnet = RBLNMultiControlNetModel(multicontrolnet, config=controlnet[0].config)
+                controlnet_dict = ("optimum.rbln", "RBLNMultiControlNetModel")
+            else:
+                controlnet = RBLNControlNetModel.from_pretrained(
+                    model_id=controlnet.config._name_or_path,
+                    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_kwargs,
+                    **rbln_constructor_kwargs,
+                )
+                controlnet_dict = ("optimum.rbln", "RBLNControlNetModel")
+
+        if model_save_dir is not None:
+            # To skip saving original pytorch modules
+            del (model.vae, model.text_encoder, model.unet, model.controlnet)
+
+            # Direct calling of `save_pretrained` causes config.unet = (None, None).
+            # So config must be saved again, later.
+            model.save_pretrained(model_save_dir)
 
         # replace modules
         model.vae = vae
@@ -158,7 +203,23 @@ class RBLNStableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetImg2
         }
         model.register_to_config(**update_dict)
 
-        model.models = [vae.model[0], vae.model[1], text_encoder.model[0], unet.model[0], controlnet.model[0]]
+        if model_save_dir is not None:
+            # overwrite to replace incorrect config
+            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:
+            model.compiled_models = [
+                vae.compiled_models[0],
+                vae.compiled_models[1],
+                text_encoder.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