optimum-rbln 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

optimum/rbln/__init__.py

@@ -63,6 +63,9 @@ _import_structure = {
         "RBLNStableDiffusionControlNetImg2ImgPipeline",
         "RBLNMultiControlNetModel",
         "RBLNStableDiffusionXLImg2ImgPipeline",
+        "RBLNStableDiffusionControlNetPipeline",
+        "RBLNStableDiffusionXLControlNetPipeline",
+        "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
     ],
     "modeling_config": ["RBLNRuntimeConfig", "RBLNConfig"],
 }
@@ -73,8 +76,11 @@ if TYPE_CHECKING:
         RBLNControlNetModel,
         RBLNMultiControlNetModel,
         RBLNStableDiffusionControlNetImg2ImgPipeline,
+        RBLNStableDiffusionControlNetPipeline,
         RBLNStableDiffusionImg2ImgPipeline,
         RBLNStableDiffusionPipeline,
+        RBLNStableDiffusionXLControlNetImg2ImgPipeline,
+        RBLNStableDiffusionXLControlNetPipeline,
         RBLNStableDiffusionXLImg2ImgPipeline,
         RBLNStableDiffusionXLPipeline,
         RBLNUNet2DConditionModel,

optimum/rbln/__version__.py

@@ -1 +1 @@
-__version__ = '0.1.0'
+__version__ = '0.1.1'

optimum/rbln/diffusers/__init__.py

@@ -39,17 +39,24 @@ _import_structure = {
         "RBLNStableDiffusionControlNetImg2ImgPipeline",
         "RBLNMultiControlNetModel",
         "RBLNStableDiffusionXLImg2ImgPipeline",
+        "RBLNStableDiffusionControlNetPipeline",
+        "RBLNStableDiffusionXLControlNetPipeline",
+        "RBLNStableDiffusionXLControlNetImg2ImgPipeline",
     ],
     "models": ["RBLNAutoencoderKL", "RBLNUNet2DConditionModel", "RBLNControlNetModel"],
 }
 
 if TYPE_CHECKING:
+
     from .models import RBLNAutoencoderKL, RBLNControlNetModel, RBLNUNet2DConditionModel
     from .pipelines import (
         RBLNMultiControlNetModel,
         RBLNStableDiffusionControlNetImg2ImgPipeline,
+        RBLNStableDiffusionControlNetPipeline,
         RBLNStableDiffusionImg2ImgPipeline,
         RBLNStableDiffusionPipeline,
+        RBLNStableDiffusionXLControlNetImg2ImgPipeline,
+        RBLNStableDiffusionXLControlNetPipeline,
         RBLNStableDiffusionXLImg2ImgPipeline,
         RBLNStableDiffusionXLPipeline,
     )

optimum/rbln/diffusers/models/autoencoder_kl.py

@@ -88,14 +88,23 @@ class RBLNAutoencoderKL(RBLNModel):
         subfolder: str = "",
         local_files_only: bool = False,
         trust_remote_code: bool = False,
+        model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
         **kwargs,
     ) -> "RBLNAutoencoderKL":
         task = kwargs.pop("task", None)
         if task is None:
             task = TasksManager.infer_task_from_model(cls.auto_model_class)
 
-        save_dir = TemporaryDirectory()
-        save_dir_path = Path(save_dir.name)
+        if model_save_dir is None:
+            save_dir = TemporaryDirectory()
+            save_dir_path = Path(save_dir.name)
+        else:
+            save_dir = model_save_dir
+            if isinstance(save_dir, TemporaryDirectory):
+                save_dir_path = Path(model_save_dir.name)
+            else:
+                save_dir_path = Path(model_save_dir)
+                save_dir_path.mkdir(exist_ok=True)
 
         rbln_config_kwargs, rbln_constructor_kwargs = cls.pop_rbln_kwargs_from_kwargs(kwargs)
 
@@ -119,7 +128,7 @@ class RBLNAutoencoderKL(RBLNModel):
         if not isinstance(config, PretrainedConfig):  # diffusers config
             config = PretrainedConfig(**config)
 
-        config.save_pretrained(save_dir_path)
+        config.save_pretrained(save_dir_path / subfolder)
         preprocessors = maybe_save_preprocessors(model_id, save_dir_path, src_subfolder=subfolder)
 
         # Get compilation arguments
