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

optimum/rbln/diffusers/models/unet_2d_condition.py

@@ -23,16 +23,15 @@
 
 import logging
 from dataclasses import dataclass
-from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
 from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
-from optimum.exporters import TasksManager
-from transformers import AutoConfig, AutoModel, PretrainedConfig
+from transformers import PretrainedConfig
 
 from ...modeling_base import RBLNModel
 from ...modeling_config import RBLNCompileConfig, RBLNConfig
+from ...utils.context import override_auto_classes
 
 
 if TYPE_CHECKING:
@@ -126,9 +125,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)
@@ -146,29 +142,11 @@ class RBLNUNet2DConditionModel(RBLNModel):
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
-        def get_model_from_task(
-            task: str,
-            model_name_or_path: Union[str, Path],
-            **kwargs,
+        with override_auto_classes(
+            config_func=UNet2DConditionModel.load_config,
+            model_func=UNet2DConditionModel.from_pretrained,
         ):
-            return UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path=model_name_or_path, **kwargs)
-
-        tasktmp = TasksManager.get_model_from_task
-        configtmp = AutoConfig.from_pretrained
-        modeltmp = AutoModel.from_pretrained
-        TasksManager.get_model_from_task = get_model_from_task
-        if kwargs.get("export", None):
-            # This is an ad-hoc to workaround save null values of the config.
-            # if export, pure optimum(not optimum-rbln) loads config using AutoConfig
-            # and diffusers model do not support loading by AutoConfig.
-            AutoConfig.from_pretrained = lambda *args, **kwargs: None
-        else:
-            AutoConfig.from_pretrained = UNet2DConditionModel.load_config
-        AutoModel.from_pretrained = UNet2DConditionModel.from_pretrained
-        rt = super().from_pretrained(*args, **kwargs)
-        AutoConfig.from_pretrained = configtmp
-        AutoModel.from_pretrained = modeltmp
-        TasksManager.get_model_from_task = tasktmp
+            rt = super().from_pretrained(*args, **kwargs)
         return rt
 
     @classmethod
@@ -185,137 +163,68 @@ class RBLNUNet2DConditionModel(RBLNModel):
         model_config: "PretrainedConfig",
         rbln_kwargs: Dict[str, Any] = {},
     ) -> RBLNConfig:
-        rbln_max_seq_len = rbln_kwargs.get("max_seq_len", None)
-        rbln_text_model_hidden_size = rbln_kwargs.get("text_model_hidden_size", None)
-        rbln_batch_size = rbln_kwargs.get("batch_size", None)
-        rbln_in_features = rbln_kwargs.get("in_features", None)
-        rbln_use_encode = rbln_kwargs.get("use_encode", None)
-        rbln_img_width = rbln_kwargs.get("img_width", None)
-        rbln_img_height = rbln_kwargs.get("img_height", None)
-        rbln_vae_scale_factor = rbln_kwargs.get("vae_scale_factor", None)
-        rbln_is_controlnet = rbln_kwargs.get("is_controlnet", None)
-
-        if rbln_max_seq_len is None:
-            rbln_max_seq_len = 77
-        if rbln_batch_size is None:
-            rbln_batch_size = 1
-
-        if rbln_use_encode:
-            if rbln_img_width is None or rbln_img_height is None or rbln_vae_scale_factor is None:
-                raise ValueError(
-                    "rbln_img_width, rbln_img_height, and rbln_vae_scale_factor must be provided when rbln_use_encode is True"
-                )
-            input_width = rbln_img_width // rbln_vae_scale_factor
-            input_height = rbln_img_height // rbln_vae_scale_factor
-        else:
-            input_width, input_height = model_config.sample_size, model_config.sample_size
+        batch_size = rbln_kwargs.get("batch_size")
+        max_seq_len = rbln_kwargs.get("max_seq_len")
+        sample_size = rbln_kwargs.get("sample_size")
+        is_controlnet = rbln_kwargs.get("is_controlnet")
+        rbln_in_features = None
+
+        if batch_size is None:
+            batch_size = 1
+
+        if sample_size is None:
+            sample_size = model_config.sample_size
+
+        if isinstance(sample_size, int):
+            sample_size = (sample_size, sample_size)
+
+        if max_seq_len is None:
+            raise ValueError("`rbln_max_seq_len` (ex. text_encoder's max_position_embeddings )must be specified")
 
         input_info = [
-            (
-                "sample",
-                [
-                    rbln_batch_size,
-                    model_config.in_channels,
-                    input_height,
-                    input_width,
-                ],
-                "float32",
-            ),
+            ("sample", [batch_size, model_config.in_channels, sample_size[0], sample_size[1]], "float32"),
             ("timestep", [], "float32"),
-            (
-                "encoder_hidden_states",
-                [
-                    rbln_batch_size,
-                    rbln_max_seq_len,
-                    model_config.cross_attention_dim,
-                ],
-                "float32",
-            ),
+            ("encoder_hidden_states", [batch_size, max_seq_len, model_config.cross_attention_dim], "float32"),
         ]
 
