optimum-rbln 0.1.0__py3-none-any.whl

optimum/rbln/diffusers/models/controlnet.py

@@ -0,0 +1,180 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional, Union
+
+import rebel
+import torch
+from diffusers import ControlNetModel
+from optimum.exporters import TasksManager
+from transformers import AutoConfig, AutoModel, PretrainedConfig
+
+from ...modeling_base import RBLNModel
+from ...modeling_config import RBLNConfig, RBLNRuntimeConfig
+
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
+
+
+class _ControlNetModel(torch.nn.Module):
+    def __init__(self, controlnet: "ControlNetModel"):
+        super().__init__()
+        self.controlnet = controlnet
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        controlnet_cond: torch.Tensor,
+        conditioning_scale,
+    ):
+        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,
+            return_dict=False,
+        )
+        return down_block_res_samples, mid_block_res_sample
+
+
+class RBLNControlNetModel(RBLNModel):
+    model_type = "rbln_model"
+    auto_model_class = AutoModel  # feature extraction
+
+    def __post_init__(self, **kwargs):
+        self.dtype = torch.float32
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        def get_model_from_task(
+            task: str,
+            model_name_or_path: Union[str, Path],
+            **kwargs,
+        ):
+            return ControlNetModel.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
+        AutoConfig.from_pretrained = ControlNetModel.load_config
+        AutoModel.from_pretrained = ControlNetModel.from_pretrained
+        rt = super().from_pretrained(*args, **kwargs)
+        AutoConfig.from_pretrained = configtmp
+        AutoModel.from_pretrained = modeltmp
+        TasksManager.get_model_from_task = tasktmp
+        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
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        preprocessors: Union["AutoFeatureExtractor", "AutoProcessor", "AutoTokenizer"],
+        model_config: "PretrainedConfig",
+        rbln_max_seq_len: Optional[int] = None,
+        rbln_batch_size: Optional[int] = None,
+        rbln_img_width: Optional[int] = None,
+        rbln_img_height: Optional[int] = None,
+        rbln_vae_scale_factor: Optional[int] = None,
+    ) -> RBLNConfig:
+        meta = {"type": "controlnet"}
+
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        if rbln_max_seq_len is None:
+            rbln_max_seq_len = 77
+
+        input_width = rbln_img_width // rbln_vae_scale_factor
+        input_height = rbln_img_height // rbln_vae_scale_factor
+
+        rbln_runtime_config = RBLNRuntimeConfig(
+            input_info=[
+                (
+                    "sample",
+                    [
+                        rbln_batch_size,
+                        model_config.in_channels,
+                        input_width,
+                        input_height,
+                    ],
+                    "float32",
+                ),
+                ("timestep", [], "float32"),
+                (
+                    "encoder_hidden_states",
+                    [
+                        rbln_batch_size,
+                        rbln_max_seq_len,
+                        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_config = RBLNConfig.from_rbln_runtime_configs([rbln_runtime_config], _rbln_meta=meta)
+        return rbln_config
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        controlnet_cond: torch.FloatTensor,
+        conditioning_scale: torch.Tensor = 1.0,
+        **kwargs,
+    ):
+        """
+        The [`ControlNetModel`] forward method.
+        """
+        output = super().forward(
+            sample.contiguous(),
+            timestep.float(),
+            encoder_hidden_states,
+            controlnet_cond,
+            torch.tensor(conditioning_scale),
+        )
+        down_block_res_samples = output[:-1]
+        mid_block_res_sample = output[-1]
+
+        return down_block_res_samples, mid_block_res_sample

optimum/rbln/diffusers/models/unet_2d_condition.py

@@ -0,0 +1,352 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+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.unet_2d_condition import UNet2DConditionModel
+from optimum.exporters import TasksManager
+from transformers import AutoConfig, AutoModel, PretrainedConfig
+
+from ...modeling_base import RBLNModel
+from ...modeling_config import RBLNConfig, RBLNRuntimeConfig
+
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from transformers import AutoFeatureExtractor, AutoProcessor, AutoTokenizer
+
+
+class _UNet_SD(torch.nn.Module):
+    def __init__(self, unet: "UNet2DConditionModel"):
+        super().__init__()
+        self.unet = unet
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        *down_and_mid_block_additional_residuals: Optional[Tuple[torch.Tensor]],
+        text_embeds: Optional[torch.Tensor] = None,
+        time_ids: Optional[torch.Tensor] = None,
+    ) -> 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:
+            down_block_additional_residuals, mid_block_additional_residual = (
+                down_and_mid_block_additional_residuals[:-1],
+                down_and_mid_block_additional_residuals[-1],
+            )
+        else:
+            down_block_additional_residuals, mid_block_additional_residual = None, None
+
+        unet_out = self.unet(
+            sample=sample,
+            timestep=timestep,
+            encoder_hidden_states=encoder_hidden_states,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+            added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )
+        return unet_out
+
+
+class _UNet_SDXL(torch.nn.Module):
+    def __init__(self, unet: "UNet2DConditionModel"):
+        super().__init__()
+        self.unet = unet
+
+    def forward(
+        self,
+        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:
+            down_block_additional_residuals, mid_block_additional_residual = (
+                down_and_mid_block_additional_residuals[:-1],
+                down_and_mid_block_additional_residuals[-1],
+            )
+        else:
+            down_block_additional_residuals, mid_block_additional_residual = None, None
+
+        unet_out = self.unet(
+            sample=sample,
+            timestep=timestep,
+            encoder_hidden_states=encoder_hidden_states,
+            down_block_additional_residuals=down_block_additional_residuals,
+            mid_block_additional_residual=mid_block_additional_residual,
+            added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )
+        return unet_out
+
+
+class RBLNUNet2DConditionModel(RBLNModel):
+    model_type = "rbln_model"
+    auto_model_class = AutoModel  # feature extraction
+
+    def __post_init__(self, **kwargs):
+        self.dtype = torch.float32
+        self.in_features = self.rbln_config.meta.get("in_features", None)
+        if self.in_features is not None:
+
+            @dataclass
+            class LINEAR1:
+                in_features: int
+
+            @dataclass
+            class ADDEMBEDDING:
+                linear_1: LINEAR1
+
+            self.add_embedding = ADDEMBEDDING(LINEAR1(self.in_features))
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        def get_model_from_task(
+            task: str,
+            model_name_or_path: Union[str, Path],
+            **kwargs,
+        ):
+            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
+        return rt
+
+    @classmethod
+    def wrap_model_if_needed(cls, model: torch.nn.Module) -> torch.nn.Module:
+        if model.config.addition_embed_type == "text_time":
+            return _UNet_SDXL(model).eval()
+        else:
+            return _UNet_SD(model).eval()
+
+    @classmethod
+    def _get_rbln_config(
+        cls,
+        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_in_features: Optional[int] = None,
+        rbln_use_encode: Optional[bool] = None,
+        rbln_img_width: Optional[int] = None,
+        rbln_img_height: Optional[int] = None,
+        rbln_vae_scale_factor: Optional[int] = None,
+        rbln_is_controlnet: Optional[bool] = None,
+    ) -> RBLNConfig:
+        meta = {"type": "unet"}
+        if rbln_batch_size is None:
+            rbln_batch_size = 1
+
+        if rbln_max_seq_len is None:
+            rbln_max_seq_len = 77
+
+        meta["rbln_use_encode"] = rbln_use_encode
+
+        if rbln_use_encode:
+            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
+
+        input_info = [
+            (
+                "sample",
+                [
+                    rbln_batch_size,
+                    model_config.in_channels,
+                    input_width,
+                    input_height,
+                ],
+                "float32",
+            ),
+            ("timestep", [], "float32"),
+            (
+                "encoder_hidden_states",
+                [
+                    rbln_batch_size,
+                    rbln_max_seq_len,
+                    model_config.cross_attention_dim,
+                ],
+                "float32",
+            ),
+        ]
+        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",
+                )
+            )
+            input_info.extend(
+                [
+                    (
+                        f"down_block_additional_residuals_{i}",
+                        [rbln_batch_size, model_config.block_out_channels[1], input_width // 2, input_height // 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(
+                [
+                    (
+                        f"down_block_additional_residuals_{i}",
+                        [rbln_batch_size, model_config.block_out_channels[3], input_width // 8, input_height // 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],
+                    "float32",
+                )
+            )
+
+        rbln_runtime_config = RBLNRuntimeConfig(
+            input_info=input_info,
+            batch_size=rbln_batch_size,
+        )
+
+        if hasattr(model_config, "addition_embed_type") and model_config.addition_embed_type == "text_time":
+            # In case of sdxl
+            if rbln_text_model_hidden_size is None:
+                rbln_text_model_hidden_size = 768
+            if rbln_in_features is None:
+                rbln_in_features = 2816
+            meta["in_features"] = rbln_in_features
+            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
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs: Dict[str, torch.Tensor] = {},
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        """
+        arg order : latent_model_input, t, prompt_embeds
+        """
+        added_cond_kwargs = {} if added_cond_kwargs is None else added_cond_kwargs
+
+        if down_block_additional_residuals is not None:
+            down_block_additional_residuals = [t.contiguous() for t in down_block_additional_residuals]
+            return (
+                super().forward(
+                    sample.contiguous(),
+                    timestep.float(),
+                    encoder_hidden_states,
+                    *down_block_additional_residuals,
+                    mid_block_additional_residual,
+                    **added_cond_kwargs,
+                ),
+            )
+
+        return (
+            super().forward(
+                sample,
+                timestep.float(),
+                encoder_hidden_states,
+                **added_cond_kwargs,
+            ),
+        )

optimum/rbln/diffusers/pipelines/__init__.py

@@ -0,0 +1,30 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .controlnet import RBLNMultiControlNetModel
+from .stable_diffusion import (
+    RBLNStableDiffusionControlNetImg2ImgPipeline,
+    RBLNStableDiffusionImg2ImgPipeline,
+    RBLNStableDiffusionPipeline,
+)
+from .stable_diffusion_xl import RBLNStableDiffusionXLImg2ImgPipeline, RBLNStableDiffusionXLPipeline

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

@@ -0,0 +1,24 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from .multicontrolnet import RBLNMultiControlNetModel