@@ -137,8 +146,12 @@ class RBLNAutoencoderKL(RBLNModel):
             enc_compiled_model = cls.compile(encoder_model, rbln_runtime_config=rbln_config["encoder"][0])
             dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["decoder"][0])
 
-            enc_compiled_model.save(save_dir_path / f"{rbln_config['encoder'][0].compiled_model_name}.rbln")
-            dec_compiled_model.save(save_dir_path / f"{rbln_config['decoder'][0].compiled_model_name}.rbln")
+            enc_compiled_model.save(
+                save_dir_path / subfolder / f"{rbln_config['encoder'][0].compiled_model_name}.rbln"
+            )
+            dec_compiled_model.save(
+                save_dir_path / subfolder / f"{rbln_config['decoder'][0].compiled_model_name}.rbln"
+            )
 
         def compile_text2img():
             decoder_model = _VAEDecoder(model)
@@ -146,19 +159,27 @@ class RBLNAutoencoderKL(RBLNModel):
 
             dec_compiled_model = cls.compile(decoder_model, rbln_runtime_config=rbln_config["compiled_model"][0])
 
-            dec_compiled_model.save(save_dir_path / f"{rbln_config['compiled_model'][0].compiled_model_name}.rbln")
+            dec_compiled_model.save(
+                save_dir_path / subfolder / f"{rbln_config['compiled_model'][0].compiled_model_name}.rbln"
+            )
 
         if rbln_config_kwargs.get("rbln_use_encode"):
             compile_img2img()
         else:
             compile_text2img()
 
-        rbln_config.save(save_dir_path)
+        rbln_config.save(save_dir_path / subfolder)
 
         return cls._from_pretrained(
             model_id=save_dir_path,
             config=config,
             model_save_dir=save_dir,
+            use_auth_token=use_auth_token,
+            revision=revision,
+            force_download=force_download,
+            cache_dir=cache_dir,
+            subfolder=subfolder,
+            local_files_only=local_files_only,
             **rbln_constructor_kwargs,
             **kwargs,
         )
@@ -216,7 +237,7 @@ class RBLNAutoencoderKL(RBLNModel):
             meta["rbln_img_height"] = rbln_img_height
 
             vae_enc_input_info = [
-                ("x", [rbln_batch_size, model_config.in_channels, rbln_img_width, rbln_img_height], "float32")
+                ("x", [rbln_batch_size, model_config.in_channels, rbln_img_height, rbln_img_width], "float32")
             ]
             vae_dec_input_info = [
                 (
@@ -224,8 +245,8 @@ class RBLNAutoencoderKL(RBLNModel):
                     [
                         rbln_batch_size,
                         model_config.latent_channels,
-                        rbln_img_width // rbln_vae_scale_factor,
                         rbln_img_height // rbln_vae_scale_factor,
+                        rbln_img_width // rbln_vae_scale_factor,
                     ],
                     "float32",
                 )

optimum/rbln/diffusers/models/controlnet.py

@@ -23,9 +23,8 @@
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Dict, Optional, Union
 
-import rebel
 import torch
 from diffusers import ControlNetModel
 from optimum.exporters import TasksManager
@@ -46,6 +45,37 @@ class _ControlNetModel(torch.nn.Module):
         super().__init__()
         self.controlnet = controlnet
 
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: torch.Tensor,
+        controlnet_cond: torch.Tensor,
+        conditioning_scale,
+        text_embeds: Optional[torch.Tensor] = None,
+        time_ids: Optional[torch.Tensor] = None,
+    ):
+        if text_embeds is not None and time_ids is not None:
+            added_cond_kwargs = {"text_embeds": text_embeds, "time_ids": time_ids}
+        else:
+            added_cond_kwargs = {}
+
+        down_block_res_samples, mid_block_res_sample = self.controlnet(
+            sample=sample,
+            timestep=timestep,
+            encoder_hidden_states=None,
+            controlnet_cond=controlnet_cond,
+            conditioning_scale=conditioning_scale,
+            added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )
+        return down_block_res_samples, mid_block_res_sample
+
+
+class _ControlNetModel_Cross_Attention(torch.nn.Module):
+    def __init__(self, controlnet: "ControlNetModel"):
+        super().__init__()
+        self.controlnet = controlnet
+
     def forward(
         self,
         sample: torch.Tensor,
@@ -53,13 +83,21 @@ class _ControlNetModel(torch.nn.Module):
         encoder_hidden_states: torch.Tensor,
         controlnet_cond: torch.Tensor,
         conditioning_scale,
+        text_embeds: Optional[torch.Tensor] = None,
+        time_ids: Optional[torch.Tensor] = None,
     ):