-        if rbln_is_controlnet:
-            if len(model_config.block_out_channels) > 0:
-                input_info.extend(
-                    [
-                        (
-                            f"down_block_additional_residuals_{i}",
-                            [rbln_batch_size, model_config.block_out_channels[0], input_height, input_width],
-                            "float32",
-                        )
-                        for i in range(3)
-                    ]
-                )
-            if len(model_config.block_out_channels) > 1:
-                input_info.append(
-                    (
-                        "down_block_additional_residuals_3",
-                        [rbln_batch_size, model_config.block_out_channels[0], input_height // 2, input_width // 2],
-                        "float32",
-                    )
-                )
-                input_info.extend(
-                    [
-                        (
-                            f"down_block_additional_residuals_{i}",
-                            [rbln_batch_size, model_config.block_out_channels[1], input_height // 2, input_width // 2],
-                            "float32",
-                        )
-                        for i in range(4, 6)
-                    ]
-                )
-            if len(model_config.block_out_channels) > 2:
-                input_info.append(
-                    (
-                        f"down_block_additional_residuals_{6}",
-                        [rbln_batch_size, model_config.block_out_channels[1], input_height // 4, input_width // 4],
-                        "float32",
-                    )
-                )
-                input_info.extend(
-                    [
-                        (
-                            f"down_block_additional_residuals_{i}",
-                            [rbln_batch_size, model_config.block_out_channels[2], input_height // 4, input_width // 4],
-                            "float32",
-                        )
-                        for i in range(7, 9)
-                    ]
-                )
-            if len(model_config.block_out_channels) > 3:
-                input_info.extend(
-                    [
-                        (
-                            f"down_block_additional_residuals_{i}",
-                            [rbln_batch_size, model_config.block_out_channels[3], input_height // 8, input_width // 8],
-                            "float32",
-                        )
-                        for i in range(9, 12)
-                    ]
-                )
-            input_info.append(
-                (
-                    "mid_block_additional_residual",
-                    [
-                        rbln_batch_size,
-                        model_config.block_out_channels[-1],
-                        input_height // 2 ** (len(model_config.block_out_channels) - 1),
-                        input_width // 2 ** (len(model_config.block_out_channels) - 1),
-                    ],
-                    "float32",
-                )
-            )
+        if is_controlnet:
+            # down block addtional residuals
+            first_shape = [batch_size, model_config.block_out_channels[0], sample_size[0], sample_size[1]]
+            height, width = sample_size[0], sample_size[1]
+            input_info.append(("down_block_additional_residuals_0", first_shape, "float32"))
+            name_idx = 1
+            for idx, _ in enumerate(model_config.down_block_types):
+                shape = [batch_size, model_config.block_out_channels[idx], height, width]
+                for _ in range(model_config.layers_per_block):
+                    input_info.append((f"down_block_additional_residuals_{name_idx}", shape, "float32"))
+                    name_idx += 1
+                if idx != len(model_config.down_block_types) - 1:
+                    height = height // 2
+                    width = width // 2
+                    shape = [batch_size, model_config.block_out_channels[idx], height, width]
+                    input_info.append((f"down_block_additional_residuals_{name_idx}", shape, "float32"))
+                    name_idx += 1
+
+            # mid block addtional residual
+            num_cross_attn_blocks = model_config.down_block_types.count("CrossAttnDownBlock2D")
+            out_channels = model_config.block_out_channels[-1]
+            shape = [
+                batch_size,
+                out_channels,
+                sample_size[0] // 2**num_cross_attn_blocks,
+                sample_size[1] // 2**num_cross_attn_blocks,
+            ]
+            input_info.append(("mid_block_additional_residual", shape, "float32"))
 
         rbln_compile_config = RBLNCompileConfig(input_info=input_info)
 
         if hasattr(model_config, "addition_embed_type") and model_config.addition_embed_type == "text_time":
-            if rbln_text_model_hidden_size is None:
-                rbln_text_model_hidden_size = 768
-            if rbln_in_features is None:
-                rbln_in_features = model_config.projection_class_embeddings_input_dim
+            rbln_text_model_hidden_size = rbln_kwargs["text_model_hidden_size"]
+            rbln_in_features = model_config.projection_class_embeddings_input_dim
             rbln_compile_config.input_info.append(
-                ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")
+                ("text_embeds", [batch_size, rbln_text_model_hidden_size], "float32")
             )
-            rbln_compile_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+            rbln_compile_config.input_info.append(("time_ids", [batch_size, 6], "float32"))
 
         rbln_config = RBLNConfig(
             rbln_cls=cls.__name__,
@@ -323,14 +232,6 @@ class RBLNUNet2DConditionModel(RBLNModel):
             rbln_kwargs=rbln_kwargs,
         )
 
-        rbln_config.model_cfg.update(
-            {
-                "max_seq_len": rbln_max_seq_len,
-                "batch_size": rbln_batch_size,
-                "use_encode": rbln_use_encode,
-            }
-        )
-
         if rbln_in_features is not None:
             rbln_config.model_cfg["in_features"] = rbln_in_features
 

optimum/rbln/diffusers/pipelines/__init__.py

@@ -20,16 +20,44 @@
 # are the intellectual property of Rebellions Inc. and may not be
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
+from typing import TYPE_CHECKING
 
-from .controlnet import (
-    RBLNMultiControlNetModel,
-    RBLNStableDiffusionControlNetImg2ImgPipeline,
-    RBLNStableDiffusionControlNetPipeline,
-    RBLNStableDiffusionXLControlNetImg2ImgPipeline,
-    RBLNStableDiffusionXLControlNetPipeline,
-)
-from .stable_diffusion import (
-    RBLNStableDiffusionImg2ImgPipeline,
-    RBLNStableDiffusionPipeline,
-)
-from .stable_diffusion_xl import RBLNStableDiffusionXLImg2ImgPipeline, RBLNStableDiffusionXLPipeline
+from transformers.utils import _LazyModule
+
+
+_import_structure = {
+    "controlnet": [
+        "RBLNMultiControlNetModel",
+        "RBLNStableDiffusionControlNetImg2ImgPipeline",
+        "RBLNStableDiffusionControlNetPipeline",
+        "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
+        "RBLNStableDiffusionXLControlNetPipeline",
+    ],
+    "stable_diffusion": [
+        "RBLNStableDiffusionImg2ImgPipeline",
+        "RBLNStableDiffusionPipeline",
+    ],
+    "stable_diffusion_xl": ["RBLNStableDiffusionXLImg2ImgPipeline", "RBLNStableDiffusionXLPipeline"],
+}
+if TYPE_CHECKING:
+    from .controlnet import (
+        RBLNMultiControlNetModel,
+        RBLNStableDiffusionControlNetImg2ImgPipeline,
+        RBLNStableDiffusionControlNetPipeline,
+        RBLNStableDiffusionXLControlNetImg2ImgPipeline,
+        RBLNStableDiffusionXLControlNetPipeline,
+    )
+    from .stable_diffusion import (
+        RBLNStableDiffusionImg2ImgPipeline,
+        RBLNStableDiffusionPipeline,
+    )
+    from .stable_diffusion_xl import RBLNStableDiffusionXLImg2ImgPipeline, RBLNStableDiffusionXLPipeline
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )

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

@@ -52,6 +52,13 @@ class RBLNMultiControlNetModel(RBLNModel):
         self.nets = models
         self.dtype = torch.float32
 
+    @property
+    def compiled_models(self):
+        cm = []
+        for net in self.nets:
+            cm.extend(net.compiled_models)
+        return cm
+
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
         def get_model_from_task(
@@ -102,6 +109,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

@@ -26,205 +26,25 @@ from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
 import torch.nn.functional as F
-from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionControlNetPipeline
+from diffusers import 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
-from ....utils.runtime_utils import ContextRblnConfig
-from ...models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
+from ....modeling_diffusers import RBLNDiffusionMixin
+from ....utils.decorator_utils import remove_compile_time_kwargs
+from ...models import RBLNControlNetModel
 from ...pipelines.controlnet.multicontrolnet import RBLNMultiControlNetModel
 
 
 logger = logging.get_logger(__name__)
 
 
-class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
-    @classmethod
-    def from_pretrained(cls, model_id, **kwargs):
-        """
-        Pipeline for text-to-image generation using Stable Diffusion with ControlNet.
-
-        This model inherits from [`StableDiffusionControlNetPipeline`]. Check the superclass documentation for the generic methods
-        implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-        It implements the methods to convert a pre-trained Stable Diffusion Controlnet pipeline into a RBLNStableDiffusionControlNet pipeline by:
-        - transferring the checkpoint weights of the original into an optimized RBLN graph,
-        - compiling the resulting graph using the RBLN compiler.
-
-        Args:
-            model_id (`Union[str, Path]`):
-                Can be either:
-                - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
-                - 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)
-        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)
-
-        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)
-
-        kwargs_dict = {
-            "pretrained_model_name_or_path": model_id,
-            **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 {}
-                ),
-            }
-        )
-
-        with ContextRblnConfig(
-            device=device,
-            device_map=device_map,
-            create_runtimes=create_runtimes,
-            optimze_host_mem=optimize_host_memory,
-        ):
-            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:
-            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
-        )
-
-        # compile model, create runtime
-        if not isinstance(vae, RBLNAutoencoderKL):
-            vae = RBLNAutoencoderKL.from_pretrained(
-                model_id=model_id,
-                subfolder="vae",
-                export=True,
-                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={**rbln_kwargs},
-            )
-
-        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={**rbln_kwargs},
-            )
-
-        batch_size = rbln_kwargs.pop("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={**rbln_kwargs},
-            )
-
-        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={**rbln_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={**rbln_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
-        model.text_encoder = text_encoder
-        model.unet = unet
-        model.controlnet = controlnet
-
-        # update config to be able to load from file.
-        update_dict = {
-            "vae": ("optimum.rbln", "RBLNAutoencoderKL"),
-            "text_encoder": ("optimum.rbln", "RBLNCLIPTextModel"),
-            "unet": ("optimum.rbln", "RBLNUNet2DConditionModel"),
-            "controlnet": controlnet_dict,
-        }
-        model.register_to_config(**update_dict)
-
-        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 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):
-                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
+class RBLNStableDiffusionControlNetPipeline(RBLNDiffusionMixin, StableDiffusionControlNetPipeline):
+    original_class = StableDiffusionControlNetPipeline
+    _submodules = ["text_encoder", "unet", "vae", "controlnet"]
 
     def check_inputs(
         self,
@@ -390,6 +210,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
                 )
 
     @torch.no_grad()
+    @remove_compile_time_kwargs
     def __call__(
         self,
         prompt: Union[str, List[str]] = None,
@@ -599,6 +420,7 @@ class RBLNStableDiffusionControlNetPipeline(StableDiffusionControlNetPipeline):
         text_encoder_lora_scale = (
             self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
         )
+
         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
             prompt,
             device,