+        if text_embeds is not None and time_ids is not None:
+            added_cond_kwargs = {"text_embeds": text_embeds, "time_ids": time_ids}
+        else:
+            added_cond_kwargs = {}
+
         down_block_res_samples, mid_block_res_sample = self.controlnet(
             sample=sample,
             timestep=timestep,
             encoder_hidden_states=encoder_hidden_states,
             controlnet_cond=controlnet_cond,
             conditioning_scale=conditioning_scale,
+            added_cond_kwargs=added_cond_kwargs,
             return_dict=False,
         )
         return down_block_res_samples, mid_block_res_sample
@@ -71,6 +109,9 @@ class RBLNControlNetModel(RBLNModel):
 
     def __post_init__(self, **kwargs):
         self.dtype = torch.float32
+        self.use_encoder_hidden_states = any(
+            item[0] == "encoder_hidden_states" for item in self.rbln_config["compiled_model"][0].input_info
+        )
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
@@ -94,14 +135,16 @@ class RBLNControlNetModel(RBLNModel):
         return rt
 
     @classmethod
-    def compile(cls, model, rbln_runtime_config: Optional[RBLNRuntimeConfig] = None):
-        compiled_model = rebel.compile_from_torch(
-            _ControlNetModel(model),
-            input_info=rbln_runtime_config.input_info,
-            batch_size=rbln_runtime_config.batch_size,
-            fusion=rbln_runtime_config.fusion,
-        )
-        return compiled_model
+    def wrap_model_if_needed(cls, model: torch.nn.Module) -> torch.nn.Module:
+        use_encoder_hidden_states = False
+        for down_block in model.down_blocks:
+            if use_encoder_hidden_states := getattr(down_block, "has_cross_attention", False):
+                break
+
+        if use_encoder_hidden_states:
+            return _ControlNetModel_Cross_Attention(model).eval()
+        else:
+            return _ControlNetModel(model).eval()
 
     @classmethod
     def _get_rbln_config(
@@ -109,6 +152,7 @@ class RBLNControlNetModel(RBLNModel):
         preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
         model_config: "PretrainedConfig",
         rbln_max_seq_len: Optional[int] = None,
+        rbln_text_model_hidden_size: Optional[int] = None,
         rbln_batch_size: Optional[int] = None,
         rbln_img_width: Optional[int] = None,
         rbln_img_height: Optional[int] = None,
@@ -132,12 +176,18 @@ class RBLNControlNetModel(RBLNModel):
                     [
                         rbln_batch_size,
                         model_config.in_channels,
-                        input_width,
                         input_height,
+                        input_width,
                     ],
                     "float32",
                 ),
                 ("timestep", [], "float32"),
+            ],
+            batch_size=rbln_batch_size,
+        )
+        use_encoder_hidden_states = any(element != "DownBlock2D" for element in model_config.down_block_types)
+        if use_encoder_hidden_states:
+            rbln_runtime_config.input_info.append(
                 (
                     "encoder_hidden_states",
                     [
@@ -146,12 +196,20 @@ class RBLNControlNetModel(RBLNModel):
                         model_config.cross_attention_dim,
                     ],
                     "float32",
-                ),
-                ("controlnet_cond", [rbln_batch_size, 3, rbln_img_width, rbln_img_height], "float32"),
-                ("conditioning_scale", [], "float32"),
-            ],
-            batch_size=rbln_batch_size,
+                )
+            )
+        rbln_runtime_config.input_info.append(
+            ("controlnet_cond", [rbln_batch_size, 3, rbln_img_height, rbln_img_width], "float32")
         )
+        rbln_runtime_config.input_info.append(("conditioning_scale", [], "float32"))
+        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
+            rbln_runtime_config.input_info.append(
+                ("text_embeds", [rbln_batch_size, rbln_text_model_hidden_size], "float32")
+            )
+            rbln_runtime_config.input_info.append(("time_ids", [rbln_batch_size, 6], "float32"))
+
         rbln_config = RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
         return rbln_config
 
@@ -162,18 +220,30 @@ class RBLNControlNetModel(RBLNModel):
         encoder_hidden_states: torch.Tensor,
         controlnet_cond: torch.FloatTensor,
         conditioning_scale: torch.Tensor = 1.0,
+        added_cond_kwargs: Dict[str, torch.Tensor] = {},
         **kwargs,
     ):
         """
         The [`ControlNetModel`] forward method.
         """
-        output = super().forward(
-            sample.contiguous(),
-            timestep.float(),
-            encoder_hidden_states,
-            controlnet_cond,
-            torch.tensor(conditioning_scale),
-        )
+        added_cond_kwargs = {} if added_cond_kwargs is None else added_cond_kwargs
+        if self.use_encoder_hidden_states:
+            output = super().forward(
+                sample.contiguous(),
+                timestep.float(),
+                encoder_hidden_states,
+                controlnet_cond,
+                torch.tensor(conditioning_scale),
+                **added_cond_kwargs,
+            )
+        else:
+            output = super().forward(
+                sample.contiguous(),
+                timestep.float(),
+                controlnet_cond,
+                torch.tensor(conditioning_scale),
+                **added_cond_kwargs,
+            )
         down_block_res_samples = output[:-1]
         mid_block_res_sample = output[-1]
 

optimum/rbln/diffusers/models/unet_2d_condition.py

@@ -27,7 +27,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
-from diffusers.models.unet_2d_condition import UNet2DConditionModel
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
 from optimum.exporters import TasksManager
 from transformers import AutoConfig, AutoModel, PretrainedConfig
 
@@ -90,22 +90,28 @@ class _UNet_SDXL(torch.nn.Module):
         sample: torch.Tensor,
         timestep: Union[torch.Tensor, float, int],
         encoder_hidden_states: torch.Tensor,
-        text_embeds: Optional[torch.Tensor] = None,
-        time_ids: Optional[torch.Tensor] = None,
         *down_and_mid_block_additional_residuals: Optional[Tuple[torch.Tensor]],
     ) -> torch.Tensor:
-        if text_embeds is not None and time_ids is not None:
-            added_cond_kwargs = {"text_embeds": text_embeds, "time_ids": time_ids}
-        else:
-            added_cond_kwargs = {}
-
-        if len(down_and_mid_block_additional_residuals) != 0:
+        if len(down_and_mid_block_additional_residuals) == 2:
+            added_cond_kwargs = {
+                "text_embeds": down_and_mid_block_additional_residuals[0],
+                "time_ids": down_and_mid_block_additional_residuals[1],
+            }
+            down_block_additional_residuals = None
+            mid_block_additional_residual = None
+        elif len(down_and_mid_block_additional_residuals) > 2:
+            added_cond_kwargs = {
+                "text_embeds": down_and_mid_block_additional_residuals[-2],
+                "time_ids": down_and_mid_block_additional_residuals[-1],
+            }
             down_block_additional_residuals, mid_block_additional_residual = (
-                down_and_mid_block_additional_residuals[:-1],
-                down_and_mid_block_additional_residuals[-1],
+                down_and_mid_block_additional_residuals[:-3],
+                down_and_mid_block_additional_residuals[-3],
             )
         else:
-            down_block_additional_residuals, mid_block_additional_residual = None, None
+            added_cond_kwargs = {}
+            down_block_additional_residuals = None
+            mid_block_additional_residual = None
 
         unet_out = self.unet(
             sample=sample,
@@ -197,9 +203,11 @@ class RBLNUNet2DConditionModel(RBLNModel):
         meta["rbln_use_encode"] = rbln_use_encode
 
         if rbln_use_encode:
+            # FIXME :: robust img shape getter
             input_width = rbln_img_width // rbln_vae_scale_factor
             input_height = rbln_img_height // rbln_vae_scale_factor
         else:
+            # FIXME :: model_config.sample_size can be tuple or list
             input_width, input_height = model_config.sample_size, model_config.sample_size
 
         input_info = [
@@ -208,8 +216,8 @@ class RBLNUNet2DConditionModel(RBLNModel):
                 [
                     rbln_batch_size,
                     model_config.in_channels,
-                    input_width,
                     input_height,
+                    input_width,
                 ],
                 "float32",
             ),
@@ -225,64 +233,73 @@ class RBLNUNet2DConditionModel(RBLNModel):
             ),
         ]
         if rbln_is_controlnet:
-            input_info.extend(
-                [
-                    (
-                        f"down_block_additional_residuals_{i}",
-                        [rbln_batch_size, model_config.block_out_channels[0], input_width, input_height],
-                        "float32",
-                    )
-                    for i in range(3)
-                ]
-            )
-            input_info.append(
-                (
-                    f"down_block_additional_residuals_{3}",
-                    [rbln_batch_size, model_config.block_out_channels[0], input_width // 2, input_height // 2],
-                    "float32",
+            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)
+                    ]
                 )
-            )
-            input_info.extend(
-                [
+                input_info.append(
                     (
-                        f"down_block_additional_residuals_{i}",
-                        [rbln_batch_size, model_config.block_out_channels[1], input_width // 2, input_height // 2],
+                        "down_block_additional_residuals_3",
+                        [rbln_batch_size, model_config.block_out_channels[0], input_height // 2, input_width // 2],
                         "float32",
                     )
-                    for i in range(4, 6)
-                ]
-            )
-            input_info.append(
-                (
-                    f"down_block_additional_residuals_{6}",
-                    [rbln_batch_size, model_config.block_out_channels[1], input_width // 4, input_height // 4],
-                    "float32",
                 )
-            )
-            input_info.extend(
-                [
-                    (
-                        f"down_block_additional_residuals_{i}",
-                        [rbln_batch_size, model_config.block_out_channels[2], input_width // 4, input_height // 4],
-                        "float32",
-                    )
-                    for i in range(7, 9)
-                ]
-            )
-            input_info.extend(
-                [
+            if len(model_config.block_out_channels) > 1:
+                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)
+                    ]
+                )
+                input_info.append(
                     (
-                        f"down_block_additional_residuals_{i}",
-                        [rbln_batch_size, model_config.block_out_channels[3], input_width // 8, input_height // 8],
+                        f"down_block_additional_residuals_{6}",
+                        [rbln_batch_size, model_config.block_out_channels[1], input_height // 4, input_width // 4],
                         "float32",
                     )
-                    for i in range(9, 12)
-                ]
-            )
+                )
+            if len(model_config.block_out_channels) > 2:
+                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[3], input_width // 8, input_height // 8],
+                    [
+                        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",
                 )
             )
@@ -344,7 +361,7 @@ class RBLNUNet2DConditionModel(RBLNModel):
 
         return (
             super().forward(
-                sample,
+                sample.contiguous(),
                 timestep.float(),
                 encoder_hidden_states,
                 **added_cond_kwargs,

optimum/rbln/diffusers/pipelines/__init__.py

@@ -21,9 +21,14 @@
 # copied, modified, or distributed without prior written permission
 # from Rebellions Inc.
 
-from .controlnet import RBLNMultiControlNetModel
-from .stable_diffusion import (
+from .controlnet import (
+    RBLNMultiControlNetModel,
     RBLNStableDiffusionControlNetImg2ImgPipeline,
+    RBLNStableDiffusionControlNetPipeline,
+    RBLNStableDiffusionXLControlNetImg2ImgPipeline,
+    RBLNStableDiffusionXLControlNetPipeline,
+)
+from .stable_diffusion import (
     RBLNStableDiffusionImg2ImgPipeline,
     RBLNStableDiffusionPipeline,
 )

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

@@ -22,3 +22,7 @@
 # from Rebellions Inc.
 
 from .multicontrolnet import RBLNMultiControlNetModel
+from .pipeline_controlnet import RBLNStableDiffusionControlNetPipeline
+from .pipeline_controlnet_img2img import RBLNStableDiffusionControlNetImg2ImgPipeline
+from .pipeline_controlnet_sd_xl import RBLNStableDiffusionXLControlNetPipeline
+from .pipeline_controlnet_sd_xl_img2img import RBLNStableDiffusionXLControlNetImg2ImgPipeline