From 15adfa2f64d2a15e8a23c6231bb1d46d072668eb Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Wed, 30 Aug 2023 10:13:41 +0800 Subject: [PATCH 01/13] version to 1.9.0rc0 --- modelscope/version.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modelscope/version.py b/modelscope/version.py index bb262e90..0a98fac3 100644 --- a/modelscope/version.py +++ b/modelscope/version.py @@ -1,5 +1,5 @@ # Make sure to modify __release_datetime__ to release time when making official release. -__version__ = '1.8.1' +__version__ = '1.9.0rc0' # default release datetime for branches under active development is set # to be a time far-far-away-into-the-future -__release_datetime__ = '2099-10-13 08:56:12' +__release_datetime__ = '2023-09-03 00:00:00' From 3fcc6e86bf75a74b56d6c8933fdce31d49188e5e Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Wed, 30 Aug 2023 15:04:43 +0800 Subject: [PATCH 02/13] remove cv diffusers version Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13843028 * remove cv diffusers version --- requirements/cv.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/cv.txt b/requirements/cv.txt index 0b89dece..ee9f5582 100644 --- a/requirements/cv.txt +++ b/requirements/cv.txt @@ -6,7 +6,7 @@ chumpy clip>=1.0 control_ldm ddpm_guided_diffusion -diffusers==0.18.0 +diffusers easydict easyrobust edit_distance From df3d781427ca76bfed9e9200f827fa6dc4ba3f7d Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Wed, 30 Aug 2023 15:33:54 +0800 Subject: [PATCH 03/13] pick support a10 video to video model --- .../multi_modal/video_to_video/video_to_video_model.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/modelscope/models/multi_modal/video_to_video/video_to_video_model.py b/modelscope/models/multi_modal/video_to_video/video_to_video_model.py index bd485da5..1a13a3fc 100755 --- a/modelscope/models/multi_modal/video_to_video/video_to_video_model.py +++ b/modelscope/models/multi_modal/video_to_video/video_to_video_model.py @@ -112,7 +112,7 @@ class VideoToVideo(TorchModel): generator.eval() load_dict = torch.load(cfg.model_path, map_location='cpu') ret = generator.load_state_dict(load_dict['state_dict'], strict=True) - self.generator = generator + self.generator = generator.half() logger.info('Load model {} path {}, with local status {}'.format( cfg.UNet.type, cfg.model_path, ret)) @@ -175,7 +175,7 @@ class VideoToVideo(TorchModel): video_data = rearrange(video_data, 'b f c h w -> (b f) c h w') video_data_list = torch.chunk( - video_data, video_data.shape[0] // 2, dim=0) + video_data, video_data.shape[0] // 1, dim=0) with torch.no_grad(): decode_data = [] for vd_data in video_data_list: @@ -185,6 +185,7 @@ class VideoToVideo(TorchModel): video_data_feature = torch.cat(decode_data, dim=0) video_data_feature = rearrange( video_data_feature, '(b f) c h w -> b c f h w', b=batch_size) + torch.cuda.empty_cache() with amp.autocast(enabled=True): total_noise_levels = 600 @@ -209,6 +210,7 @@ class VideoToVideo(TorchModel): t_min=0, discretization='trailing') + torch.cuda.empty_cache() scale_factor = 0.18215 vid_tensor_feature = 1. / scale_factor * gen_vid From 239672c043e07f4c75334d0b6e71270bf392ffa4 Mon Sep 17 00:00:00 2001 From: "wenmeng.zwm" Date: Wed, 30 Aug 2023 14:58:32 +0800 Subject: [PATCH 04/13] add python311 support for whl --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 98b12888..dbac6e77 100644 --- a/setup.py +++ b/setup.py @@ -219,6 +219,7 @@ if __name__ == '__main__': 'Programming Language :: Python :: 3.8', 'Programming Language :: Python :: 3.9', 'Programming Language :: Python :: 3.10', + 'Programming Language :: Python :: 3.11', ], license='Apache License 2.0', tests_require=parse_requirements('requirements/tests.txt'), From 4449de9d0346f51a1a28556c8d4fd3b63df55776 Mon Sep 17 00:00:00 2001 From: "baiguan.yt" Date: Thu, 7 Sep 2023 16:53:54 +0800 Subject: [PATCH 05/13] add PASD image super-resolution for facechain and more Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13924655 * add parameters height and width * update pipeline --- modelscope/metainfo.py | 1 + modelscope/models/cv/__init__.py | 9 +- .../image_super_resolution_pasd/__init__.py | 24 + .../image_super_resolution_pasd/attention.py | 421 +++ .../image_super_resolution_pasd/controlnet.py | 738 ++++ .../cv/image_super_resolution_pasd/misc.py | 159 + .../transformer_2d.py | 365 ++ .../unet_2d_blocks.py | 2983 +++++++++++++++++ .../unet_2d_condition.py | 887 +++++ modelscope/outputs/outputs.py | 1 + modelscope/pipeline_inputs.py | 4 + modelscope/pipelines/cv/__init__.py | 8 +- .../image_super_resolution_pasd_pipeline.py | 229 ++ .../multi_modal/diffusers_wrapped/devices.py | 126 + .../diffusers_wrapped/pasd_pipeline.py | 1131 +++++++ .../multi_modal/diffusers_wrapped/vaehook.py | 762 +++++ modelscope/utils/constant.py | 1 + .../test_image_super_resolution_pasd.py | 43 + tests/run_config.yaml | 1 + 19 files changed, 7887 insertions(+), 6 deletions(-) create mode 100644 modelscope/models/cv/image_super_resolution_pasd/__init__.py create mode 100644 modelscope/models/cv/image_super_resolution_pasd/attention.py create mode 100644 modelscope/models/cv/image_super_resolution_pasd/controlnet.py create mode 100644 modelscope/models/cv/image_super_resolution_pasd/misc.py create mode 100644 modelscope/models/cv/image_super_resolution_pasd/transformer_2d.py create mode 100644 modelscope/models/cv/image_super_resolution_pasd/unet_2d_blocks.py create mode 100644 modelscope/models/cv/image_super_resolution_pasd/unet_2d_condition.py create mode 100644 modelscope/pipelines/cv/image_super_resolution_pasd_pipeline.py create mode 100644 modelscope/pipelines/multi_modal/diffusers_wrapped/devices.py create mode 100644 modelscope/pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py create mode 100644 modelscope/pipelines/multi_modal/diffusers_wrapped/vaehook.py create mode 100644 tests/pipelines/test_image_super_resolution_pasd.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index d50c2b52..b5ec251a 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -339,6 +339,7 @@ class Pipelines(object): image_colorization = 'unet-image-colorization' image_style_transfer = 'AAMS-style-transfer' image_super_resolution = 'rrdb-image-super-resolution' + image_super_resolution_pasd = 'image-super-resolution-pasd' image_debanding = 'rrdb-image-debanding' face_image_generation = 'gan-face-image-generation' product_retrieval_embedding = 'resnet50-product-retrieval-embedding' diff --git a/modelscope/models/cv/__init__.py b/modelscope/models/cv/__init__.py index d3d946f5..2417f33b 100644 --- a/modelscope/models/cv/__init__.py +++ b/modelscope/models/cv/__init__.py @@ -14,10 +14,11 @@ from . import (action_recognition, animal_recognition, bad_image_detecting, image_quality_assessment_degradation, image_quality_assessment_man, image_quality_assessment_mos, image_reid_person, image_restoration, - image_semantic_segmentation, image_to_image_generation, - image_to_image_translation, language_guided_video_summarization, - movie_scene_segmentation, object_detection, - panorama_depth_estimation, pedestrian_attribute_recognition, + image_semantic_segmentation, image_super_resolution_pasd, + image_to_image_generation, image_to_image_translation, + language_guided_video_summarization, movie_scene_segmentation, + object_detection, panorama_depth_estimation, + pedestrian_attribute_recognition, pointcloud_sceneflow_estimation, product_retrieval_embedding, referring_video_object_segmentation, robust_image_classification, salient_detection, diff --git a/modelscope/models/cv/image_super_resolution_pasd/__init__.py b/modelscope/models/cv/image_super_resolution_pasd/__init__.py new file mode 100644 index 00000000..1448c348 --- /dev/null +++ b/modelscope/models/cv/image_super_resolution_pasd/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .unet_2d_condition import UNet2DConditionModel + from .controlnet import ControlNetModel + +else: + _import_structure = { + 'unet_2d_condition': ['UNet2DConditionModel'], + 'controlnet': ['ControlNetModel'] + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/cv/image_super_resolution_pasd/attention.py b/modelscope/models/cv/image_super_resolution_pasd/attention.py new file mode 100644 index 00000000..1ec566de --- /dev/null +++ b/modelscope/models/cv/image_super_resolution_pasd/attention.py @@ -0,0 +1,421 @@ +# Part of the implementation is borrowed and modified from diffusers, +# publicly available at https://github.com/huggingface/diffusers/tree/main/src/diffusers/models/attention.py +import math +from typing import Callable, Optional + +import torch +import torch.nn.functional as F +from diffusers.models.attention_processor import Attention +from diffusers.models.embeddings import CombinedTimestepLabelEmbeddings +from diffusers.utils import maybe_allow_in_graph +from torch import nn + + +@maybe_allow_in_graph +class BasicTransformerBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + pixelwise_cross_attention_dim: Optional[int] = None, + activation_fn: str = 'geglu', + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = 'layer_norm', + final_dropout: bool = False, + use_pixelwise_attention: bool = False, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + + self.use_ada_layer_norm_zero = ( + num_embeds_ada_norm is not None) and norm_type == 'ada_norm_zero' + self.use_ada_layer_norm = (num_embeds_ada_norm + is not None) and norm_type == 'ada_norm' + + if norm_type in ('ada_norm', + 'ada_norm_zero') and num_embeds_ada_norm is None: + raise ValueError( + f'`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to' + f' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.' + ) + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if self.use_ada_layer_norm: + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_zero: + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + else: + self.norm1 = nn.LayerNorm( + dim, elementwise_affine=norm_elementwise_affine) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim + if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm else nn.LayerNorm( + dim, elementwise_affine=norm_elementwise_affine)) + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim + if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2 = None + self.attn2 = None + + # 2+. pixelwise-Attn + if use_pixelwise_attention: + self.norm2_plus = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm else nn.LayerNorm( + dim, elementwise_affine=norm_elementwise_affine)) + self.attn2_plus = Attention( + query_dim=dim, + cross_attention_dim=pixelwise_cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2_plus = None + self.attn2_plus = None + + # 3. Feed-forward + self.norm3 = nn.LayerNorm( + dim, elementwise_affine=norm_elementwise_affine) + self.ff = FeedForward( + dim, + dropout=dropout, + activation_fn=activation_fn, + final_dropout=final_dropout) + + def forward( + self, + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + encoder_pixelwise_hidden_states=None, + encoder_pixelwise_attention_mask=None, + timestep=None, + cross_attention_kwargs=None, + class_labels=None, + ): + # Notice that normalization is always applied before the real computation in the following blocks. + # 1. Self-Attention + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, + timestep, + class_labels, + hidden_dtype=hidden_states.dtype) + else: + norm_hidden_states = self.norm1(hidden_states) + + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states + if self.only_cross_attention else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = attn_output + hidden_states + + # 2. Cross-Attention + if self.attn2 is not None: + norm_hidden_states = ( + self.norm2(hidden_states, timestep) + if self.use_ada_layer_norm else self.norm2(hidden_states)) + # TODO (Birch-San): Here we should prepare the encoder_attention mask correctly + # prepare attention mask here + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 2+. pixelwise-Attention + if self.attn2_plus is not None: + norm_hidden_states = ( + self.norm2_plus(hidden_states, timestep) + if self.use_ada_layer_norm else self.norm2_plus(hidden_states)) + + attn_output = self.attn2_plus( + norm_hidden_states, + encoder_hidden_states=encoder_pixelwise_hidden_states, + attention_mask=encoder_pixelwise_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 3. Feed-forward + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = norm_hidden_states * ( + 1 + scale_mlp[:, None]) + shift_mlp[:, None] + + ff_output = self.ff(norm_hidden_states) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = ff_output + hidden_states + + return hidden_states + + +class FeedForward(nn.Module): + r""" + A feed-forward layer. + + Parameters: + dim (`int`): The number of channels in the input. + dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. + mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. + """ + + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + dropout: float = 0.0, + activation_fn: str = 'geglu', + final_dropout: bool = False, + ): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + if activation_fn == 'gelu': + act_fn = GELU(dim, inner_dim) + if activation_fn == 'gelu-approximate': + act_fn = GELU(dim, inner_dim, approximate='tanh') + elif activation_fn == 'geglu': + act_fn = GEGLU(dim, inner_dim) + elif activation_fn == 'geglu-approximate': + act_fn = ApproximateGELU(dim, inner_dim) + + self.net = nn.ModuleList([]) + # project in + self.net.append(act_fn) + # project dropout + self.net.append(nn.Dropout(dropout)) + # project out + self.net.append(nn.Linear(inner_dim, dim_out)) + # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout + if final_dropout: + self.net.append(nn.Dropout(dropout)) + + def forward(self, hidden_states): + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +class GELU(nn.Module): + r""" + GELU activation function with tanh approximation support with `approximate="tanh"`. + """ + + def __init__(self, dim_in: int, dim_out: int, approximate: str = 'none'): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out) + self.approximate = approximate + + def gelu(self, gate): + if gate.device.type != 'mps': + return F.gelu(gate, approximate=self.approximate) + # mps: gelu is not implemented for float16 + return F.gelu( + gate.to(dtype=torch.float32), + approximate=self.approximate).to(dtype=gate.dtype) + + def forward(self, hidden_states): + hidden_states = self.proj(hidden_states) + hidden_states = self.gelu(hidden_states) + return hidden_states + + +class GEGLU(nn.Module): + r""" + A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202. + + Parameters: + dim_in (`int`): The number of channels in the input. + dim_out (`int`): The number of channels in the output. + """ + + def __init__(self, dim_in: int, dim_out: int): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2) + + def gelu(self, gate): + if gate.device.type != 'mps': + return F.gelu(gate) + # mps: gelu is not implemented for float16 + return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype) + + def forward(self, hidden_states): + hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1) + return hidden_states * self.gelu(gate) + + +class ApproximateGELU(nn.Module): + """ + The approximate form of Gaussian Error Linear Unit (GELU) + + For more details, see section 2: https://arxiv.org/abs/1606.08415 + """ + + def __init__(self, dim_in: int, dim_out: int): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out) + + def forward(self, x): + x = self.proj(x) + return x * torch.sigmoid(1.702 * x) + + +class AdaLayerNorm(nn.Module): + """ + Norm layer modified to incorporate timestep embeddings. + """ + + def __init__(self, embedding_dim, num_embeddings): + super().__init__() + self.emb = nn.Embedding(num_embeddings, embedding_dim) + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, embedding_dim * 2) + self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False) + + def forward(self, x, timestep): + emb = self.linear(self.silu(self.emb(timestep))) + scale, shift = torch.chunk(emb, 2) + x = self.norm(x) * (1 + scale) + shift + return x + + +class AdaLayerNormZero(nn.Module): + """ + Norm layer adaptive layer norm zero (adaLN-Zero). + """ + + def __init__(self, embedding_dim, num_embeddings): + super().__init__() + + self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, + embedding_dim) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + self.norm = nn.LayerNorm( + embedding_dim, elementwise_affine=False, eps=1e-6) + + def forward(self, x, timestep, class_labels, hidden_dtype=None): + emb = self.linear( + self.silu( + self.emb(timestep, class_labels, hidden_dtype=hidden_dtype))) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk( + 6, dim=1) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class AdaGroupNorm(nn.Module): + """ + GroupNorm layer modified to incorporate timestep embeddings. + """ + + def __init__(self, + embedding_dim: int, + out_dim: int, + num_groups: int, + act_fn: Optional[str] = None, + eps: float = 1e-5): + super().__init__() + self.num_groups = num_groups + self.eps = eps + self.act = None + if act_fn == 'swish': + self.act = lambda x: F.silu(x) + elif act_fn == 'mish': + self.act = nn.Mish() + elif act_fn == 'silu': + self.act = nn.SiLU() + elif act_fn == 'gelu': + self.act = nn.GELU() + + self.linear = nn.Linear(embedding_dim, out_dim * 2) + + def forward(self, x, emb): + if self.act: + emb = self.act(emb) + emb = self.linear(emb) + emb = emb[:, :, None, None] + scale, shift = emb.chunk(2, dim=1) + + x = F.group_norm(x, self.num_groups, eps=self.eps) + x = x * (1 + scale) + shift + return x diff --git a/modelscope/models/cv/image_super_resolution_pasd/controlnet.py b/modelscope/models/cv/image_super_resolution_pasd/controlnet.py new file mode 100644 index 00000000..f2224a1b --- /dev/null +++ b/modelscope/models/cv/image_super_resolution_pasd/controlnet.py @@ -0,0 +1,738 @@ +# Part of the implementation is borrowed and modified from diffusers, +# publicly available at https://github.com/huggingface/diffusers/tree/main/src/diffusers/models/controlnet.py +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models import ModelMixin # , ControlNetModel +from diffusers.models.attention_processor import (AttentionProcessor, + AttnProcessor) +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from diffusers.utils import BaseOutput, logging +from torch import nn +from torch.nn import functional as F +from torchvision import utils + +from modelscope.models.cv.super_resolution.rrdbnet_arch import RRDB +from .unet_2d_blocks import (CrossAttnDownBlock2D, DownBlock2D, + UNetMidBlock2DCrossAttn, get_down_block) +from .unet_2d_condition import UNet2DConditionModel + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class ControlNetOutput(BaseOutput): + controlnet_cond_mid: torch.Tensor + down_block_res_samples: Tuple[torch.Tensor] + mid_block_res_sample: torch.Tensor + + +class ControlNetConditioningEmbedding(nn.Module): + """ + Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN + [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized + training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the + convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides + (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full + model) to encode image-space conditions ... into feature maps ..." + """ + + def __init__( + self, + conditioning_embedding_channels: int, + conditioning_channels: int = 3, + block_out_channels: Tuple[int] = (16, 32, 96, 256), + return_rgbs: bool = True, + use_rrdb: bool = False, + ): + super().__init__() + + self.return_rgbs = return_rgbs + self.use_rrdb = use_rrdb + + self.conv_in = nn.Conv2d( + conditioning_channels, + block_out_channels[0], + kernel_size=3, + padding=1) + + if self.use_rrdb: + num_rrdb_block = 2 + layers = ( + RRDB(block_out_channels[0], block_out_channels[0]) + for i in range(num_rrdb_block)) + self.preprocesser = nn.Sequential(*layers) + + self.blocks = nn.ModuleList([]) + if return_rgbs: + self.to_rgbs = nn.ModuleList([]) + + for i in range(len(block_out_channels) - 1): + channel_in = block_out_channels[i] + channel_out = block_out_channels[i + 1] + self.blocks.append( + nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)) + self.blocks.append( + nn.Conv2d( + channel_in, + channel_out, + kernel_size=3, + padding=1, + stride=2)) + + if return_rgbs: + self.to_rgbs.append( + nn.Conv2d(channel_out, 3, kernel_size=3, padding=1)) + + self.conv_out = zero_module( + nn.Conv2d( + block_out_channels[-1], + conditioning_embedding_channels, + kernel_size=3, + padding=1)) + + def forward(self, conditioning): + embedding = self.conv_in(conditioning) + embedding = F.silu(embedding) + + if self.use_rrdb: + embedding = self.preprocesser(embedding) + + out_rgbs = [] + for i, block in enumerate(self.blocks): + embedding = block(embedding) + embedding = F.silu(embedding) + + if i % 2 and self.return_rgbs: # 0 + out_rgbs.append(self.to_rgbs[i // 2](embedding)) + + embedding = self.conv_out(embedding) + + if self.return_rgbs: + return embedding, out_rgbs + else: + return embedding + + +class ControlNetModel(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 4, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + 'CrossAttnDownBlock2D', + 'CrossAttnDownBlock2D', + 'CrossAttnDownBlock2D', + 'DownBlock2D', + ), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: int = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + act_fn: str = 'silu', + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: int = 1280, + attention_head_dim: Union[int, Tuple[int]] = 8, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = 'default', + projection_class_embeddings_input_dim: Optional[int] = None, + controlnet_conditioning_channel_order: str = 'rgb', + conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, + 32, + 96, + 256), + global_pool_conditions: bool = False, + return_rgbs: bool = False, + use_rrdb: bool = False): + super().__init__() + + # Check inputs + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f'Must provide the same number of `block_out_channels` as `down_block_types`. \ + `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}.' + ) + + if not isinstance( + only_cross_attention, + bool) and len(only_cross_attention) != len(down_block_types): + raise ValueError( + f'Must provide the same number of `only_cross_attention` as `down_block_types`. \ + `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}.' + ) + + if not isinstance( + attention_head_dim, + int) and len(attention_head_dim) != len(down_block_types): + raise ValueError( + f'Must provide the same number of `attention_head_dim` as `down_block_types`. \ + `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}.' + ) + + # input + self.return_rgbs = return_rgbs + conv_in_kernel = 3 + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d( + in_channels, + block_out_channels[0], + kernel_size=conv_in_kernel, + padding=conv_in_padding) + + # time + time_embed_dim = block_out_channels[0] * 4 + + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, + freq_shift) + timestep_input_dim = block_out_channels[0] + + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + ) + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, + time_embed_dim) + elif class_embed_type == 'timestep': + self.class_embedding = TimestepEmbedding(timestep_input_dim, + time_embed_dim) + elif class_embed_type == 'identity': + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + elif class_embed_type == 'projection': + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" + ) + # The projection `class_embed_type` is the same as the timestep `class_embed_type` except + # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings + # 2. it projects from an arbitrary input dimension. + # + # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. + # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. + # As a result, `TimestepEmbedding` can be passed arbitrary vectors. + self.class_embedding = TimestepEmbedding( + projection_class_embeddings_input_dim, time_embed_dim) + else: + self.class_embedding = None + + # control net conditioning embedding + self.controlnet_cond_embedding = ControlNetConditioningEmbedding( + conditioning_embedding_channels=block_out_channels[0], + block_out_channels=conditioning_embedding_out_channels, + return_rgbs=return_rgbs, + use_rrdb=use_rrdb, + ) + + self.down_blocks = nn.ModuleList([]) + self.controlnet_down_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + only_cross_attention = [only_cross_attention + ] * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim, ) * len(down_block_types) + + # down + output_channel = block_out_channels[0] + + controlnet_block = nn.Conv2d( + output_channel, output_channel, kernel_size=1) + controlnet_block = zero_module(controlnet_block) + self.controlnet_down_blocks.append(controlnet_block) + + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + temb_channels=time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attention_head_dim[i], + downsample_padding=downsample_padding, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + self.down_blocks.append(down_block) + + for _ in range(layers_per_block): + controlnet_block = nn.Conv2d( + output_channel, output_channel, kernel_size=1) + controlnet_block = zero_module(controlnet_block) + self.controlnet_down_blocks.append(controlnet_block) + + if not is_final_block: + controlnet_block = nn.Conv2d( + output_channel, output_channel, kernel_size=1) + controlnet_block = zero_module(controlnet_block) + self.controlnet_down_blocks.append(controlnet_block) + + # mid + mid_block_channel = block_out_channels[-1] + + controlnet_block = nn.Conv2d( + mid_block_channel, mid_block_channel, kernel_size=1) + controlnet_block = zero_module(controlnet_block) + self.controlnet_mid_block = controlnet_block + + self.mid_block = UNetMidBlock2DCrossAttn( + in_channels=mid_block_channel, + temb_channels=time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attention_head_dim[-1], + resnet_groups=norm_num_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + ) + + @classmethod + def from_unet( + cls, + unet: UNet2DConditionModel, + controlnet_conditioning_channel_order: str = 'rgb', + conditioning_embedding_out_channels: Optional[Tuple[int]] = (16, 32, + 96, 256), + load_weights_from_unet: bool = True, + ): + r""" + Instantiate Controlnet class from UNet2DConditionModel. + + Parameters: + unet (`UNet2DConditionModel`): + UNet model which weights are copied to the ControlNet. Note that all configuration options are also + copied where applicable. + """ + controlnet = cls( + in_channels=unet.config.in_channels, + flip_sin_to_cos=unet.config.flip_sin_to_cos, + freq_shift=unet.config.freq_shift, + down_block_types=unet.config.down_block_types, + only_cross_attention=unet.config.only_cross_attention, + block_out_channels=unet.config.block_out_channels, + layers_per_block=unet.config.layers_per_block, + downsample_padding=unet.config.downsample_padding, + mid_block_scale_factor=unet.config.mid_block_scale_factor, + act_fn=unet.config.act_fn, + norm_num_groups=unet.config.norm_num_groups, + norm_eps=unet.config.norm_eps, + cross_attention_dim=unet.config.cross_attention_dim, + attention_head_dim=unet.config.attention_head_dim, + use_linear_projection=unet.config.use_linear_projection, + class_embed_type=unet.config.class_embed_type, + num_class_embeds=unet.config.num_class_embeds, + upcast_attention=unet.config.upcast_attention, + resnet_time_scale_shift=unet.config.resnet_time_scale_shift, + projection_class_embeddings_input_dim=unet.config. + projection_class_embeddings_input_dim, + controlnet_conditioning_channel_order= + controlnet_conditioning_channel_order, + conditioning_embedding_out_channels= + conditioning_embedding_out_channels, + ) + + if load_weights_from_unet: + controlnet.conv_in.load_state_dict(unet.conv_in.state_dict()) + controlnet.time_proj.load_state_dict(unet.time_proj.state_dict()) + controlnet.time_embedding.load_state_dict( + unet.time_embedding.state_dict()) + + if controlnet.class_embedding: + controlnet.class_embedding.load_state_dict( + unet.class_embedding.state_dict()) + + controlnet.down_blocks.load_state_dict( + unet.down_blocks.state_dict()) + controlnet.mid_block.load_state_dict(unet.mid_block.state_dict()) + + if controlnet.sr_model is not None: + load_net = torch.load( + 'annotator/ckpts/RealESRNet_x4plus.pth', + map_location=lambda storage, loc: storage) + if 'params_ema' in load_net: + load_net = load_net['params_ema'] + elif 'params' in load_net: + load_net = load_net['params'] + # remove unnecessary 'module.' + for k, v in deepcopy(load_net).items(): + if k.startswith('module.'): + load_net[k[7:]] = v + load_net.pop(k) + controlnet.sr_model.load_state_dict(load_net, strict=True) + + return controlnet + + @property + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, + processors: Dict[str, + AttentionProcessor]): + if hasattr(module, 'set_processor'): + processors[f'{name}.processor'] = module.processor + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f'{name}.{sub_name}', child, + processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor(self, processor: Union[AttentionProcessor, + Dict[str, + AttentionProcessor]]): + r""" + Parameters: + `processor (`dict` of `AttentionProcessor` or `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + of **all** `Attention` layers. + In case `processor` is a dict, the key needs to define the path to + the corresponding cross attention processor. + This is strongly recommended when setting trainable attention processors.: + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f'A dict of processors was passed, but the number of processors {len(processor)} does not match the' + f' number of attention layers: {count}. Please make sure to pass {count} processor classes.' + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, + processor): + if hasattr(module, 'set_processor'): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f'{name}.processor')) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f'{name}.{sub_name}', child, + processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + self.set_attn_processor(AttnProcessor()) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attention_slice + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module will split the input tensor in slices, to compute attention + in several steps. This is useful to save some memory in exchange for a small speed decrease. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If + `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): + if hasattr(module, 'set_attention_slice'): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_sliceable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_sliceable_dims(module) + + num_sliceable_layers = len(sliceable_head_dims) + + if slice_size == 'auto': + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == 'max': + # make smallest slice possible + slice_size = num_sliceable_layers * [1] + + slice_size = num_sliceable_layers * [slice_size] if not isinstance( + slice_size, list) else slice_size + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f'You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different' + f' attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}.' + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError( + f'size {size} has to be smaller or equal to {dim}.') + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice(module: torch.nn.Module, + slice_size: List[int]): + if hasattr(module, 'set_attention_slice'): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D)): + module.gradient_checkpointing = value + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + controlnet_cond: torch.FloatTensor, + fg_mask: Optional[torch.FloatTensor] = None, + conditioning_scale_fg: float = 1.0, + conditioning_scale_bg: float = 1.0, + 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, + guess_mode: bool = False, + return_dict: bool = True, + ) -> Union[ControlNetOutput, Tuple]: + # check channel order + channel_order = self.config.controlnet_conditioning_channel_order + + if channel_order == 'rgb': + # in rgb order by default + ... + elif channel_order == 'bgr': + controlnet_cond = torch.flip(controlnet_cond, dims=[1]) + else: + raise ValueError( + f'unknown `controlnet_conditioning_channel_order`: {channel_order}' + ) + + # prepare attention_mask + if attention_mask is not None: + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == 'mps' + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], + dtype=dtype, + device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # timesteps does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb, timestep_cond) + + if self.class_embedding is not None: + if class_labels is None: + raise ValueError( + 'class_labels should be provided when num_class_embeds > 0' + ) + + if self.config.class_embed_type == 'timestep': + class_labels = self.time_proj(class_labels) + + class_emb = self.class_embedding(class_labels).to(dtype=self.dtype) + emb = emb + class_emb + + # 2. pre-process + sample = self.conv_in(sample) + + controlnet_cond_mid = None + if self.return_rgbs: + controlnet_cond, controlnet_cond_mid = self.controlnet_cond_embedding( + controlnet_cond) + else: + controlnet_cond = self.controlnet_cond_embedding(controlnet_cond) + + sample = sample + controlnet_cond + + # 3. down + down_block_res_samples = (sample, ) + for downsample_block in self.down_blocks: + if hasattr(downsample_block, 'has_cross_attention' + ) and downsample_block.has_cross_attention: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + else: + sample, res_samples = downsample_block( + hidden_states=sample, temb=emb) + + down_block_res_samples += res_samples + + # 4. mid + if self.mid_block is not None: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + + # 5. Control net blocks + + controlnet_down_block_res_samples = () + + for down_block_res_sample, controlnet_block in zip( + down_block_res_samples, self.controlnet_down_blocks): + down_block_res_sample = controlnet_block(down_block_res_sample) + controlnet_down_block_res_samples = controlnet_down_block_res_samples + ( + down_block_res_sample, ) + + down_block_res_samples = controlnet_down_block_res_samples + + mid_block_res_sample = self.controlnet_mid_block(sample) + + # 6. scaling + if guess_mode and not self.config.global_pool_conditions: + scales = torch.logspace( + -1, 0, len(down_block_res_samples) + 1, + device=sample.device) # 0.1 to 1.0 + + scales = scales * conditioning_scale_fg + down_block_res_samples = [ + sample * scale + for sample, scale in zip(down_block_res_samples, scales) + ] + mid_block_res_sample = mid_block_res_sample * scales[ + -1] # last one + else: + if fg_mask is None: + down_block_res_samples = [ + sample * conditioning_scale_fg + for sample in down_block_res_samples + ] + mid_block_res_sample = mid_block_res_sample * conditioning_scale_fg + else: + down_block_masks = [ + torch.zeros_like(sample) + conditioning_scale_bg + for i, sample in enumerate(down_block_res_samples) + ] + mid_block_mask = torch.zeros_like( + mid_block_res_sample) + conditioning_scale_bg + + for i, sample in enumerate(down_block_masks): + tmp_mask = F.interpolate( + fg_mask, + size=sample.shape[-2:]).repeat(sample.shape[0], + sample.shape[1], 1, + 1).bool() + down_block_masks[i] = sample.masked_fill( + tmp_mask, conditioning_scale_fg) + + tmp_mask = F.interpolate( + fg_mask, size=mid_block_mask.shape[-2:]).repeat( + mid_block_mask.shape[0], mid_block_mask.shape[1], 1, + 1).bool() + mid_block_mask = mid_block_mask.masked_fill( + tmp_mask, conditioning_scale_fg) + + down_block_res_samples = [ + sample * down_block_mask for sample, down_block_mask in + zip(down_block_res_samples, down_block_masks) + ] + mid_block_res_sample = mid_block_res_sample * mid_block_mask + + if self.config.global_pool_conditions: + down_block_res_samples = [ + torch.mean(sample, dim=(2, 3), keepdim=True) + for sample in down_block_res_samples + ] + mid_block_res_sample = torch.mean( + mid_block_res_sample, dim=(2, 3), keepdim=True) + + if not return_dict: + return (controlnet_cond_mid, down_block_res_samples, + mid_block_res_sample) + + return ControlNetOutput( + controlnet_cond_mid=controlnet_cond_mid, + down_block_res_samples=down_block_res_samples, + mid_block_res_sample=mid_block_res_sample) + + +def zero_module(module): + for p in module.parameters(): + nn.init.zeros_(p) + return module diff --git a/modelscope/models/cv/image_super_resolution_pasd/misc.py b/modelscope/models/cv/image_super_resolution_pasd/misc.py new file mode 100644 index 00000000..04869064 --- /dev/null +++ b/modelscope/models/cv/image_super_resolution_pasd/misc.py @@ -0,0 +1,159 @@ +# Part of the implementation is borrowed and modified from StableSR, +# publicly available at https://github.com/IceClear/StableSR/blob/main/scripts/wavelet_color_fix.py +import torch +from diffusers.pipelines.stable_diffusion.convert_from_ckpt import ( + convert_ldm_unet_checkpoint, convert_ldm_vae_checkpoint) +from PIL import Image +from safetensors import safe_open +from torch import Tensor +from torch.nn import functional as F +from torchvision.transforms import ToPILImage, ToTensor + + +def adain_color_fix(target: Image, source: Image): + # Convert images to tensors + to_tensor = ToTensor() + target_tensor = to_tensor(target).unsqueeze(0) + source_tensor = to_tensor(source).unsqueeze(0) + + # Apply adaptive instance normalization + result_tensor = adaptive_instance_normalization(target_tensor, + source_tensor) + + # Convert tensor back to image + to_image = ToPILImage() + result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0)) + + return result_image + + +def wavelet_color_fix(target: Image, source: Image): + # Convert images to tensors + to_tensor = ToTensor() + target_tensor = to_tensor(target).unsqueeze(0) + source_tensor = to_tensor(source).unsqueeze(0) + + # Apply wavelet reconstruction + result_tensor = wavelet_reconstruction(target_tensor, source_tensor) + + # Convert tensor back to image + to_image = ToPILImage() + result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0)) + + return result_image + + +def calc_mean_std(feat: Tensor, eps=1e-5): + """Calculate mean and std for adaptive_instance_normalization. + Args: + feat (Tensor): 4D tensor. + eps (float): A small value added to the variance to avoid + divide-by-zero. Default: 1e-5. + """ + size = feat.size() + assert len(size) == 4, 'The input feature should be 4D tensor.' + b, c = size[:2] + feat_var = feat.reshape(b, c, -1).var(dim=2) + eps + feat_std = feat_var.sqrt().reshape(b, c, 1, 1) + feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1) + return feat_mean, feat_std + + +def adaptive_instance_normalization(content_feat: Tensor, style_feat: Tensor): + """Adaptive instance normalization. + Adjust the reference features to have the similar color and illuminations + as those in the degradate features. + Args: + content_feat (Tensor): The reference feature. + style_feat (Tensor): The degradate features. + """ + size = content_feat.size() + style_mean, style_std = calc_mean_std(style_feat) + content_mean, content_std = calc_mean_std(content_feat) + normalized_feat = (content_feat + - content_mean.expand(size)) / content_std.expand(size) + return normalized_feat * style_std.expand(size) + style_mean.expand(size) + + +def wavelet_blur(image: Tensor, radius: int): + """ + Apply wavelet blur to the input tensor. + """ + # input shape: (1, 3, H, W) + # convolution kernel + kernel_vals = [ + [0.0625, 0.125, 0.0625], + [0.125, 0.25, 0.125], + [0.0625, 0.125, 0.0625], + ] + kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device) + # add channel dimensions to the kernel to make it a 4D tensor + kernel = kernel[None, None] + # repeat the kernel across all input channels + kernel = kernel.repeat(3, 1, 1, 1) + image = F.pad(image, (radius, radius, radius, radius), mode='replicate') + # apply convolution + output = F.conv2d(image, kernel, groups=3, dilation=radius) + return output + + +def wavelet_decomposition(image: Tensor, levels=5): + """ + Apply wavelet decomposition to the input tensor. + This function only returns the low frequency & the high frequency. + """ + high_freq = torch.zeros_like(image) + for i in range(levels): + radius = 2**i + low_freq = wavelet_blur(image, radius) + high_freq += (image - low_freq) + image = low_freq + + return high_freq, low_freq + + +def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor): + """ + Apply wavelet decomposition, so that the content will have the same color as the style. + """ + # calculate the wavelet decomposition of the content feature + content_high_freq, content_low_freq = wavelet_decomposition(content_feat) + del content_low_freq + # calculate the wavelet decomposition of the style feature + style_high_freq, style_low_freq = wavelet_decomposition(style_feat) + del style_high_freq + # reconstruct the content feature with the style's high frequency + return content_high_freq + style_low_freq + + +def load_dreambooth_lora(unet, vae=None, model_path=None, model_base=''): + if model_path is None: + return unet + + if model_path.endswith('.ckpt'): + base_state_dict = torch.load(model_path)['state_dict'] + elif model_path.endswith('.safetensors'): + state_dict = {} + with safe_open(model_path, framework='pt', device='cpu') as f: + for key in f.keys(): + state_dict[key] = f.get_tensor(key) + + is_lora = all('lora' in k for k in state_dict.keys()) + if not is_lora: + base_state_dict = state_dict + else: + base_state_dict = {} + with safe_open(model_base, framework='pt', device='cpu') as f: + for key in f.keys(): + base_state_dict[key] = f.get_tensor(key) + + converted_unet_checkpoint = convert_ldm_unet_checkpoint( + base_state_dict, unet.config) + unet.load_state_dict(converted_unet_checkpoint, strict=False) + + if vae is not None: + converted_vae_checkpoint = convert_ldm_vae_checkpoint( + base_state_dict, vae.config) + vae.load_state_dict(converted_vae_checkpoint) + + return unet, vae diff --git a/modelscope/models/cv/image_super_resolution_pasd/transformer_2d.py b/modelscope/models/cv/image_super_resolution_pasd/transformer_2d.py new file mode 100644 index 00000000..ad22d212 --- /dev/null +++ b/modelscope/models/cv/image_super_resolution_pasd/transformer_2d.py @@ -0,0 +1,365 @@ +# Part of the implementation is borrowed and modified from diffusers, +# publicly available at https://github.com/huggingface/diffusers/tree/main/src/diffusers/models/transformer_2d.py +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models import ModelMixin +from diffusers.models.embeddings import ImagePositionalEmbeddings, PatchEmbed +from diffusers.utils import BaseOutput, deprecate +from torch import nn + +from .attention import BasicTransformerBlock + + +@dataclass +class Transformer2DModelOutput(BaseOutput): + """ + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or + `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete): + Hidden states conditioned on `encoder_hidden_states` input. If discrete, returns probability distributions + for the unnoised latent pixels. + """ + + sample: torch.FloatTensor + + +class Transformer2DModel(ModelMixin, ConfigMixin): + """ + Transformer model for image-like data. Takes either discrete (classes of vector embeddings) or continuous (actual + embeddings) inputs. + + When input is continuous: First, project the input (aka embedding) and reshape to b, t, d. Then apply standard + transformer action. Finally, reshape to image. + + When input is discrete: First, input (classes of latent pixels) is converted to embeddings and has positional + embeddings applied, see `ImagePositionalEmbeddings`. Then apply standard transformer action. Finally, predict + classes of unnoised image. + + Note that it is assumed one of the input classes is the masked latent pixel. The predicted classes of the unnoised + image do not contain a prediction for the masked pixel as the unnoised image cannot be masked. + + Parameters: + num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. + attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. + in_channels (`int`, *optional*): + Pass if the input is continuous. The number of channels in the input and output. + num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The number of encoder_hidden_states dimensions to use. + sample_size (`int`, *optional*): Pass if the input is discrete. The width of the latent images. + Note that this is fixed at training time as it is used for learning a number of position embeddings. See + `ImagePositionalEmbeddings`. + num_vector_embeds (`int`, *optional*): + Pass if the input is discrete. The number of classes of the vector embeddings of the latent pixels. + Includes the class for the masked latent pixel. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm ( `int`, *optional*): Pass if at least one of the norm_layers is `AdaLayerNorm`. + The number of diffusion steps used during training. Note that this is fixed at training time as it is used + to learn a number of embeddings that are added to the hidden states. During inference, you can denoise for + up to but not more than steps than `num_embeds_ada_norm`. + attention_bias (`bool`, *optional*): + Configure if the TransformerBlocks' attention should contain a bias parameter. + """ + + @register_to_config + def __init__( + self, + num_attention_heads: int = 16, + attention_head_dim: int = 88, + in_channels: Optional[int] = None, + out_channels: Optional[int] = None, + num_layers: int = 1, + dropout: float = 0.0, + norm_num_groups: int = 32, + cross_attention_dim: Optional[int] = None, + pixelwise_cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + sample_size: Optional[int] = None, + num_vector_embeds: Optional[int] = None, + patch_size: Optional[int] = None, + activation_fn: str = 'geglu', + num_embeds_ada_norm: Optional[int] = None, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + norm_type: str = 'layer_norm', + norm_elementwise_affine: bool = True, + use_pixelwise_attention: bool = False, + ): + super().__init__() + self.use_linear_projection = use_linear_projection + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + + # 1. Transformer2DModel can process both standard continuous images of shape + # `(batch_size, num_channels, width, height)` + # as well as quantized image embeddings of shape `(batch_size, num_image_vectors)` + # Define whether input is continuous or discrete depending on configuration + self.is_input_continuous = (in_channels + is not None) and (patch_size is None) + self.is_input_vectorized = num_vector_embeds is not None + self.is_input_patches = in_channels is not None and patch_size is not None + + if norm_type == 'layer_norm' and num_embeds_ada_norm is not None: + deprecation_message = ( + f'The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or' + " incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config." + ' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect' + ' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it' + ' would be very nice if you could open a Pull request for the `transformer/config.json` file' + ) + deprecate( + 'norm_type!=num_embeds_ada_norm', + '1.0.0', + deprecation_message, + standard_warn=False) + norm_type = 'ada_norm' + + if self.is_input_continuous and self.is_input_vectorized: + raise ValueError( + f'Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make' + ' sure that either `in_channels` or `num_vector_embeds` is None.' + ) + elif self.is_input_vectorized and self.is_input_patches: + raise ValueError( + f'Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make' + ' sure that either `num_vector_embeds` or `num_patches` is None.' + ) + elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches: + raise ValueError( + f'Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:' + f' {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None.' + ) + + # 2. Define input layers + if self.is_input_continuous: + self.in_channels = in_channels + + self.norm = torch.nn.GroupNorm( + num_groups=norm_num_groups, + num_channels=in_channels, + eps=1e-6, + affine=True) + if use_linear_projection: + self.proj_in = nn.Linear(in_channels, inner_dim) + if use_pixelwise_attention: + self.proj_in_plus = nn.Linear( + pixelwise_cross_attention_dim, + pixelwise_cross_attention_dim) + else: + self.proj_in_plus = None + else: + self.proj_in = nn.Conv2d( + in_channels, inner_dim, kernel_size=1, stride=1, padding=0) + if use_pixelwise_attention: + self.proj_in_plus = nn.Conv2d( + pixelwise_cross_attention_dim, + pixelwise_cross_attention_dim, + kernel_size=1, + stride=1, + padding=0) + else: + self.proj_in_plus = None + + elif self.is_input_vectorized: + assert sample_size is not None, 'Transformer2DModel over discrete input must provide sample_size' + assert num_vector_embeds is not None, 'Transformer2DModel over discrete input must provide num_embed' + + self.height = sample_size + self.width = sample_size + self.num_vector_embeds = num_vector_embeds + self.num_latent_pixels = self.height * self.width + + self.latent_image_embedding = ImagePositionalEmbeddings( + num_embed=num_vector_embeds, + embed_dim=inner_dim, + height=self.height, + width=self.width) + elif self.is_input_patches: + assert sample_size is not None, 'Transformer2DModel over patched input must provide sample_size' + + self.height = sample_size + self.width = sample_size + + self.patch_size = patch_size + self.pos_embed = PatchEmbed( + height=sample_size, + width=sample_size, + patch_size=patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + ) + + # 3. Define transformers blocks + self.transformer_blocks = nn.ModuleList([ + BasicTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + pixelwise_cross_attention_dim=pixelwise_cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + use_pixelwise_attention=use_pixelwise_attention, + ) for d in range(num_layers) + ]) + + # 4. Define output layers + self.out_channels = in_channels if out_channels is None else out_channels + if self.is_input_continuous: + # TODO: should use out_channels for continuous projections + if use_linear_projection: + self.proj_out = nn.Linear(inner_dim, in_channels) + else: + self.proj_out = nn.Conv2d( + inner_dim, in_channels, kernel_size=1, stride=1, padding=0) + elif self.is_input_vectorized: + self.norm_out = nn.LayerNorm(inner_dim) + self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1) + elif self.is_input_patches: + self.norm_out = nn.LayerNorm( + inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim) + self.proj_out_2 = nn.Linear( + inner_dim, patch_size * patch_size * self.out_channels) + + def forward( + self, + hidden_states, + encoder_hidden_states=None, + encoder_pixelwise_hidden_states=None, + timestep=None, + class_labels=None, + cross_attention_kwargs=None, + return_dict: bool = True, + ): + """ + Args: + hidden_states ( When discrete, `torch.LongTensor` of shape `(batch size, num latent pixels)`. + When continuous, `torch.FloatTensor` of shape `(batch size, channel, height, width)`): Input + hidden_states + encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*): + Conditional embeddings for cross attention layer. If not given, cross-attention defaults to + self-attention. + timestep ( `torch.long`, *optional*): + Optional timestep to be applied as an embedding in AdaLayerNorm's. Used to indicate denoising step. + class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): + Optional class labels to be applied as an embedding in AdaLayerZeroNorm. Used to indicate class labels + conditioning. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. + + Returns: + [`~models.transformer_2d.Transformer2DModelOutput`] or `tuple`: + [`~models.transformer_2d.Transformer2DModelOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is the sample tensor. + """ + # 1. Input + if self.is_input_continuous: + batch, _, height, width = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + if not self.use_linear_projection: + hidden_states = self.proj_in(hidden_states) + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch, height * width, inner_dim) + + if self.proj_in_plus is not None: + encoder_pixelwise_hidden_states = self.proj_in_plus( + encoder_pixelwise_hidden_states) + pixelwise_inner_dim = encoder_pixelwise_hidden_states.shape[ + 1] + encoder_pixelwise_hidden_states = encoder_pixelwise_hidden_states.permute( + 0, 2, 3, 1).reshape(batch, height * width, + pixelwise_inner_dim) + else: + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch, height * width, inner_dim) + hidden_states = self.proj_in(hidden_states) + if self.proj_in_plus is not None: + pixelwise_inner_dim = encoder_pixelwise_hidden_states.shape[ + 1] + encoder_pixelwise_hidden_states = encoder_pixelwise_hidden_states.permute( + 0, 2, 3, 1).reshape(batch, height * width, + pixelwise_inner_dim) + encoder_pixelwise_hidden_states = self.proj_in_plus( + encoder_pixelwise_hidden_states) + + elif self.is_input_vectorized: + hidden_states = self.latent_image_embedding(hidden_states) + elif self.is_input_patches: + hidden_states = self.pos_embed(hidden_states) + + # 2. Blocks + for block in self.transformer_blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_pixelwise_hidden_states=encoder_pixelwise_hidden_states, + timestep=timestep, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + ) + + # 3. Output + if self.is_input_continuous: + if not self.use_linear_projection: + hidden_states = hidden_states.reshape(batch, height, width, + inner_dim).permute( + 0, 3, 1, + 2).contiguous() + hidden_states = self.proj_out(hidden_states) + else: + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.reshape(batch, height, width, + inner_dim).permute( + 0, 3, 1, + 2).contiguous() + + output = hidden_states + residual + elif self.is_input_vectorized: + hidden_states = self.norm_out(hidden_states) + logits = self.out(hidden_states) + # (batch, self.num_vector_embeds - 1, self.num_latent_pixels) + logits = logits.permute(0, 2, 1) + + # log(p(x_0)) + output = F.log_softmax(logits.double(), dim=1).float() + elif self.is_input_patches: + # TODO: cleanup! + conditioning = self.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype) + shift, scale = self.proj_out_1(F.silu(conditioning)).chunk( + 2, dim=1) + hidden_states = self.norm_out(hidden_states) * ( + 1 + scale[:, None]) + shift[:, None] + hidden_states = self.proj_out_2(hidden_states) + + # unpatchify + height = width = int(hidden_states.shape[1]**0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, self.patch_size, self.patch_size, + self.out_channels)) + hidden_states = torch.einsum('nhwpqc->nchpwq', hidden_states) + output = hidden_states.reshape( + shape=(-1, self.out_channels, height * self.patch_size, + width * self.patch_size)) + + if not return_dict: + return (output, ) + + return Transformer2DModelOutput(sample=output) diff --git a/modelscope/models/cv/image_super_resolution_pasd/unet_2d_blocks.py b/modelscope/models/cv/image_super_resolution_pasd/unet_2d_blocks.py new file mode 100644 index 00000000..c95639e6 --- /dev/null +++ b/modelscope/models/cv/image_super_resolution_pasd/unet_2d_blocks.py @@ -0,0 +1,2983 @@ +# Part of the implementation is borrowed and modified from diffusers, +# publicly available at https://github.com/huggingface/diffusers/tree/main/src/diffusers/models/unet_2d_blocks.py +from typing import Optional + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers.models.attention_processor import (Attention, + AttnAddedKVProcessor, + AttnAddedKVProcessor2_0) +from diffusers.models.dual_transformer_2d import DualTransformer2DModel +from diffusers.models.resnet import (Downsample2D, FirDownsample2D, + FirUpsample2D, KDownsample2D, KUpsample2D, + ResnetBlock2D, Upsample2D) +from torch import nn + +from .attention import AdaGroupNorm +from .transformer_2d import Transformer2DModel + + +def get_down_block( + down_block_type, + num_layers, + in_channels, + out_channels, + temb_channels, + add_downsample, + resnet_eps, + resnet_act_fn, + attn_num_head_channels, + resnet_groups=None, + cross_attention_dim=None, + downsample_padding=None, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + resnet_time_scale_shift='default', + resnet_skip_time_act=False, + resnet_out_scale_factor=1.0, + cross_attention_norm=None, +): + down_block_type = down_block_type[7:] if down_block_type.startswith( + 'UNetRes') else down_block_type + if down_block_type == 'DownBlock2D': + return DownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == 'ResnetDownsampleBlock2D': + return ResnetDownsampleBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + ) + elif down_block_type == 'AttnDownBlock2D': + return AttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + attn_num_head_channels=attn_num_head_channels, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == 'CrossAttnDownBlock2D': + if cross_attention_dim is None: + raise ValueError( + 'cross_attention_dim must be specified for CrossAttnDownBlock2D' + ) + return CrossAttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == 'SimpleCrossAttnDownBlock2D': + if cross_attention_dim is None: + raise ValueError( + 'cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D' + ) + return SimpleCrossAttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif down_block_type == 'SkipDownBlock2D': + return SkipDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == 'AttnSkipDownBlock2D': + return AttnSkipDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + downsample_padding=downsample_padding, + attn_num_head_channels=attn_num_head_channels, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == 'DownEncoderBlock2D': + return DownEncoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == 'AttnDownEncoderBlock2D': + return AttnDownEncoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + attn_num_head_channels=attn_num_head_channels, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == 'KDownBlock2D': + return KDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + ) + elif down_block_type == 'KCrossAttnDownBlock2D': + return KCrossAttnDownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + add_self_attention=True if not add_downsample else False, + ) + raise ValueError(f'{down_block_type} does not exist.') + + +def get_up_block( + up_block_type, + num_layers, + in_channels, + out_channels, + prev_output_channel, + temb_channels, + add_upsample, + resnet_eps, + resnet_act_fn, + attn_num_head_channels, + resnet_groups=None, + cross_attention_dim=None, + pixelwise_cross_attention_dim=None, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + resnet_time_scale_shift='default', + resnet_skip_time_act=False, + resnet_out_scale_factor=1.0, + cross_attention_norm=None, +): + up_block_type = up_block_type[7:] if up_block_type.startswith( + 'UNetRes') else up_block_type + if up_block_type == 'UpBlock2D': + return UpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == 'ResnetUpsampleBlock2D': + return ResnetUpsampleBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + ) + elif up_block_type == 'CrossAttnUpBlock2D': + if cross_attention_dim is None: + raise ValueError( + 'cross_attention_dim must be specified for CrossAttnUpBlock2D') + return CrossAttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + pixelwise_cross_attention_dim=pixelwise_cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + use_pixelwise_attention=True, + ) + elif up_block_type == 'SimpleCrossAttnUpBlock2D': + if cross_attention_dim is None: + raise ValueError( + 'cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D' + ) + return SimpleCrossAttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + output_scale_factor=resnet_out_scale_factor, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif up_block_type == 'AttnUpBlock2D': + return AttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + attn_num_head_channels=attn_num_head_channels, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == 'SkipUpBlock2D': + return SkipUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == 'AttnSkipUpBlock2D': + return AttnSkipUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + attn_num_head_channels=attn_num_head_channels, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == 'UpDecoderBlock2D': + return UpDecoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == 'AttnUpDecoderBlock2D': + return AttnUpDecoderBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + attn_num_head_channels=attn_num_head_channels, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == 'KUpBlock2D': + return KUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + ) + elif up_block_type == 'KCrossAttnUpBlock2D': + return KCrossAttnUpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + ) + + raise ValueError(f'{up_block_type} does not exist.') + + +class UNetMidBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attn_num_head_channels=1, + output_scale_factor=1.0, + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min( + in_channels // 4, 32) + self.add_attention = add_attention + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attn_num_head_channels + if attn_num_head_channels is not None else 1, + dim_head=attn_num_head_channels + if attn_num_head_channels is not None else in_channels, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + else: + attentions.append(None) + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states, temb=None): + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + hidden_states = attn(hidden_states) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class UNetMidBlock2DCrossAttn(nn.Module): + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=1.0, + cross_attention_dim=1280, + dual_cross_attention=False, + use_linear_projection=False, + upcast_attention=False, + ): + super().__init__() + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + resnet_groups = resnet_groups if resnet_groups is not None else min( + in_channels // 4, 32) + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + for _ in range(num_layers): + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + attn_num_head_channels, + in_channels // attn_num_head_channels, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + )) + else: + attentions.append( + DualTransformer2DModel( + attn_num_head_channels, + in_channels // attn_num_head_channels, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + )) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, + hidden_states, + temb=None, + encoder_hidden_states=None, + attention_mask=None, + cross_attention_kwargs=None): + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class UNetMidBlock2DSimpleCrossAttn(nn.Module): + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=1.0, + cross_attention_dim=1280, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + + self.has_cross_attention = True + + self.attn_num_head_channels = attn_num_head_channels + resnet_groups = resnet_groups if resnet_groups is not None else min( + in_channels // 4, 32) + + self.num_heads = in_channels // self.attn_num_head_channels + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + ) + ] + attentions = [] + + for _ in range(num_layers): + processor = ( + AttnAddedKVProcessor2_0() if hasattr( + F, 'scaled_dot_product_attention') else + AttnAddedKVProcessor()) + + attentions.append( + Attention( + query_dim=in_channels, + cross_attention_dim=in_channels, + heads=self.num_heads, + dim_head=attn_num_head_channels, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + )) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, + hidden_states, + temb=None, + encoder_hidden_states=None, + attention_mask=None, + cross_attention_kwargs=None): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + # attn + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + # resnet + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class AttnDownBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=1.0, + downsample_padding=1, + add_downsample=True, + ): + super().__init__() + resnets = [] + attentions = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + attentions.append( + Attention( + out_channels, + heads=out_channels // attn_num_head_channels + if attn_num_head_channels is not None else 1, + dim_head=attn_num_head_channels + if attn_num_head_channels is not None else out_channels, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name='op') + ]) + else: + self.downsamplers = None + + def forward(self, hidden_states, temb=None, upsample_size=None): + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states) + output_states += (hidden_states, ) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states += (hidden_states, ) + + return hidden_states, output_states + + +class CrossAttnDownBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + downsample_padding=1, + add_downsample=True, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + attn_num_head_channels, + out_channels // attn_num_head_channels, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + )) + else: + attentions.append( + DualTransformer2DModel( + attn_num_head_channels, + out_channels // attn_num_head_channels, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + )) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name='op') + ]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, + hidden_states, + temb=None, + encoder_hidden_states=None, + attention_mask=None, + cross_attention_kwargs=None): + # TODO(Patrick, William) - attention mask is not used + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + cross_attention_kwargs, + )[0] + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + + output_states = output_states + (hidden_states, ) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states = output_states + (hidden_states, ) + + return hidden_states, output_states + + +class DownBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name='op') + ]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None): + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + else: + hidden_states = resnet(hidden_states, temb) + + output_states = output_states + (hidden_states, ) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states = output_states + (hidden_states, ) + + return hidden_states, output_states + + +class DownEncoderBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name='op') + ]) + else: + self.downsamplers = None + + def forward(self, hidden_states): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class AttnDownEncoderBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + resnets = [] + attentions = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + attentions.append( + Attention( + out_channels, + heads=out_channels // attn_num_head_channels + if attn_num_head_channels is not None else 1, + dim_head=attn_num_head_channels + if attn_num_head_channels is not None else out_channels, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name='op') + ]) + else: + self.downsamplers = None + + def forward(self, hidden_states): + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb=None) + hidden_states = attn(hidden_states) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class AttnSkipDownBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=np.sqrt(2.0), + downsample_padding=1, + add_downsample=True, + ): + super().__init__() + self.attentions = nn.ModuleList([]) + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(in_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + self.attentions.append( + Attention( + out_channels, + heads=out_channels // attn_num_head_channels + if attn_num_head_channels is not None else 1, + dim_head=attn_num_head_channels + if attn_num_head_channels is not None else out_channels, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=32, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + + if add_downsample: + self.resnet_down = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + down=True, + kernel='fir', + ) + self.downsamplers = nn.ModuleList( + [FirDownsample2D(out_channels, out_channels=out_channels)]) + self.skip_conv = nn.Conv2d( + 3, out_channels, kernel_size=(1, 1), stride=(1, 1)) + else: + self.resnet_down = None + self.downsamplers = None + self.skip_conv = None + + def forward(self, hidden_states, temb=None, skip_sample=None): + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states) + output_states += (hidden_states, ) + + if self.downsamplers is not None: + hidden_states = self.resnet_down(hidden_states, temb) + for downsampler in self.downsamplers: + skip_sample = downsampler(skip_sample) + + hidden_states = self.skip_conv(skip_sample) + hidden_states + + output_states += (hidden_states, ) + + return hidden_states, output_states, skip_sample + + +class SkipDownBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_pre_norm: bool = True, + output_scale_factor=np.sqrt(2.0), + add_downsample=True, + downsample_padding=1, + ): + super().__init__() + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(in_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + if add_downsample: + self.resnet_down = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + down=True, + kernel='fir', + ) + self.downsamplers = nn.ModuleList( + [FirDownsample2D(out_channels, out_channels=out_channels)]) + self.skip_conv = nn.Conv2d( + 3, out_channels, kernel_size=(1, 1), stride=(1, 1)) + else: + self.resnet_down = None + self.downsamplers = None + self.skip_conv = None + + def forward(self, hidden_states, temb=None, skip_sample=None): + output_states = () + + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb) + output_states += (hidden_states, ) + + if self.downsamplers is not None: + hidden_states = self.resnet_down(hidden_states, temb) + for downsampler in self.downsamplers: + skip_sample = downsampler(skip_sample) + + hidden_states = self.skip_conv(skip_sample) + hidden_states + + output_states += (hidden_states, ) + + return hidden_states, output_states, skip_sample + + +class ResnetDownsampleBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + skip_time_act=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + down=True, + ) + ]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None): + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + else: + hidden_states = resnet(hidden_states, temb) + + output_states = output_states + (hidden_states, ) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, temb) + + output_states = output_states + (hidden_states, ) + + return hidden_states, output_states + + +class SimpleCrossAttnDownBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_downsample=True, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + + self.has_cross_attention = True + + resnets = [] + attentions = [] + + self.attn_num_head_channels = attn_num_head_channels + self.num_heads = out_channels // self.attn_num_head_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + )) + + processor = ( + AttnAddedKVProcessor2_0() if hasattr( + F, 'scaled_dot_product_attention') else + AttnAddedKVProcessor()) + + attentions.append( + Attention( + query_dim=out_channels, + cross_attention_dim=out_channels, + heads=self.num_heads, + dim_head=attn_num_head_channels, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + )) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + down=True, + ) + ]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, + hidden_states, + temb=None, + encoder_hidden_states=None, + attention_mask=None, + cross_attention_kwargs=None): + output_states = () + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + for resnet, attn in zip(self.resnets, self.attentions): + # resnet + hidden_states = resnet(hidden_states, temb) + + # attn + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + output_states = output_states + (hidden_states, ) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, temb) + + output_states = output_states + (hidden_states, ) + + return hidden_states, output_states + + +class KDownBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_eps: float = 1e-5, + resnet_act_fn: str = 'gelu', + resnet_group_size: int = 32, + add_downsample=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + temb_channels=temb_channels, + groups=groups, + groups_out=groups_out, + eps=resnet_eps, + non_linearity=resnet_act_fn, + time_embedding_norm='ada_group', + conv_shortcut_bias=False, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + # YiYi's comments- might be able to use FirDownsample2D, look into details later + self.downsamplers = nn.ModuleList([KDownsample2D()]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None): + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + else: + hidden_states = resnet(hidden_states, temb) + + output_states += (hidden_states, ) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states, output_states + + +class KCrossAttnDownBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + cross_attention_dim: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_group_size: int = 32, + add_downsample=True, + attn_num_head_channels: int = 64, + add_self_attention: bool = False, + resnet_eps: float = 1e-5, + resnet_act_fn: str = 'gelu', + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + temb_channels=temb_channels, + groups=groups, + groups_out=groups_out, + eps=resnet_eps, + non_linearity=resnet_act_fn, + time_embedding_norm='ada_group', + conv_shortcut_bias=False, + )) + attentions.append( + KAttentionBlock( + out_channels, + out_channels // attn_num_head_channels, + attn_num_head_channels, + cross_attention_dim=cross_attention_dim, + temb_channels=temb_channels, + attention_bias=True, + add_self_attention=add_self_attention, + cross_attention_norm='layer_norm', + group_size=resnet_group_size, + )) + + self.resnets = nn.ModuleList(resnets) + self.attentions = nn.ModuleList(attentions) + + if add_downsample: + self.downsamplers = nn.ModuleList([KDownsample2D()]) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, + hidden_states, + temb=None, + encoder_hidden_states=None, + attention_mask=None, + cross_attention_kwargs=None): + output_states = () + + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + attention_mask, + cross_attention_kwargs, + ) + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + + if self.downsamplers is None: + output_states += (None, ) + else: + output_states += (hidden_states, ) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states, output_states + + +class AttnUpBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=1.0, + add_upsample=True, + ): + super().__init__() + resnets = [] + attentions = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers + - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + attentions.append( + Attention( + out_channels, + heads=out_channels // attn_num_head_channels + if attn_num_head_channels is not None else 1, + dim_head=attn_num_head_channels + if attn_num_head_channels is not None else out_channels, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + Upsample2D( + out_channels, use_conv=True, out_channels=out_channels) + ]) + else: + self.upsamplers = None + + def forward(self, + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None): + for resnet, attn in zip(self.resnets, self.attentions): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], + dim=1) + + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class CrossAttnUpBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + cross_attention_dim=1280, + pixelwise_cross_attention_dim=None, + output_scale_factor=1.0, + add_upsample=True, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + use_pixelwise_attention=False, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers + - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + attn_num_head_channels, + out_channels // attn_num_head_channels, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + pixelwise_cross_attention_dim=res_skip_channels, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + use_pixelwise_attention=use_pixelwise_attention, + )) + else: + attentions.append( + DualTransformer2DModel( + attn_num_head_channels, + out_channels // attn_num_head_channels, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + )) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + Upsample2D( + out_channels, use_conv=True, out_channels=out_channels) + ]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + res_hidden_states_tuple, + temb=None, + encoder_hidden_states=None, + encoder_pixelwise_hidden_states_tuple=None, + cross_attention_kwargs=None, + upsample_size=None, + attention_mask=None, + ): + # TODO(Patrick, William) - attention mask is not used + for resnet, attn in zip(self.resnets, self.attentions): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], + dim=1) + + encoder_pixelwise_hidden_states = encoder_pixelwise_hidden_states_tuple[ + -1] + encoder_pixelwise_hidden_states_tuple = encoder_pixelwise_hidden_states_tuple[: + -1] + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + encoder_pixelwise_hidden_states, + cross_attention_kwargs, + )[0] + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_pixelwise_hidden_states= + encoder_pixelwise_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + +class UpBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers + - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + Upsample2D( + out_channels, use_conv=True, out_channels=out_channels) + ]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], + dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + else: + hidden_states = resnet(hidden_states, temb) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + +class UpDecoderBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + Upsample2D( + out_channels, use_conv=True, out_channels=out_channels) + ]) + else: + self.upsamplers = None + + def forward(self, hidden_states): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class AttnUpDecoderBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=1.0, + add_upsample=True, + ): + super().__init__() + resnets = [] + attentions = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + attentions.append( + Attention( + out_channels, + heads=out_channels // attn_num_head_channels + if attn_num_head_channels is not None else 1, + dim_head=attn_num_head_channels + if attn_num_head_channels is not None else out_channels, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=resnet_groups, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + Upsample2D( + out_channels, use_conv=True, out_channels=out_channels) + ]) + else: + self.upsamplers = None + + def forward(self, hidden_states): + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb=None) + hidden_states = attn(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class AttnSkipUpBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=np.sqrt(2.0), + upsample_padding=1, + add_upsample=True, + ): + super().__init__() + self.attentions = nn.ModuleList([]) + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers + - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + self.resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(resnet_in_channels + res_skip_channels // 4, + 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.attentions.append( + Attention( + out_channels, + heads=out_channels // attn_num_head_channels + if attn_num_head_channels is not None else 1, + dim_head=attn_num_head_channels + if attn_num_head_channels is not None else out_channels, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=32, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + + self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) + if add_upsample: + self.resnet_up = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + up=True, + kernel='fir', + ) + self.skip_conv = nn.Conv2d( + out_channels, + 3, + kernel_size=(3, 3), + stride=(1, 1), + padding=(1, 1)) + self.skip_norm = torch.nn.GroupNorm( + num_groups=min(out_channels // 4, 32), + num_channels=out_channels, + eps=resnet_eps, + affine=True) + self.act = nn.SiLU() + else: + self.resnet_up = None + self.skip_conv = None + self.skip_norm = None + self.act = None + + def forward(self, + hidden_states, + res_hidden_states_tuple, + temb=None, + skip_sample=None): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], + dim=1) + + hidden_states = resnet(hidden_states, temb) + + hidden_states = self.attentions[0](hidden_states) + + if skip_sample is not None: + skip_sample = self.upsampler(skip_sample) + else: + skip_sample = 0 + + if self.resnet_up is not None: + skip_sample_states = self.skip_norm(hidden_states) + skip_sample_states = self.act(skip_sample_states) + skip_sample_states = self.skip_conv(skip_sample_states) + + skip_sample = skip_sample + skip_sample_states + + hidden_states = self.resnet_up(hidden_states, temb) + + return hidden_states, skip_sample + + +class SkipUpBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_pre_norm: bool = True, + output_scale_factor=np.sqrt(2.0), + add_upsample=True, + upsample_padding=1, + ): + super().__init__() + self.resnets = nn.ModuleList([]) + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers + - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + self.resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min((resnet_in_channels + res_skip_channels) // 4, + 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) + if add_upsample: + self.resnet_up = ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=min(out_channels // 4, 32), + groups_out=min(out_channels // 4, 32), + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_in_shortcut=True, + up=True, + kernel='fir', + ) + self.skip_conv = nn.Conv2d( + out_channels, + 3, + kernel_size=(3, 3), + stride=(1, 1), + padding=(1, 1)) + self.skip_norm = torch.nn.GroupNorm( + num_groups=min(out_channels // 4, 32), + num_channels=out_channels, + eps=resnet_eps, + affine=True) + self.act = nn.SiLU() + else: + self.resnet_up = None + self.skip_conv = None + self.skip_norm = None + self.act = None + + def forward(self, + hidden_states, + res_hidden_states_tuple, + temb=None, + skip_sample=None): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], + dim=1) + + hidden_states = resnet(hidden_states, temb) + + if skip_sample is not None: + skip_sample = self.upsampler(skip_sample) + else: + skip_sample = 0 + + if self.resnet_up is not None: + skip_sample_states = self.skip_norm(hidden_states) + skip_sample_states = self.act(skip_sample_states) + skip_sample_states = self.skip_conv(skip_sample_states) + + skip_sample = skip_sample + skip_sample_states + + hidden_states = self.resnet_up(hidden_states, temb) + + return hidden_states, skip_sample + + +class ResnetUpsampleBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + skip_time_act=False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers + - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + up=True, + ) + ]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None): + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], + dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + else: + hidden_states = resnet(hidden_states, temb) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, temb) + + return hidden_states + + +class SimpleCrossAttnUpBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = 'default', + resnet_act_fn: str = 'swish', + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_upsample=True, + skip_time_act=False, + only_cross_attention=False, + cross_attention_norm=None, + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + + self.num_heads = out_channels // self.attn_num_head_channels + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers + - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + )) + + processor = ( + AttnAddedKVProcessor2_0() if hasattr( + F, 'scaled_dot_product_attention') else + AttnAddedKVProcessor()) + + attentions.append( + Attention( + query_dim=out_channels, + cross_attention_dim=out_channels, + heads=self.num_heads, + dim_head=attn_num_head_channels, + added_kv_proj_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + bias=True, + upcast_softmax=True, + only_cross_attention=only_cross_attention, + cross_attention_norm=cross_attention_norm, + processor=processor, + )) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + ResnetBlock2D( + in_channels=out_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + skip_time_act=skip_time_act, + up=True, + ) + ]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + res_hidden_states_tuple, + temb=None, + encoder_hidden_states=None, + upsample_size=None, + attention_mask=None, + cross_attention_kwargs=None, + ): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + for resnet, attn in zip(self.resnets, self.attentions): + # resnet + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], + dim=1) + + hidden_states = resnet(hidden_states, temb) + + # attn + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, temb) + + return hidden_states + + +class KUpBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 5, + resnet_eps: float = 1e-5, + resnet_act_fn: str = 'gelu', + resnet_group_size: Optional[int] = 32, + add_upsample=True, + ): + super().__init__() + resnets = [] + k_in_channels = 2 * out_channels + k_out_channels = in_channels + num_layers = num_layers - 1 + + for i in range(num_layers): + in_channels = k_in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=k_out_channels if + (i == num_layers - 1) else out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=groups, + groups_out=groups_out, + dropout=dropout, + non_linearity=resnet_act_fn, + time_embedding_norm='ada_group', + conv_shortcut_bias=False, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([KUpsample2D()]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None): + res_hidden_states_tuple = res_hidden_states_tuple[-1] + if res_hidden_states_tuple is not None: + hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], + dim=1) + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + else: + hidden_states = resnet(hidden_states, temb) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class KCrossAttnUpBlock2D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 4, + resnet_eps: float = 1e-5, + resnet_act_fn: str = 'gelu', + resnet_group_size: int = 32, + attn_num_head_channels=1, # attention dim_head + cross_attention_dim: int = 768, + add_upsample: bool = True, + upcast_attention: bool = False, + ): + super().__init__() + resnets = [] + attentions = [] + + is_first_block = in_channels == out_channels == temb_channels + is_middle_block = in_channels != out_channels + add_self_attention = True if is_first_block else False + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + + # in_channels, and out_channels for the block (k-unet) + k_in_channels = out_channels if is_first_block else 2 * out_channels + k_out_channels = in_channels + + num_layers = num_layers - 1 + + for i in range(num_layers): + in_channels = k_in_channels if i == 0 else out_channels + groups = in_channels // resnet_group_size + groups_out = out_channels // resnet_group_size + + if is_middle_block and (i == num_layers - 1): + conv_2d_out_channels = k_out_channels + else: + conv_2d_out_channels = None + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + conv_2d_out_channels=conv_2d_out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=groups, + groups_out=groups_out, + dropout=dropout, + non_linearity=resnet_act_fn, + time_embedding_norm='ada_group', + conv_shortcut_bias=False, + )) + attentions.append( + KAttentionBlock( + k_out_channels if (i == num_layers - 1) else out_channels, + k_out_channels // attn_num_head_channels if + (i == num_layers + - 1) else out_channels // attn_num_head_channels, + attn_num_head_channels, + cross_attention_dim=cross_attention_dim, + temb_channels=temb_channels, + attention_bias=True, + add_self_attention=add_self_attention, + cross_attention_norm='layer_norm', + upcast_attention=upcast_attention, + )) + + self.resnets = nn.ModuleList(resnets) + self.attentions = nn.ModuleList(attentions) + + if add_upsample: + self.upsamplers = nn.ModuleList([KUpsample2D()]) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + res_hidden_states_tuple, + temb=None, + encoder_hidden_states=None, + cross_attention_kwargs=None, + upsample_size=None, + attention_mask=None, + ): + res_hidden_states_tuple = res_hidden_states_tuple[-1] + if res_hidden_states_tuple is not None: + hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], + dim=1) + + for resnet, attn in zip(self.resnets, self.attentions): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + attention_mask, + cross_attention_kwargs, + )[0] + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + emb=temb, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +# can potentially later be renamed to `No-feed-forward` attention +class KAttentionBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout: float = 0.0, + cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + upcast_attention: bool = False, + temb_channels: int = 768, # for ada_group_norm + add_self_attention: bool = False, + cross_attention_norm: Optional[str] = None, + group_size: int = 32, + ): + super().__init__() + self.add_self_attention = add_self_attention + + # 1. Self-Attn + if add_self_attention: + self.norm1 = AdaGroupNorm(temb_channels, dim, + max(1, dim // group_size)) + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=None, + cross_attention_norm=None, + ) + + # 2. Cross-Attn + self.norm2 = AdaGroupNorm(temb_channels, dim, + max(1, dim // group_size)) + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + cross_attention_norm=cross_attention_norm, + ) + + def _to_3d(self, hidden_states, height, weight): + return hidden_states.permute(0, 2, 3, + 1).reshape(hidden_states.shape[0], + height * weight, -1) + + def _to_4d(self, hidden_states, height, weight): + return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], + -1, height, weight) + + def forward( + self, + hidden_states, + encoder_hidden_states=None, + emb=None, + attention_mask=None, + cross_attention_kwargs=None, + ): + cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {} + + # 1. Self-Attention + if self.add_self_attention: + norm_hidden_states = self.norm1(hidden_states, emb) + + height, weight = norm_hidden_states.shape[2:] + norm_hidden_states = self._to_3d(norm_hidden_states, height, + weight) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=None, + **cross_attention_kwargs, + ) + attn_output = self._to_4d(attn_output, height, weight) + + hidden_states = attn_output + hidden_states + + # 2. Cross-Attention/None + norm_hidden_states = self.norm2(hidden_states, emb) + + height, weight = norm_hidden_states.shape[2:] + norm_hidden_states = self._to_3d(norm_hidden_states, height, weight) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + **cross_attention_kwargs, + ) + attn_output = self._to_4d(attn_output, height, weight) + + hidden_states = attn_output + hidden_states + + return hidden_states diff --git a/modelscope/models/cv/image_super_resolution_pasd/unet_2d_condition.py b/modelscope/models/cv/image_super_resolution_pasd/unet_2d_condition.py new file mode 100644 index 00000000..539d3d1c --- /dev/null +++ b/modelscope/models/cv/image_super_resolution_pasd/unet_2d_condition.py @@ -0,0 +1,887 @@ +# Part of the implementation is borrowed and modified from diffusers, +# publicly available at https://github.com/huggingface/diffusers/tree/main/src/diffusers/models/unet_2d_condition.py +import os +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import json +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import UNet2DConditionLoadersMixin +from diffusers.models import ModelMixin +from diffusers.models.attention_processor import (AttentionProcessor, + AttnProcessor) +from diffusers.models.embeddings import (GaussianFourierProjection, + TextTimeEmbedding, TimestepEmbedding, + Timesteps) +from diffusers.utils import BaseOutput, logging + +from .unet_2d_blocks import (CrossAttnDownBlock2D, CrossAttnUpBlock2D, + DownBlock2D, UNetMidBlock2DCrossAttn, + UNetMidBlock2DSimpleCrossAttn, UpBlock2D, + get_down_block, get_up_block) + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class UNet2DConditionOutput(BaseOutput): + """ + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model. + """ + + sample: torch.FloatTensor + + +class UNet2DConditionModel(ModelMixin, ConfigMixin, + UNet2DConditionLoadersMixin): + r""" + UNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a timestep + and returns sample shaped output. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library + implements for all the models (such as downloading or saving, etc.) + + Parameters: + sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + Height and width of input/output sample. + in_channels (`int`, *optional*, defaults to 4): The number of channels in the input sample. + out_channels (`int`, *optional*, defaults to 4): The number of channels in the output. + center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. + flip_sin_to_cos (`bool`, *optional*, defaults to `False`): + Whether to flip the sin to cos in the time embedding. + freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. + down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", "DownBlock2D")`): + The tuple of downsample blocks to use. + mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): + The mid block type. Choose from `UNetMidBlock2DCrossAttn` or `UNetMidBlock2DSimpleCrossAttn`, will skip the + mid block layer if `None`. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", + "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`): + The tuple of upsample blocks to use. + only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`): + Whether to include self-attention in the basic transformer blocks, see + [`~models.attention.BasicTransformerBlock`]. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + The tuple of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. + downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. + mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. + If `None`, it will skip the normalization and activation layers in post-processing + norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. + cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + The dimension of the cross attention features. + encoder_hid_dim (`int`, *optional*, defaults to None): + If given, `encoder_hidden_states` will be projected from this dimension to `cross_attention_dim`. + attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads. + resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config + for resnet blocks, see [`~models.resnet.ResnetBlock2D`]. Choose from `default` or `scale_shift`. + class_embed_type (`str`, *optional*, defaults to None): + The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, + `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. + addition_embed_type (`str`, *optional*, defaults to None): + Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or + "text". "text" will use the `TextTimeEmbedding` layer. + num_class_embeds (`int`, *optional*, defaults to None): + Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing + class conditioning with `class_embed_type` equal to `None`. + time_embedding_type (`str`, *optional*, default to `positional`): + The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. + time_embedding_dim (`int`, *optional*, default to `None`): + An optional override for the dimension of the projected time embedding. + time_embedding_act_fn (`str`, *optional*, default to `None`): + Optional activation function to use on the time embeddings only one time before they as passed to the rest + of the unet. Choose from `silu`, `mish`, `gelu`, and `swish`. + timestep_post_act (`str, *optional*, default to `None`): + The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. + time_cond_proj_dim (`int`, *optional*, default to `None`): + The dimension of `cond_proj` layer in timestep embedding. + conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. + conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer. + projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when + using the "projection" `class_embed_type`. Required when using the "projection" `class_embed_type`. + class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time + embeddings with the class embeddings. + mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`): + Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If + `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is None, the + `only_cross_attention` value will be used as the value for `mid_block_only_cross_attention`. Else, it will + default to `False`. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + 'CrossAttnDownBlock2D', + 'CrossAttnDownBlock2D', + 'CrossAttnDownBlock2D', + 'DownBlock2D', + ), + mid_block_type: Optional[str] = 'UNetMidBlock2DCrossAttn', + up_block_types: Tuple[str] = ('UpBlock2D', 'CrossAttnUpBlock2D', + 'CrossAttnUpBlock2D', + 'CrossAttnUpBlock2D'), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: Union[int, Tuple[int]] = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + act_fn: str = 'silu', + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: Union[int, Tuple[int]] = 1280, + encoder_hid_dim: Optional[int] = None, + attention_head_dim: Union[int, Tuple[int]] = 8, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + addition_embed_type: Optional[str] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = 'default', + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: int = 1.0, + time_embedding_type: str = 'positional', + time_embedding_dim: Optional[int] = None, + time_embedding_act_fn: Optional[str] = None, + timestep_post_act: Optional[str] = None, + time_cond_proj_dim: Optional[int] = None, + conv_in_kernel: int = 3, + conv_out_kernel: int = 3, + projection_class_embeddings_input_dim: Optional[int] = None, + class_embeddings_concat: bool = False, + mid_block_only_cross_attention: Optional[bool] = None, + cross_attention_norm: Optional[str] = None, + addition_embed_type_num_heads=64, + ): + super().__init__() + + self.sample_size = sample_size + + # Check inputs + if len(down_block_types) != len(up_block_types): + raise ValueError( + f'Must provide the same number of `down_block_types` as `up_block_types`. \ + `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}.' + ) + + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f'Must provide the same number of `block_out_channels` as `down_block_types`. \ + `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}.' + ) + + if not isinstance( + only_cross_attention, + bool) and len(only_cross_attention) != len(down_block_types): + raise ValueError( + f'Must provide the same number of `only_cross_attention` as `down_block_types`. \ + `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}.' + ) + + if not isinstance( + attention_head_dim, + int) and len(attention_head_dim) != len(down_block_types): + raise ValueError( + f'Must provide the same number of `attention_head_dim` as `down_block_types`. \ + `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}.' + ) + + if isinstance( + cross_attention_dim, + list) and len(cross_attention_dim) != len(down_block_types): + raise ValueError( + f'Must provide the same number of `cross_attention_dim` as `down_block_types`. \ + `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}.' + ) + + if not isinstance( + layers_per_block, + int) and len(layers_per_block) != len(down_block_types): + raise ValueError( + f'Must provide the same number of `layers_per_block` as `down_block_types`. \ + `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}.' + ) + + # input + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d( + in_channels, + block_out_channels[0], + kernel_size=conv_in_kernel, + padding=conv_in_padding) + + # time + if time_embedding_type == 'fourier': + time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 + if time_embed_dim % 2 != 0: + raise ValueError( + f'`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.' + ) + self.time_proj = GaussianFourierProjection( + time_embed_dim // 2, + set_W_to_weight=False, + log=False, + flip_sin_to_cos=flip_sin_to_cos) + timestep_input_dim = time_embed_dim + elif time_embedding_type == 'positional': + time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 + + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, + freq_shift) + timestep_input_dim = block_out_channels[0] + else: + raise ValueError( + f'{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`.' + ) + + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + post_act_fn=timestep_post_act, + cond_proj_dim=time_cond_proj_dim, + ) + + if encoder_hid_dim is not None: + self.encoder_hid_proj = nn.Linear(encoder_hid_dim, + cross_attention_dim) + else: + self.encoder_hid_proj = None + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, + time_embed_dim) + elif class_embed_type == 'timestep': + self.class_embedding = TimestepEmbedding( + timestep_input_dim, time_embed_dim, act_fn=act_fn) + elif class_embed_type == 'identity': + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + elif class_embed_type == 'projection': + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" + ) + # The projection `class_embed_type` is the same as the timestep `class_embed_type` except + # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings + # 2. it projects from an arbitrary input dimension. + # + # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. + # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. + # As a result, `TimestepEmbedding` can be passed arbitrary vectors. + self.class_embedding = TimestepEmbedding( + projection_class_embeddings_input_dim, time_embed_dim) + elif class_embed_type == 'simple_projection': + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" + ) + self.class_embedding = nn.Linear( + projection_class_embeddings_input_dim, time_embed_dim) + else: + self.class_embedding = None + + if addition_embed_type == 'text': + if encoder_hid_dim is not None: + text_time_embedding_from_dim = encoder_hid_dim + else: + text_time_embedding_from_dim = cross_attention_dim + + self.add_embedding = TextTimeEmbedding( + text_time_embedding_from_dim, + time_embed_dim, + num_heads=addition_embed_type_num_heads) + elif addition_embed_type is not None: + raise ValueError( + f"addition_embed_type: {addition_embed_type} must be None or 'text'." + ) + + if time_embedding_act_fn is None: + self.time_embed_act = None + elif time_embedding_act_fn == 'swish': + self.time_embed_act = lambda x: F.silu(x) + elif time_embedding_act_fn == 'mish': + self.time_embed_act = nn.Mish() + elif time_embedding_act_fn == 'silu': + self.time_embed_act = nn.SiLU() + elif time_embedding_act_fn == 'gelu': + self.time_embed_act = nn.GELU() + else: + raise ValueError( + f'Unsupported activation function: {time_embedding_act_fn}') + + self.down_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = only_cross_attention + + only_cross_attention = [only_cross_attention + ] * len(down_block_types) + + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = False + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim, ) * len(down_block_types) + + if isinstance(cross_attention_dim, int): + cross_attention_dim = ( + cross_attention_dim, ) * len(down_block_types) + + if isinstance(layers_per_block, int): + layers_per_block = [layers_per_block] * len(down_block_types) + + if class_embeddings_concat: + # The time embeddings are concatenated with the class embeddings. The dimension of the + # time embeddings passed to the down, middle, and up blocks is twice the dimension of the + # regular time embeddings + blocks_time_embed_dim = time_embed_dim * 2 + else: + blocks_time_embed_dim = time_embed_dim + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=blocks_time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim[i], + attn_num_head_channels=attention_head_dim[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == 'UNetMidBlock2DCrossAttn': + self.mid_block = UNetMidBlock2DCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim[-1], + attn_num_head_channels=attention_head_dim[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + ) + elif mid_block_type == 'UNetMidBlock2DSimpleCrossAttn': + self.mid_block = UNetMidBlock2DSimpleCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + cross_attention_dim=cross_attention_dim[-1], + attn_num_head_channels=attention_head_dim[-1], + resnet_groups=norm_num_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + skip_time_act=resnet_skip_time_act, + only_cross_attention=mid_block_only_cross_attention, + cross_attention_norm=cross_attention_norm, + ) + elif mid_block_type is None: + self.mid_block = None + else: + raise ValueError(f'unknown mid_block_type : {mid_block_type}') + + # count how many layers upsample the images + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_attention_head_dim = list(reversed(attention_head_dim)) + reversed_layers_per_block = list(reversed(layers_per_block)) + reversed_cross_attention_dim = list(reversed(cross_attention_dim)) + only_cross_attention = list(reversed(only_cross_attention)) + + reversed_pixelwise_cross_attention_dim = [-1, 1280, 640, 320] + + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[min( + i + 1, + len(block_out_channels) - 1)] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=reversed_layers_per_block[i] + 1, + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=blocks_time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=reversed_cross_attention_dim[i], + pixelwise_cross_attention_dim= + reversed_pixelwise_cross_attention_dim[i], + attn_num_head_channels=reversed_attention_head_dim[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_num_groups is not None: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=norm_eps) + + if act_fn == 'swish': + self.conv_act = lambda x: F.silu(x) + elif act_fn == 'mish': + self.conv_act = nn.Mish() + elif act_fn == 'silu': + self.conv_act = nn.SiLU() + elif act_fn == 'gelu': + self.conv_act = nn.GELU() + else: + raise ValueError(f'Unsupported activation function: {act_fn}') + + else: + self.conv_norm_out = None + self.conv_act = None + + conv_out_padding = (conv_out_kernel - 1) // 2 + self.conv_out = nn.Conv2d( + block_out_channels[0], + out_channels, + kernel_size=conv_out_kernel, + padding=conv_out_padding) + + @property + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, + processors: Dict[str, + AttentionProcessor]): + if hasattr(module, 'set_processor'): + processors[f'{name}.processor'] = module.processor + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f'{name}.{sub_name}', child, + processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor(self, processor: Union[AttentionProcessor, + Dict[str, + AttentionProcessor]]): + r""" + Parameters: + `processor (`dict` of `AttentionProcessor` or `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + of **all** `Attention` layers. + In case `processor` is a dict, the key needs to define the path to + the corresponding cross attention processor. + This is strongly recommended when setting trainable attention processors.: + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f'A dict of processors was passed, but the number of processors {len(processor)} does not match the' + f' number of attention layers: {count}. Please make sure to pass {count} processor classes.' + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, + processor): + if hasattr(module, 'set_processor'): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f'{name}.processor')) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f'{name}.{sub_name}', child, + processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + self.set_attn_processor(AttnProcessor()) + + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module will split the input tensor in slices, to compute attention + in several steps. This is useful to save some memory in exchange for a small speed decrease. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If + `"max"`, maximum amount of memory will be saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): + if hasattr(module, 'set_attention_slice'): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_sliceable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_sliceable_dims(module) + + num_sliceable_layers = len(sliceable_head_dims) + + if slice_size == 'auto': + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == 'max': + # make smallest slice possible + slice_size = num_sliceable_layers * [1] + + slice_size = num_sliceable_layers * [slice_size] if not isinstance( + slice_size, list) else slice_size + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f'You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different' + f' attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}.' + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError( + f'size {size} has to be smaller or equal to {dim}.') + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice(module: torch.nn.Module, + slice_size: List[int]): + if hasattr(module, 'set_attention_slice'): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (CrossAttnDownBlock2D, DownBlock2D, + CrossAttnUpBlock2D, UpBlock2D)): + module.gradient_checkpointing = value + + def forward( + self, + sample: torch.FloatTensor, + 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, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + return_dict: bool = True, + ) -> Union[UNet2DConditionOutput, Tuple]: + r""" + Args: + sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor + timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps + encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py). + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): + logger.info( + 'Forward upsample size to force interpolation output size.') + forward_upsample_size = True + + # prepare attention_mask + if attention_mask is not None: + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # 0. center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == 'mps' + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], + dtype=dtype, + device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb, timestep_cond) + + if self.class_embedding is not None: + if class_labels is None: + raise ValueError( + 'class_labels should be provided when num_class_embeds > 0' + ) + + if self.config.class_embed_type == 'timestep': + class_labels = self.time_proj(class_labels) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # there might be better ways to encapsulate this. + class_labels = class_labels.to(dtype=sample.dtype) + + class_emb = self.class_embedding(class_labels).to( + dtype=sample.dtype) + + if self.config.class_embeddings_concat: + emb = torch.cat([emb, class_emb], dim=-1) + else: + emb = emb + class_emb + + if self.config.addition_embed_type == 'text': + aug_emb = self.add_embedding(encoder_hidden_states) + emb = emb + aug_emb + + if self.time_embed_act is not None: + emb = self.time_embed_act(emb) + + if self.encoder_hid_proj is not None: + encoder_hidden_states = self.encoder_hid_proj( + encoder_hidden_states) + + # 2. pre-process + sample = self.conv_in(sample) + + # 3. down + down_block_res_samples = (sample, ) + for downsample_block in self.down_blocks: + if hasattr(downsample_block, 'has_cross_attention' + ) and downsample_block.has_cross_attention: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + else: + sample, res_samples = downsample_block( + hidden_states=sample, temb=emb) + + down_block_res_samples += res_samples + + if down_block_additional_residuals is not None: + new_down_block_res_samples = () + + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals): + down_block_res_sample = down_block_res_sample + down_block_additional_residual + new_down_block_res_samples = new_down_block_res_samples + ( + down_block_res_sample, ) + + down_block_res_samples = new_down_block_res_samples + + # 4. mid + if self.mid_block is not None: + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + ) + + if mid_block_additional_residual is not None: + sample = sample + mid_block_additional_residual + + # 5. up + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets):] + down_block_res_samples = down_block_res_samples[:-len( + upsample_block.resnets)] + + down_block_additional_residual = down_block_additional_residuals[ + -len(upsample_block.resnets):] + down_block_additional_residuals = down_block_additional_residuals[:-len( + upsample_block.resnets)] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if hasattr(upsample_block, 'has_cross_attention' + ) and upsample_block.has_cross_attention: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + encoder_pixelwise_hidden_states_tuple= + down_block_additional_residual, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size) + + # 6. post-process + if self.conv_norm_out: + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + if not return_dict: + return (sample, ) + + return UNet2DConditionOutput(sample=sample) + + @classmethod + def from_pretrained_(cls, pretrained_model_path, subfolder=None, **kwargs): + if subfolder is not None: + pretrained_model_path = os.path.join(pretrained_model_path, + subfolder) + + config_file = os.path.join(pretrained_model_path, 'config.json') + if not os.path.isfile(config_file): + raise RuntimeError(f'{config_file} does not exist') + with open(config_file, 'r') as f: + config = json.load(f) + + from diffusers.utils import WEIGHTS_NAME + model = cls.from_config(config) + model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME) + if not os.path.isfile(model_file): + raise RuntimeError(f'{model_file} does not exist') + state_dict = torch.load(model_file, map_location='cpu') + for k, v in model.state_dict().items(): + if 'attn2_plus' in k: + state_dict.update({k: v}) + model.load_state_dict(state_dict, strict=False) + + return model diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index e0ce7b50..d9c6147e 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -727,6 +727,7 @@ TASK_OUTPUTS = { # {"output_img": np.array with shape (h, w, 3)} Tasks.skin_retouching: [OutputKeys.OUTPUT_IMG], Tasks.image_super_resolution: [OutputKeys.OUTPUT_IMG], + Tasks.image_super_resolution_pasd: [OutputKeys.OUTPUT_IMG], Tasks.image_colorization: [OutputKeys.OUTPUT_IMG], Tasks.image_color_enhancement: [OutputKeys.OUTPUT_IMG], Tasks.image_denoising: [OutputKeys.OUTPUT_IMG], diff --git a/modelscope/pipeline_inputs.py b/modelscope/pipeline_inputs.py index 784d67d4..df0d4794 100644 --- a/modelscope/pipeline_inputs.py +++ b/modelscope/pipeline_inputs.py @@ -182,6 +182,10 @@ TASK_INPUTS = { InputType.IMAGE, Tasks.crowd_counting: InputType.IMAGE, + Tasks.image_super_resolution_pasd: { + 'image': InputType.IMAGE, + 'prompt': InputType.TEXT, + }, Tasks.image_inpainting: { 'img': InputType.IMAGE, 'mask': InputType.IMAGE, diff --git a/modelscope/pipelines/cv/__init__.py b/modelscope/pipelines/cv/__init__.py index 126bc22a..15ebb80e 100644 --- a/modelscope/pipelines/cv/__init__.py +++ b/modelscope/pipelines/cv/__init__.py @@ -38,6 +38,7 @@ if TYPE_CHECKING: from .image_semantic_segmentation_pipeline import ImageSemanticSegmentationPipeline from .image_style_transfer_pipeline import ImageStyleTransferPipeline from .image_super_resolution_pipeline import ImageSuperResolutionPipeline + from .image_super_resolution_pasd_pipeline import ImageSuperResolutionPASDPipeline from .image_to_image_generate_pipeline import Image2ImageGenerationPipeline from .image_to_image_translation_pipeline import Image2ImageTranslationPipeline from .image_inpainting_pipeline import ImageInpaintingPipeline @@ -151,6 +152,8 @@ else: ['ImageSemanticSegmentationPipeline'], 'image_style_transfer_pipeline': ['ImageStyleTransferPipeline'], 'image_super_resolution_pipeline': ['ImageSuperResolutionPipeline'], + 'image_super_resolution_pasd_pipeline': + ['ImageSuperResolutionPASDPipeline'], 'image_to_image_translation_pipeline': ['Image2ImageTranslationPipeline'], 'product_retrieval_embedding_pipeline': @@ -185,8 +188,9 @@ else: ['FaceAttributeRecognitionPipeline'], 'mtcnn_face_detection_pipeline': ['MtcnnFaceDetectionPipeline'], 'hand_static_pipeline': ['HandStaticPipeline'], - 'referring_video_object_segmentation_pipeline': - ['ReferringVideoObjectSegmentationPipeline'], + 'referring_video_object_segmentation_pipeline': [ + 'ReferringVideoObjectSegmentationPipeline' + ], 'language_guided_video_summarization_pipeline': [ 'LanguageGuidedVideoSummarizationPipeline' ], diff --git a/modelscope/pipelines/cv/image_super_resolution_pasd_pipeline.py b/modelscope/pipelines/cv/image_super_resolution_pasd_pipeline.py new file mode 100644 index 00000000..c4451483 --- /dev/null +++ b/modelscope/pipelines/cv/image_super_resolution_pasd_pipeline.py @@ -0,0 +1,229 @@ +# Copyright © Alibaba, Inc. and its affiliates. + +import os +import tempfile +from typing import Any, Dict, Optional, Union + +import numpy as np +import PIL +import torch +from diffusers import AutoencoderKL, UniPCMultistepScheduler +from torchvision.models import ResNet50_Weights, resnet50 +from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer + +from modelscope.metainfo import Pipelines +from modelscope.models.cv.image_portrait_enhancement.retinaface import \ + detection +from modelscope.models.cv.image_super_resolution_pasd import ( + ControlNetModel, UNet2DConditionModel) +from modelscope.models.cv.image_super_resolution_pasd.misc import ( + load_dreambooth_lora, wavelet_color_fix) +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input, Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.multi_modal.diffusers_wrapped.pasd_pipeline import \ + PixelAwareStableDiffusionPipeline +from modelscope.preprocessors.image import load_image +from modelscope.utils.config import Config +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.device import create_device + + +@PIPELINES.register_module( + Tasks.image_super_resolution_pasd, + module_name=Pipelines.image_super_resolution_pasd) +class ImageSuperResolutionPASDPipeline(Pipeline): + """ Pixel-Aware Stable Diffusion for Realistic Image Super-Resolution Pipeline. + + Example: + + >>> import cv2 + >>> from modelscope.outputs import OutputKeys + >>> from modelscope.pipelines import pipeline + >>> from modelscope.utils.constant import Tasks + + >>> input_location = 'example_image.png' + + >>> input = { + >>> 'image': input_location, + >>> 'upscale': 2, + >>> 'prompt': prompt, + >>> 'fidelity_scale_fg': 1.5, + >>> 'fidelity_scale_bg': 0.7 + >>> } + >>> pasd = pipeline(Tasks.image_super_resolution_pasd, model='damo/PASD_image_super_resolutions') + >>> output = pasd(input)[OutputKeys.OUTPUT_IMG] + >>> cv2.imwrite('result.png', output) + + """ + + def __init__(self, model: str, device_name: str = 'cuda', **kwargs): + """ + use `model` to create a image super resolution pipeline for prediction + Args: + model: model id on modelscope hub. + """ + super().__init__(model=model, **kwargs) + + torch_dtype = kwargs.get('torch_dtype', torch.float16) + self.device = create_device(device_name) + self.config = Config.from_file( + os.path.join(model, ModelFile.CONFIGURATION)) + cfg = self.config.model_cfg + dreambooth_lora_ckpt = cfg['dreambooth_lora_ckpt'] + tiled_size = cfg['tiled_size'] + self.process_size = cfg['process_size'] + + scheduler = UniPCMultistepScheduler.from_pretrained( + model, subfolder='scheduler') + text_encoder = CLIPTextModel.from_pretrained( + model, subfolder='text_encoder') + tokenizer = CLIPTokenizer.from_pretrained(model, subfolder='tokenizer') + vae = AutoencoderKL.from_pretrained(model, subfolder='vae') + feature_extractor = CLIPImageProcessor.from_pretrained( + f'{model}/feature_extractor') + unet = UNet2DConditionModel.from_pretrained(model, subfolder='unet') + controlnet = ControlNetModel.from_pretrained( + model, subfolder='controlnet') + + unet, vae = load_dreambooth_lora(unet, vae, + f'{model}/{dreambooth_lora_ckpt}') + + vae.requires_grad_(False) + text_encoder.requires_grad_(False) + unet.requires_grad_(False) + controlnet.requires_grad_(False) + + text_encoder.to(self.device, dtype=torch_dtype) + vae.to(self.device, dtype=torch_dtype) + unet.to(self.device, dtype=torch_dtype) + controlnet.to(self.device, dtype=torch_dtype) + + self.pipeline = PixelAwareStableDiffusionPipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + feature_extractor=feature_extractor, + unet=unet, + controlnet=controlnet, + scheduler=scheduler, + safety_checker=None, + requires_safety_checker=False, + ) + self.pipeline._init_tiled_vae(decoder_tile_size=tiled_size) + self.pipeline.enable_model_cpu_offload() + + self.weights = ResNet50_Weights.DEFAULT + self.resnet_preprocess = self.weights.transforms() + self.resnet = resnet50(weights=self.weights) + self.resnet.eval() + + self.threshold = 0.8 + detector_model_path = f'{model}/RetinaFace-R50.pth' + self.face_detector = detection.RetinaFaceDetection( + detector_model_path, self.device) + + def preprocess(self, input: Input): + return input + + def forward(self, inputs: Dict[str, Any]): + if not isinstance(inputs, dict): + raise ValueError( + f'Expected the input to be a dictionary, but got {type(input)}' + ) + + num_inference_steps = inputs.get('num_inference_steps', 20) + guidance_scale = inputs.get('guidance_scale', 7.5) + added_prompt = inputs.get( + 'added_prompt', + 'clean, high-resolution, 8k, best quality, masterpiece, extremely detailed' + ) + negative_prompt = inputs.get( + 'negative_prompt', + 'dotted, noise, blur, lowres, smooth, longbody, bad anatomy, bad hands, missing fingers, extra digit, \ + fewer digits, cropped, worst quality, low quality') + eta = inputs.get('eta', 0.0) + prompt = inputs.get('prompt', '') + upscale = inputs.get('upscale', 2) + fidelity_scale_fg = inputs.get('fidelity_scale_fg', 1.5) + fidelity_scale_bg = inputs.get('fidelity_scale_bg', 0.7) + + input_image = load_image(inputs['image']).convert('RGB') + + with torch.no_grad(): + generator = torch.Generator(device=self.device) + + batch = self.resnet_preprocess(input_image).unsqueeze(0) + prediction = self.resnet(batch).squeeze(0).softmax(0) + class_id = prediction.argmax().item() + score = prediction[class_id].item() + category_name = self.weights.meta['categories'][class_id] + if score >= 0.1: + prompt += f'{category_name}' if prompt == '' else f', {category_name}' + + prompt = added_prompt if prompt == '' else f'{prompt}, {added_prompt}' + + ori_width, ori_height = input_image.size + resize_flag = False + rscale = upscale + if ori_width < self.process_size // rscale or ori_height < self.process_size // rscale: + scale = (self.process_size // rscale) / min( + ori_width, ori_height) + tmp_image = input_image.resize( + (int(scale * ori_width), int(scale * ori_height))) + + input_image = tmp_image + resize_flag = True + + input_image = input_image.resize( + (input_image.size[0] * rscale, input_image.size[1] * rscale)) + input_image = input_image.resize( + (input_image.size[0] // 8 * 8, input_image.size[1] // 8 * 8)) + width, height = input_image.size + + fg_mask = None + if fidelity_scale_fg != fidelity_scale_bg: + fg_mask = torch.zeros([1, 1, height, width]) + facebs, _ = self.face_detector.detect(np.array(input_image)) + for fb in facebs: + if fb[-1] < self.threshold: + continue + fb = list(map(int, fb)) + fg_mask[:, :, fb[1]:fb[3], fb[0]:fb[2]] = 1 + fg_mask = fg_mask.to(self.device) + + if fg_mask is None: + fidelity_scale = min( + max(fidelity_scale_fg, fidelity_scale_bg), 1) + fidelity_scale_fg = fidelity_scale_bg = fidelity_scale + + try: + image = self.pipeline( + prompt, + input_image, + num_inference_steps=num_inference_steps, + generator=generator, + height=height, + width=width, + guidance_scale=guidance_scale, + negative_prompt=negative_prompt, + conditioning_scale_fg=fidelity_scale_fg, + conditioning_scale_bg=fidelity_scale_bg, + fg_mask=fg_mask, + eta=eta, + ).images[0] + + image = wavelet_color_fix(image, input_image) + + if resize_flag: + image = image.resize( + (ori_width * rscale, ori_height * rscale)) + except Exception as e: + print(e) + image = PIL.Image.new('RGB', (512, 512), (0, 0, 0)) + + return {'result': image} + + def postprocess(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + result = np.array(inputs['result']) + return {OutputKeys.OUTPUT_IMG: result[:, :, ::-1]} diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/devices.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/devices.py new file mode 100644 index 00000000..d7cebd32 --- /dev/null +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/devices.py @@ -0,0 +1,126 @@ +# The implementation is adopted from stable-diffusion-webui, made public available under the Apache 2.0 License +# at https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/devices.py + +import contextlib +import sys + +import torch + +if sys.platform == 'darwin': + from modules import mac_specific + + +def has_mps() -> bool: + if sys.platform != 'darwin': + return False + else: + return mac_specific.has_mps + + +def get_cuda_device_string(): + return 'cuda' + + +def get_optimal_device_name(): + if torch.cuda.is_available(): + return get_cuda_device_string() + + if has_mps(): + return 'mps' + + return 'cpu' + + +def get_optimal_device(): + return torch.device(get_optimal_device_name()) + + +def get_device_for(task): + return get_optimal_device() + + +def torch_gc(): + + if torch.cuda.is_available(): + with torch.cuda.device(get_cuda_device_string()): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + if has_mps(): + mac_specific.torch_mps_gc() + + +def enable_tf32(): + if torch.cuda.is_available(): + + # enabling benchmark option seems to enable a range of cards to do fp16 when they otherwise can't + # see https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4407 + if any( + torch.cuda.get_device_capability(devid) == (7, 5) + for devid in range(0, torch.cuda.device_count())): + torch.backends.cudnn.benchmark = True + + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + +enable_tf32() + +cpu = torch.device('cpu') +device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = torch.device( + 'cuda') +dtype = torch.float16 +dtype_vae = torch.float16 +dtype_unet = torch.float16 +unet_needs_upcast = False + + +def cond_cast_unet(input): + return input.to(dtype_unet) if unet_needs_upcast else input + + +def cond_cast_float(input): + return input.float() if unet_needs_upcast else input + + +def randn(seed, shape): + torch.manual_seed(seed) + return torch.randn(shape, device=device) + + +def randn_without_seed(shape): + return torch.randn(shape, device=device) + + +def autocast(disable=False): + if disable: + return contextlib.nullcontext() + + return torch.autocast('cuda') + + +def without_autocast(disable=False): + return torch.autocast('cuda', enabled=False) if torch.is_autocast_enabled() and \ + not disable else contextlib.nullcontext() + + +class NansException(Exception): + pass + + +def test_for_nans(x, where): + if not torch.all(torch.isnan(x)).item(): + return + + if where == 'unet': + message = 'A tensor with all NaNs was produced in Unet.' + + elif where == 'vae': + message = 'A tensor with all NaNs was produced in VAE.' + + else: + message = 'A tensor with all NaNs was produced.' + + message += ' Use --disable-nan-check commandline argument to disable this check.' + + raise NansException(message) diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py new file mode 100644 index 00000000..7ebd3ce2 --- /dev/null +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py @@ -0,0 +1,1131 @@ +# Part of the implementation is borrowed and modified from diffusers, publicly available at +# https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/controlnet/pipeline_controlnet.py + +import inspect +import os +import warnings +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F +from diffusers.image_processor import VaeImageProcessor +from diffusers.loaders import TextualInversionLoaderMixin +from diffusers.models import (AutoencoderKL, ControlNetModel, + UNet2DConditionModel) +from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput +from diffusers.pipelines.stable_diffusion.safety_checker import \ + StableDiffusionSafetyChecker +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import (PIL_INTERPOLATION, is_accelerate_available, + is_accelerate_version, is_compiled_module, + logging, randn_tensor, replace_example_docstring) +from torchvision.utils import save_image +from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer + +from .vaehook import VAEHook + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> # !pip install opencv-python transformers accelerate + >>> from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler + >>> from diffusers.utils import load_image + >>> import numpy as np + >>> import torch + + >>> import cv2 + >>> from PIL import Image + + >>> # download an image + >>> image = load_image( + ... "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png" + ... ) + >>> image = np.array(image) + + >>> # get canny image + >>> image = cv2.Canny(image, 100, 200) + >>> image = image[:, :, None] + >>> image = np.concatenate([image, image, image], axis=2) + >>> canny_image = Image.fromarray(image) + + >>> # load control net and stable diffusion v1-5 + >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) + >>> pipe = StableDiffusionControlNetPipeline.from_pretrained( + ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 + ... ) + + >>> # speed up diffusion process with faster scheduler and memory optimization + >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) + >>> # remove following line if xformers is not installed + >>> pipe.enable_xformers_memory_efficient_attention() + + >>> pipe.enable_model_cpu_offload() + + >>> # generate image + >>> generator = torch.manual_seed(0) + >>> image = pipe( + ... "futuristic-looking woman", num_inference_steps=20, generator=generator, image=canny_image + ... ).images[0] + ``` +""" + + +class PixelAwareStableDiffusionPipeline(DiffusionPipeline, + TextualInversionLoaderMixin): + r""" + Pipeline for text-to-image generation using Stable Diffusion with ControlNet guidance. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the + library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) + + In addition the pipeline inherits the following loading methods: + - *Textual-Inversion*: [`loaders.TextualInversionLoaderMixin.load_textual_inversion`] + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + text_encoder ([`CLIPTextModel`]): + Frozen text-encoder. Stable Diffusion uses the text portion of + [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically + the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. + tokenizer (`CLIPTokenizer`): + Tokenizer of class + [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). + unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. + controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): + Provides additional conditioning to the unet during the denoising process. If you set multiple ControlNets + as a list, the outputs from each ControlNet are added together to create one combined additional + conditioning. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of + [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. + safety_checker ([`StableDiffusionSafetyChecker`]): + Classification module that estimates whether generated images could be considered offensive or harmful. + Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details. + feature_extractor ([`CLIPImageProcessor`]): + Model that extracts features from generated images to be used as inputs for the `safety_checker`. + """ + _optional_components = ['safety_checker', 'feature_extractor'] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + controlnet: Union[ControlNetModel, List[ControlNetModel], + Tuple[ControlNetModel], MultiControlNetModel], + scheduler: KarrasDiffusionSchedulers, + safety_checker: StableDiffusionSafetyChecker, + feature_extractor: CLIPImageProcessor, + requires_safety_checker: bool = True, + ): + super().__init__() + + if safety_checker is None and requires_safety_checker: + logger.warning( + f'You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure' + ' that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered' + ' results in services or applications open to the public. Both the diffusers team and Hugging Face' + ' strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling' + ' it only for use-cases that involve analyzing network behavior or auditing its results. For more' + ' information, please have a look at https://github.com/huggingface/diffusers/pull/254 .' + ) + + if safety_checker is not None and feature_extractor is None: + raise ValueError( + 'Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety' + " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." + ) + + if isinstance(controlnet, (list, tuple)): + controlnet = MultiControlNetModel(controlnet) + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + controlnet=controlnet, + scheduler=scheduler, + safety_checker=safety_checker, + feature_extractor=feature_extractor, + ) + self.vae_scale_factor = 2**( + len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor) + self.register_to_config( + requires_safety_checker=requires_safety_checker) + + def _init_tiled_vae(self, + encoder_tile_size=256, + decoder_tile_size=256, + fast_decoder=False, + fast_encoder=False, + color_fix=False, + vae_to_gpu=True): + # save original forward (only once) + if not hasattr(self.vae.encoder, 'original_forward'): + setattr(self.vae.encoder, 'original_forward', + self.vae.encoder.forward) + if not hasattr(self.vae.decoder, 'original_forward'): + setattr(self.vae.decoder, 'original_forward', + self.vae.decoder.forward) + + encoder = self.vae.encoder + decoder = self.vae.decoder + + self.vae.encoder.forward = VAEHook( + encoder, + encoder_tile_size, + is_decoder=False, + fast_decoder=fast_decoder, + fast_encoder=fast_encoder, + color_fix=color_fix, + to_gpu=vae_to_gpu) + self.vae.decoder.forward = VAEHook( + decoder, + decoder_tile_size, + is_decoder=True, + fast_decoder=fast_decoder, + fast_encoder=fast_encoder, + color_fix=color_fix, + to_gpu=vae_to_gpu) + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline.enable_vae_slicing + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. + + When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several + steps. This is useful to save some memory and allow larger batch sizes. + """ + self.vae.enable_slicing() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline.disable_vae_slicing + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to + computing decoding in one step. + """ + self.vae.disable_slicing() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline.enable_vae_tiling + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. + + When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in + several steps. This is useful to save a large amount of memory and to allow the processing of larger images. + """ + self.vae.enable_tiling() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline.disable_vae_tiling + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to + computing decoding in one step. + """ + self.vae.disable_tiling() + + def enable_sequential_cpu_offload(self, gpu_id=0): + r""" + Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, + text_encoder, vae, controlnet, and safety checker have their state dicts saved to CPU and then are moved to a + `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called. + Note that offloading happens on a submodule basis. Memory savings are higher than with + `enable_model_cpu_offload`, but performance is lower. + """ + if is_accelerate_available(): + from accelerate import cpu_offload + else: + raise ImportError( + 'Please install accelerate via `pip install accelerate`') + + device = torch.device(f'cuda:{gpu_id}') + + for cpu_offloaded_model in [ + self.unet, self.text_encoder, self.vae, self.controlnet + ]: + cpu_offload(cpu_offloaded_model, device) + + if self.safety_checker is not None: + cpu_offload( + self.safety_checker, + execution_device=device, + offload_buffers=True) + + def enable_model_cpu_offload(self, gpu_id=0): + r""" + Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared + to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward` + method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with + `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`. + """ + if is_accelerate_available() and is_accelerate_version( + '>=', '0.17.0.dev0'): + from accelerate import cpu_offload_with_hook + else: + raise ImportError( + '`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.' + ) + + device = torch.device(f'cuda:{gpu_id}') + + hook = None + for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]: + _, hook = cpu_offload_with_hook( + cpu_offloaded_model, device, prev_module_hook=hook) + + if self.safety_checker is not None: + # the safety checker can offload the vae again + _, hook = cpu_offload_with_hook( + self.safety_checker, device, prev_module_hook=hook) + + # control net hook has be manually offloaded as it alternates with unet + cpu_offload_with_hook(self.controlnet, device) + + # We'll offload the last model manually. + self.final_offload_hook = hook + + @property + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline._execution_device + def _execution_device(self): + r""" + Returns the device on which the pipeline's models will be executed. After calling + `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module + hooks. + """ + if not hasattr(self.unet, '_hf_hook'): + return self.device + for module in self.unet.modules(): + if (hasattr(module, '_hf_hook') + and hasattr(module._hf_hook, 'execution_device') + and module._hf_hook.execution_device is not None): + return torch.device(module._hf_hook.execution_device) + return self.device + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline._encode_prompt + def _encode_prompt( + self, + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt=None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + num_images_per_prompt (`int`): + number of images that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + """ + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + # textual inversion: procecss multi-vector tokens if necessary + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, self.tokenizer) + + text_inputs = self.tokenizer( + prompt, + padding='max_length', + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors='pt', + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer( + prompt, padding='longest', return_tensors='pt').input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[ + -1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer.batch_decode( + untruncated_ids[:, self.tokenizer.model_max_length - 1:-1]) + logger.warning( + 'The following part of your input was truncated because CLIP can only handle sequences up to' + f' {self.tokenizer.model_max_length} tokens: {removed_text}' + ) + + if hasattr(self.text_encoder.config, 'use_attention_mask' + ) and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + prompt_embeds = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + prompt_embeds = prompt_embeds[0] + + prompt_embeds = prompt_embeds.to( + dtype=self.text_encoder.dtype, device=device) + + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, + seq_len, -1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance and negative_prompt_embeds is None: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [''] * batch_size + elif prompt is not None and type(prompt) is not type( + negative_prompt): + raise TypeError( + f'`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=' + f' {type(prompt)}.') + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f'`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:' + f' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches' + ' the batch size of `prompt`.') + else: + uncond_tokens = negative_prompt + + # textual inversion: procecss multi-vector tokens if necessary + if isinstance(self, TextualInversionLoaderMixin): + uncond_tokens = self.maybe_convert_prompt( + uncond_tokens, self.tokenizer) + + max_length = prompt_embeds.shape[1] + uncond_input = self.tokenizer( + uncond_tokens, + padding='max_length', + max_length=max_length, + truncation=True, + return_tensors='pt', + ) + + if hasattr(self.text_encoder.config, 'use_attention_mask' + ) and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + negative_prompt_embeds = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + negative_prompt_embeds = negative_prompt_embeds[0] + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + negative_prompt_embeds = negative_prompt_embeds.to( + dtype=self.text_encoder.dtype, device=device) + + negative_prompt_embeds = negative_prompt_embeds.repeat( + 1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view( + batch_size * num_images_per_prompt, seq_len, -1) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + + return prompt_embeds + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline.run_safety_checker + def run_safety_checker(self, image, device, dtype): + if self.safety_checker is None: + has_nsfw_concept = None + else: + if torch.is_tensor(image): + feature_extractor_input = self.image_processor.postprocess( + image, output_type='pil') + else: + feature_extractor_input = self.image_processor.numpy_to_pil( + image) + safety_checker_input = self.feature_extractor( + feature_extractor_input, return_tensors='pt').to(device) + image, has_nsfw_concept = self.safety_checker( + images=image, + clip_input=safety_checker_input.pixel_values.to(dtype)) + return image, has_nsfw_concept + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline.decode_latents + def decode_latents(self, latents): + warnings.warn( + 'The decode_latents method is deprecated and will be removed in a future version. Please' + ' use VaeImageProcessor instead', + FutureWarning, + ) + latents = 1 / self.vae.config.scaling_factor * latents + image = self.vae.decode(latents, return_dict=False)[0] + image = (image / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + image = image.cpu().permute(0, 2, 3, 1).float().numpy() + return image + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline.prepare_extra_step_kwargs + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = 'eta' in set( + inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs['eta'] = eta + + # check if the scheduler accepts generator + accepts_generator = 'generator' in set( + inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs['generator'] = generator + return extra_step_kwargs + + def check_inputs( + self, + prompt, + image, + height, + width, + callback_steps, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + controlnet_conditioning_scale=1.0, + ): + if height % 8 != 0 or width % 8 != 0: + raise ValueError( + f'`height` and `width` have to be divisible by 8 but are {height} and {width}.' + ) + + flag = callback_steps is not None and ( + not isinstance(callback_steps, int) or callback_steps <= 0) + if (callback_steps is None) or flag: + raise ValueError( + f'`callback_steps` has to be a positive integer but is {callback_steps} of type' + f' {type(callback_steps)}.') + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f'Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to' + ' only forward one of the two.') + elif prompt is None and prompt_embeds is None: + raise ValueError( + 'Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.' + ) + elif prompt is not None and (not isinstance(prompt, str) + and not isinstance(prompt, list)): + raise ValueError( + f'`prompt` has to be of type `str` or `list` but is {type(prompt)}' + ) + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f'Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:' + f' {negative_prompt_embeds}. Please make sure to only forward one of the two.' + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + '`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but' + f' got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`' + f' {negative_prompt_embeds.shape}.') + + # `prompt` needs more sophisticated handling when there are multiple + # conditionings. + if isinstance(self.controlnet, MultiControlNetModel): + if isinstance(prompt, list): + logger.warning( + f'You have {len(self.controlnet.nets)} ControlNets and you have passed {len(prompt)}' + ' prompts. The conditionings will be fixed across the prompts.' + ) + + # Check `image` + is_compiled = hasattr( + F, 'scaled_dot_product_attention') and isinstance( + self.controlnet, torch._dynamo.eval_frame.OptimizedModule) + if (isinstance(self.controlnet, ControlNetModel) or is_compiled + and isinstance(self.controlnet._orig_mod, ControlNetModel)): + self.check_image(image, prompt, prompt_embeds) + elif (isinstance(self.controlnet, MultiControlNetModel) or is_compiled + and isinstance(self.controlnet._orig_mod, MultiControlNetModel)): + if not isinstance(image, list): + raise TypeError( + 'For multiple controlnets: `image` must be type `list`') + + # When `image` is a nested list: + # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]]) + elif any(isinstance(i, list) for i in image): + raise ValueError( + 'A single batch of multiple conditionings are supported at the moment.' + ) + elif len(image) != len(self.controlnet.nets): + raise ValueError( + 'For multiple controlnets: `image` must have the same length as the number of controlnets.' + ) + + for image_ in image: + self.check_image(image_, prompt, prompt_embeds) + else: + assert False + + # Check `controlnet_conditioning_scale` + if (isinstance(self.controlnet, ControlNetModel) or is_compiled + and isinstance(self.controlnet._orig_mod, ControlNetModel)): + if not isinstance(controlnet_conditioning_scale, float): + raise TypeError( + 'For single controlnet: `controlnet_conditioning_scale` must be type `float`.' + ) + elif (isinstance(self.controlnet, MultiControlNetModel) or is_compiled + and isinstance(self.controlnet._orig_mod, MultiControlNetModel)): + if isinstance(controlnet_conditioning_scale, list): + if any( + isinstance(i, list) + for i in controlnet_conditioning_scale): + raise ValueError( + 'A single batch of multiple conditionings are supported at the moment.' + ) + elif isinstance( + controlnet_conditioning_scale, + list) and len(controlnet_conditioning_scale) != len( + self.controlnet.nets): + raise ValueError( + 'For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`,' + 'it must have the same length as the number of controlnets' + ) + else: + assert False + + def check_image(self, image, prompt, prompt_embeds): + image_is_pil = isinstance(image, PIL.Image.Image) + image_is_tensor = isinstance(image, torch.Tensor) + image_is_pil_list = isinstance(image, list) and isinstance( + image[0], PIL.Image.Image) + image_is_tensor_list = isinstance(image, list) and isinstance( + image[0], torch.Tensor) + + if not image_is_pil and not image_is_tensor and not image_is_pil_list and not image_is_tensor_list: + raise TypeError( + 'image must be passed and be one of PIL image, torch tensor, list of PIL images,' + 'or list of torch tensors') + + if image_is_pil: + image_batch_size = 1 + elif image_is_tensor: + image_batch_size = image.shape[0] + elif image_is_pil_list: + image_batch_size = len(image) + elif image_is_tensor_list: + image_batch_size = len(image) + + if prompt is not None and isinstance(prompt, str): + prompt_batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + prompt_batch_size = len(prompt) + elif prompt_embeds is not None: + prompt_batch_size = prompt_embeds.shape[0] + + if image_batch_size != 1 and image_batch_size != prompt_batch_size: + raise ValueError( + f'If image batch size is not 1, image batch size must be same as prompt batch size. \ + image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}' + ) + + def prepare_image( + self, + image, + width, + height, + batch_size, + num_images_per_prompt, + device, + dtype, + do_classifier_free_guidance=False, + guess_mode=False, + ): + if not isinstance(image, torch.Tensor): + if isinstance(image, PIL.Image.Image): + image = [image] + + if isinstance(image[0], PIL.Image.Image): + images = [] + + for image_ in image: + image_ = image_.convert('RGB') + image_ = np.array(image_) + image_ = image_[None, :] + images.append(image_) + + image = images + + image = np.concatenate(image, axis=0) + image = np.array(image).astype(np.float32) / 255.0 + image = image.transpose(0, 3, 1, 2) + image = torch.from_numpy(image) + elif isinstance(image[0], torch.Tensor): + image = torch.cat(image, dim=0) + + image_batch_size = image.shape[0] + + if image_batch_size == 1: + repeat_by = batch_size + else: + # image batch size is the same as prompt batch size + repeat_by = num_images_per_prompt + + image = image.repeat_interleave(repeat_by, dim=0) + + image = image.to(device=device, dtype=dtype) + + if do_classifier_free_guidance and not guess_mode: + image = torch.cat([image] * 2) + + return image + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion. + # StableDiffusionPipeline.prepare_latents + def prepare_latents(self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None): + shape = (batch_size, num_channels_latents, + height // self.vae_scale_factor, + width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f'You have passed a list of generators of length {len(generator)}, but requested an effective batch' + f' size of {batch_size}. Make sure the batch size matches the length of the generators.' + ) + + if latents is None: + latents = randn_tensor( + shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + return latents + + def _default_height_width(self, height, width, image): + # NOTE: It is possible that a list of images have different + # dimensions for each image, so just checking the first image + # is not _exactly_ correct, but it is simple. + while isinstance(image, list): + image = image[0] + + if height is None: + if isinstance(image, PIL.Image.Image): + height = image.height + elif isinstance(image, torch.Tensor): + height = image.shape[2] + + height = (height // 8) * 8 # round down to nearest multiple of 8 + + if width is None: + if isinstance(image, PIL.Image.Image): + width = image.width + elif isinstance(image, torch.Tensor): + width = image.shape[3] + + width = (width // 8) * 8 # round down to nearest multiple of 8 + + return height, width + + # override DiffusionPipeline + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + safe_serialization: bool = False, + variant: Optional[str] = None, + ): + if isinstance(self.controlnet, ControlNetModel): + super().save_pretrained(save_directory, safe_serialization, + variant) + else: + raise NotImplementedError( + 'Currently, the `save_pretrained()` is not implemented for Multi-ControlNet.' + ) + + def _gaussian_weights(self, tile_width, tile_height, nbatches): + """Generates a gaussian mask of weights for tile contributions""" + from numpy import pi, exp, sqrt + import numpy as np + + latent_width = tile_width + latent_height = tile_height + + var = 0.01 + midpoint = (latent_width - 1) / 2 + x_probs = [] + for x in range(latent_width): + tmp = -(x - midpoint) * (x - midpoint) / ( + latent_width * latent_width) / (2 * var) + tmp = exp(tmp) / sqrt(2 * pi * var) + x_probs.append(tmp) + + midpoint = latent_height / 2 + y_probs = [] + for y in range(latent_height): + tmp = -(y - midpoint) * (y - midpoint) / ( + latent_height * latent_height) / (2 * var) + tmp = exp(tmp) / sqrt(2 * pi * var) + y_probs.append(tmp) + + weights = np.outer(y_probs, x_probs) + return torch.tile( + torch.tensor(weights, device=self.device), + (nbatches, self.unet.config.in_channels, 1, 1)) + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + image: Union[torch.FloatTensor, PIL.Image.Image, + List[torch.FloatTensor], List[PIL.Image.Image]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, + List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = 'pil', + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], + None]] = None, + callback_steps: int = 1, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + fg_mask: Optional[torch.FloatTensor] = None, + conditioning_scale_fg: Union[float, List[float]] = 1.0, + conditioning_scale_bg: Union[float, List[float]] = 1.0, + guess_mode: bool = False, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + image (`torch.FloatTensor`, `PIL.Image.Image`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, + `List[List[torch.FloatTensor]]`, or `List[List[PIL.Image.Image]]`): + The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If + the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can + also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If + height and/or width are passed, `image` is resized according to them. If multiple ControlNets are + specified in init, images must be passed as a list such that each element of the list can be correctly + batched for input to a single controlnet. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, *optional*, defaults to 7.5): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to + [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a + plain tuple. + callback (`Callable`, *optional*): + A function that will be called every `callback_steps` steps during inference. The function will be + called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function will be called. If not specified, the callback will be + called at every step. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py). + controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0): + The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added + to the residual in the original unet. If multiple ControlNets are specified in init, you can set the + corresponding scale as a list. + guess_mode (`bool`, *optional*, defaults to `False`): + In this mode, the ControlNet encoder will try best to recognize the content of the input image even if + you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended. + + Examples: + + Returns: + [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: + [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. + When returning a tuple, the first element is a list with the generated images, and the second element is a + list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" + (nsfw) content, according to the `safety_checker`. + """ + # 0. Default height and width to unet + height, width = self._default_height_width(height, width, image) + + # 1. Check inputs. Raise error if not correct + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + controlnet = self.controlnet._orig_mod if is_compiled_module( + self.controlnet) else self.controlnet + + guess_mode = guess_mode + + # 3. Encode input prompt + prompt_embeds = self._encode_prompt( + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + ) + + # 4. Prepare image + image = self.prepare_image( + image=image, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + guess_mode=guess_mode, + ) + + # 5. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 6. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 8. Denoising loop + num_warmup_steps = len( + timesteps) - num_inference_steps * self.scheduler.order + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat( + [latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t) + + # controlnet(s) inference + if guess_mode and do_classifier_free_guidance: + # Infer ControlNet only for the conditional batch. + controlnet_latent_model_input = latents + controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] + else: + controlnet_latent_model_input = latent_model_input + controlnet_prompt_embeds = prompt_embeds + + if image is not None: + rgbs, down_block_res_samples, mid_block_res_sample = self.controlnet( + controlnet_latent_model_input, + t, + encoder_hidden_states=controlnet_prompt_embeds, + controlnet_cond=image, + fg_mask=fg_mask, + conditioning_scale_fg=conditioning_scale_fg, + conditioning_scale_bg=conditioning_scale_bg, + guess_mode=guess_mode, + return_dict=False, + ) + else: + down_block_res_samples, mid_block_res_sample = [ + None + ] * 10, [None] * 10 + + if guess_mode and do_classifier_free_guidance: + # Infered ControlNet only for the conditional batch. + # To apply the output of ControlNet to both the unconditional and conditional batches, + # add 0 to the unconditional batch to keep it unchanged. + down_block_res_samples = [ + torch.cat([torch.zeros_like(d), d]) + for d in down_block_res_samples + ] + mid_block_res_sample = torch.cat([ + torch.zeros_like(mid_block_res_sample), + mid_block_res_sample + ]) + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + return_dict=False, + )[0] + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * ( + noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step( + noise_pred, + t, + latents, + **extra_step_kwargs, + return_dict=False)[0] + + # call the callback, if provided + flag = ((i + 1) > num_warmup_steps + and (i + 1) % self.scheduler.order == 0) + if i == len(timesteps) - 1 or flag: + progress_bar.update() + if callback is not None and i % callback_steps == 0: + callback(i, t, latents) + + # If we do sequential model offloading, let's offload unet and controlnet + # manually for max memory savings + if hasattr( + self, + 'final_offload_hook') and self.final_offload_hook is not None: + self.unet.to('cpu') + self.controlnet.to('cpu') + torch.cuda.empty_cache() + + has_nsfw_concept = None + if not output_type == 'latent': + image = self.vae.decode( + latents / self.vae.config.scaling_factor, return_dict=False)[0] + else: + image = latents + has_nsfw_concept = None + + if has_nsfw_concept is None: + do_denormalize = [True] * image.shape[0] + else: + do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] + + image = self.image_processor.postprocess( + image, output_type=output_type, do_denormalize=do_denormalize) + + # Offload last model to CPU + if hasattr( + self, + 'final_offload_hook') and self.final_offload_hook is not None: + self.final_offload_hook.offload() + + if not return_dict: + return (image, has_nsfw_concept) + + return StableDiffusionPipelineOutput( + images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/vaehook.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/vaehook.py new file mode 100644 index 00000000..fb03886d --- /dev/null +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/vaehook.py @@ -0,0 +1,762 @@ +# Part of the implementation is borrowed and modified from multidiffusion-upscaler-for-automatic1111, publicly available +# at https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111/blob/main/scripts/vae_optimize.py + +import gc +import math +from time import time + +import torch +import torch.nn.functional as F +import torch.version +from tqdm import tqdm + +from .devices import device, get_optimal_device, test_for_nans, torch_gc + +sd_flag = False + + +def get_recommend_encoder_tile_size(): + if torch.cuda.is_available(): + total_memory = torch.cuda.get_device_properties( + device).total_memory // 2**20 + if total_memory > 16 * 1000: + ENCODER_TILE_SIZE = 3072 + elif total_memory > 12 * 1000: + ENCODER_TILE_SIZE = 2048 + elif total_memory > 8 * 1000: + ENCODER_TILE_SIZE = 1536 + else: + ENCODER_TILE_SIZE = 960 + else: + ENCODER_TILE_SIZE = 512 + return ENCODER_TILE_SIZE + + +def get_recommend_decoder_tile_size(): + if torch.cuda.is_available(): + total_memory = torch.cuda.get_device_properties( + device).total_memory // 2**20 + if total_memory > 30 * 1000: + DECODER_TILE_SIZE = 256 + elif total_memory > 16 * 1000: + DECODER_TILE_SIZE = 192 + elif total_memory > 12 * 1000: + DECODER_TILE_SIZE = 128 + elif total_memory > 8 * 1000: + DECODER_TILE_SIZE = 96 + else: + DECODER_TILE_SIZE = 64 + else: + DECODER_TILE_SIZE = 64 + return DECODER_TILE_SIZE + + +if 'global const': + DEFAULT_ENABLED = False + DEFAULT_MOVE_TO_GPU = False + DEFAULT_FAST_ENCODER = True + DEFAULT_FAST_DECODER = True + DEFAULT_COLOR_FIX = 0 + DEFAULT_ENCODER_TILE_SIZE = get_recommend_encoder_tile_size() + DEFAULT_DECODER_TILE_SIZE = get_recommend_decoder_tile_size() + + +# inplace version of silu +def inplace_nonlinearity(x): + # Test: fix for Nans + return F.silu(x, inplace=True) + + +# extracted from ldm.modules.diffusionmodules.model + + +# from diffusers lib +def attn_forward_new(self, h_): + batch_size, channel, height, width = h_.shape + hidden_states = h_.view(batch_size, channel, + height * width).transpose(1, 2) + + attention_mask = None + encoder_hidden_states = None + batch_size, sequence_length, _ = hidden_states.shape + attention_mask = self.prepare_attention_mask(attention_mask, + sequence_length, batch_size) + + query = self.to_q(hidden_states) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + elif self.norm_cross: + encoder_hidden_states = self.norm_encoder_hidden_states( + encoder_hidden_states) + + key = self.to_k(encoder_hidden_states) + value = self.to_v(encoder_hidden_states) + + query = self.head_to_batch_dim(query) + key = self.head_to_batch_dim(key) + value = self.head_to_batch_dim(value) + + attention_probs = self.get_attention_scores(query, key, attention_mask) + hidden_states = torch.bmm(attention_probs, value) + hidden_states = self.batch_to_head_dim(hidden_states) + + # linear proj + hidden_states = self.to_out[0](hidden_states) + # dropout + hidden_states = self.to_out[1](hidden_states) + + hidden_states = hidden_states.transpose(-1, -2).reshape( + batch_size, channel, height, width) + + return hidden_states + + +def attn2task(task_queue, net): + task_queue.append(('store_res', lambda x: x)) + task_queue.append(('pre_norm', net.group_norm)) + task_queue.append(('attn', lambda x, net=net: attn_forward_new(net, x))) + task_queue.append(['add_res', None]) + + +def resblock2task(queue, block): + """ + Turn a ResNetBlock into a sequence of tasks and append to the task queue + + @param queue: the target task queue + @param block: ResNetBlock + + """ + if block.in_channels != block.out_channels: + if sd_flag: + if block.use_conv_shortcut: + queue.append(('store_res', block.conv_shortcut)) + else: + queue.append(('store_res', block.nin_shortcut)) + else: + if block.use_in_shortcut: + queue.append(('store_res', block.conv_shortcut)) + else: + queue.append(('store_res', block.nin_shortcut)) + + else: + queue.append(('store_res', lambda x: x)) + queue.append(('pre_norm', block.norm1)) + queue.append(('silu', inplace_nonlinearity)) + queue.append(('conv1', block.conv1)) + queue.append(('pre_norm', block.norm2)) + queue.append(('silu', inplace_nonlinearity)) + queue.append(('conv2', block.conv2)) + queue.append(['add_res', None]) + + +def build_sampling(task_queue, net, is_decoder): + """ + Build the sampling part of a task queue + @param task_queue: the target task queue + @param net: the network + @param is_decoder: currently building decoder or encoder + """ + if is_decoder: + if sd_flag: + resblock2task(task_queue, net.mid.block_1) + attn2task(task_queue, net.mid.attn_1) + print(task_queue) + resblock2task(task_queue, net.mid.block_2) + resolution_iter = reversed(range(net.num_resolutions)) + block_ids = net.num_res_blocks + 1 + condition = 0 + module = net.up + func_name = 'upsample' + else: + resblock2task(task_queue, net.mid_block.resnets[0]) + attn2task(task_queue, net.mid_block.attentions[0]) + resblock2task(task_queue, net.mid_block.resnets[1]) + resolution_iter = (range(len(net.up_blocks)) + ) # net.num_resolutions = 3 + block_ids = 2 + 1 + condition = len(net.up_blocks) - 1 + module = net.up_blocks + func_name = 'upsamplers' + else: + resolution_iter = range(net.num_resolutions) + block_ids = net.num_res_blocks + condition = net.num_resolutions - 1 + module = net.down + func_name = 'downsample' + + for i_level in resolution_iter: + for i_block in range(block_ids): + if sd_flag: + resblock2task(task_queue, module[i_level].block[i_block]) + else: + resblock2task(task_queue, module[i_level].resnets[i_block]) + if i_level != condition: + if sd_flag: + task_queue.append( + (func_name, getattr(module[i_level], func_name))) + else: + task_queue.append((func_name, module[i_level].upsamplers[0])) + + if not is_decoder: + if sd_flag: + resblock2task(task_queue, net.mid.block_1) + attn2task(task_queue, net.mid.attn_1) + resblock2task(task_queue, net.mid.block_2) + else: + resblock2task(task_queue, net.mid_block.resnets[0]) + attn2task(task_queue, net.mid_block.attentions[0]) + resblock2task(task_queue, net.mid_block.resnets[1]) + + +def build_task_queue(net, is_decoder): + """ + Build a single task queue for the encoder or decoder + @param net: the VAE decoder or encoder network + @param is_decoder: currently building decoder or encoder + @return: the task queue + """ + task_queue = [] + task_queue.append(('conv_in', net.conv_in)) + + # construct the sampling part of the task queue + # because encoder and decoder share the same architecture, we extract the sampling part + build_sampling(task_queue, net, is_decoder) + if is_decoder and not sd_flag: + net.give_pre_end = False + net.tanh_out = False + + if not is_decoder or not net.give_pre_end: + if sd_flag: + task_queue.append(('pre_norm', net.norm_out)) + else: + task_queue.append(('pre_norm', net.conv_norm_out)) + task_queue.append(('silu', inplace_nonlinearity)) + task_queue.append(('conv_out', net.conv_out)) + if is_decoder and net.tanh_out: + task_queue.append(('tanh', torch.tanh)) + + return task_queue + + +def clone_task_queue(task_queue): + """ + Clone a task queue + @param task_queue: the task queue to be cloned + @return: the cloned task queue + """ + return [[item for item in task] for task in task_queue] + + +def get_var_mean(input, num_groups, eps=1e-6): + """ + Get mean and var for group norm + """ + b, c = input.size(0), input.size(1) + channel_in_group = int(c / num_groups) + input_reshaped = input.contiguous().view(1, int(b * num_groups), + channel_in_group, + *input.size()[2:]) + var, mean = torch.var_mean( + input_reshaped, dim=[0, 2, 3, 4], unbiased=False) + return var, mean + + +def custom_group_norm(input, + num_groups, + mean, + var, + weight=None, + bias=None, + eps=1e-6): + """ + Custom group norm with fixed mean and var + + @param input: input tensor + @param num_groups: number of groups. by default, num_groups = 32 + @param mean: mean, must be pre-calculated by get_var_mean + @param var: var, must be pre-calculated by get_var_mean + @param weight: weight, should be fetched from the original group norm + @param bias: bias, should be fetched from the original group norm + @param eps: epsilon, by default, eps = 1e-6 to match the original group norm + + @return: normalized tensor + """ + b, c = input.size(0), input.size(1) + channel_in_group = int(c / num_groups) + input_reshaped = input.contiguous().view(1, int(b * num_groups), + channel_in_group, + *input.size()[2:]) + + out = F.batch_norm( + input_reshaped, + mean, + var, + weight=None, + bias=None, + training=False, + momentum=0, + eps=eps) + + out = out.view(b, c, *input.size()[2:]) + + # post affine transform + if weight is not None: + out *= weight.view(1, -1, 1, 1) + if bias is not None: + out += bias.view(1, -1, 1, 1) + return out + + +def crop_valid_region(x, input_bbox, target_bbox, is_decoder): + """ + Crop the valid region from the tile + @param x: input tile + @param input_bbox: original input bounding box + @param target_bbox: output bounding box + @param scale: scale factor + @return: cropped tile + """ + padded_bbox = [i * 8 if is_decoder else i // 8 for i in input_bbox] + margin = [target_bbox[i] - padded_bbox[i] for i in range(4)] + return x[:, :, margin[2]:x.size(2) + margin[3], + margin[0]:x.size(3) + margin[1]] + + +# ↓↓↓ https://github.com/Kahsolt/stable-diffusion-webui-vae-tile-infer ↓↓↓ + + +def perfcount(fn): + + def wrapper(*args, **kwargs): + ts = time() + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats(device) + torch_gc() + gc.collect() + + ret = fn(*args, **kwargs) + + torch_gc() + gc.collect() + if torch.cuda.is_available(): + vram = torch.cuda.max_memory_allocated(device) / 2**20 + torch.cuda.reset_peak_memory_stats(device) + print( + f'[Tiled VAE]: Done in {time() - ts:.3f}s, max VRAM alloc {vram:.3f} MB' + ) + else: + print(f'[Tiled VAE]: Done in {time() - ts:.3f}s') + + return ret + + return wrapper + + +# copy end :) + + +class GroupNormParam: + + def __init__(self): + self.var_list = [] + self.mean_list = [] + self.pixel_list = [] + self.weight = None + self.bias = None + + def add_tile(self, tile, layer): + var, mean = get_var_mean(tile, 32) + # For giant images, the variance can be larger than max float16 + # In this case we create a copy to float32 + if var.dtype == torch.float16 and var.isinf().any(): + fp32_tile = tile.float() + var, mean = get_var_mean(fp32_tile, 32) + # ============= DEBUG: test for infinite ============= + # if torch.isinf(var).any(): + # print('var: ', var) + # ==================================================== + self.var_list.append(var) + self.mean_list.append(mean) + self.pixel_list.append(tile.shape[2] * tile.shape[3]) + if hasattr(layer, 'weight'): + self.weight = layer.weight + self.bias = layer.bias + else: + self.weight = None + self.bias = None + + def summary(self): + """ + summarize the mean and var and return a function + that apply group norm on each tile + """ + if len(self.var_list) == 0: + return None + var = torch.vstack(self.var_list) + mean = torch.vstack(self.mean_list) + max_value = max(self.pixel_list) + pixels = torch.tensor( + self.pixel_list, dtype=torch.float32, device=device) / max_value + sum_pixels = torch.sum(pixels) + pixels = pixels.unsqueeze(1) / sum_pixels + var = torch.sum(var * pixels, dim=0) + mean = torch.sum(mean * pixels, dim=0) + return lambda x: custom_group_norm(x, 32, mean, var, self.weight, self. + bias) + + @staticmethod + def from_tile(tile, norm): + """ + create a function from a single tile without summary + """ + var, mean = get_var_mean(tile, 32) + if var.dtype == torch.float16 and var.isinf().any(): + fp32_tile = tile.float() + var, mean = get_var_mean(fp32_tile, 32) + # if it is a macbook, we need to convert back to float16 + if var.device.type == 'mps': + # clamp to avoid overflow + var = torch.clamp(var, 0, 60000) + var = var.half() + mean = mean.half() + if hasattr(norm, 'weight'): + weight = norm.weight + bias = norm.bias + else: + weight = None + bias = None + + def group_norm_func(x, mean=mean, var=var, weight=weight, bias=bias): + return custom_group_norm(x, 32, mean, var, weight, bias, 1e-6) + + return group_norm_func + + +class VAEHook: + + def __init__(self, + net, + tile_size, + is_decoder, + fast_decoder, + fast_encoder, + color_fix, + to_gpu=False): + self.net = net # encoder | decoder + self.tile_size = tile_size + self.is_decoder = is_decoder + self.fast_mode = (fast_encoder and not is_decoder) or (fast_decoder + and is_decoder) + self.color_fix = color_fix and not is_decoder + self.to_gpu = to_gpu + self.pad = 11 if is_decoder else 32 + + def __call__(self, x): + B, C, H, W = x.shape + original_device = next(self.net.parameters()).device + try: + if self.to_gpu: + self.net.to(get_optimal_device()) + if max(H, W) <= self.pad * 2 + self.tile_size: + print( + '[Tiled VAE]: the input size is tiny and unnecessary to tile.' + ) + return self.net.original_forward(x) + else: + return self.vae_tile_forward(x) + finally: + self.net.to(original_device) + + def get_best_tile_size(self, lowerbound, upperbound): + """ + Get the best tile size for GPU memory + """ + divider = 32 + while divider >= 2: + remainer = lowerbound % divider + if remainer == 0: + return lowerbound + candidate = lowerbound - remainer + divider + if candidate <= upperbound: + return candidate + divider //= 2 + return lowerbound + + def split_tiles(self, h, w): + """ + Tool function to split the image into tiles + @param h: height of the image + @param w: width of the image + @return: tile_input_bboxes, tile_output_bboxes + """ + tile_input_bboxes, tile_output_bboxes = [], [] + tile_size = self.tile_size + pad = self.pad + num_height_tiles = math.ceil((h - 2 * pad) / tile_size) + num_width_tiles = math.ceil((w - 2 * pad) / tile_size) + # If any of the numbers are 0, we let it be 1 + # This is to deal with long and thin images + num_height_tiles = max(num_height_tiles, 1) + num_width_tiles = max(num_width_tiles, 1) + + # Suggestions from https://github.com/Kahsolt: auto shrink the tile size + real_tile_height = math.ceil((h - 2 * pad) / num_height_tiles) + real_tile_width = math.ceil((w - 2 * pad) / num_width_tiles) + real_tile_height = self.get_best_tile_size(real_tile_height, tile_size) + real_tile_width = self.get_best_tile_size(real_tile_width, tile_size) + + print( + f'[Tiled VAE]: split to {num_height_tiles}x{num_width_tiles}={num_height_tiles*num_width_tiles} tiles.', + f'Optimal tile size {real_tile_width}x{real_tile_height}, original tile size {tile_size}x{tile_size}' + ) + + for i in range(num_height_tiles): + for j in range(num_width_tiles): + # bbox: [x1, x2, y1, y2] + # the padding is is unnessary for image borders. So we directly start from (32, 32) + input_bbox = [ + pad + j * real_tile_width, + min(pad + (j + 1) * real_tile_width, w), + pad + i * real_tile_height, + min(pad + (i + 1) * real_tile_height, h), + ] + + # if the output bbox is close to the image boundary, we extend it to the image boundary + output_bbox = [ + input_bbox[0] if input_bbox[0] > pad else 0, + input_bbox[1] if input_bbox[1] < w - pad else w, + input_bbox[2] if input_bbox[2] > pad else 0, + input_bbox[3] if input_bbox[3] < h - pad else h, + ] + + # scale to get the final output bbox + output_bbox = [ + x * 8 if self.is_decoder else x // 8 for x in output_bbox + ] + tile_output_bboxes.append(output_bbox) + + # indistinguishable expand the input bbox by pad pixels + tile_input_bboxes.append([ + max(0, input_bbox[0] - pad), + min(w, input_bbox[1] + pad), + max(0, input_bbox[2] - pad), + min(h, input_bbox[3] + pad), + ]) + + return tile_input_bboxes, tile_output_bboxes + + @torch.no_grad() + def estimate_group_norm(self, z, task_queue, color_fix): + device = z.device + tile = z + last_id = len(task_queue) - 1 + while last_id >= 0 and task_queue[last_id][0] != 'pre_norm': + last_id -= 1 + if last_id <= 0 or task_queue[last_id][0] != 'pre_norm': + raise ValueError('No group norm found in the task queue') + # estimate until the last group norm + for i in range(last_id + 1): + task = task_queue[i] + if task[0] == 'pre_norm': + group_norm_func = GroupNormParam.from_tile(tile, task[1]) + task_queue[i] = ('apply_norm', group_norm_func) + if i == last_id: + return True + tile = group_norm_func(tile) + elif task[0] == 'store_res': + task_id = i + 1 + while task_id < last_id and task_queue[task_id][0] != 'add_res': + task_id += 1 + if task_id >= last_id: + continue + task_queue[task_id][1] = task[1](tile) + elif task[0] == 'add_res': + tile += task[1].to(device) + task[1] = None + elif color_fix and task[0] == 'downsample': + for j in range(i, last_id + 1): + if task_queue[j][0] == 'store_res': + task_queue[j] = ('store_res_cpu', task_queue[j][1]) + return True + else: + tile = task[1](tile) + try: + test_for_nans(tile, 'vae') + except Exception as e: + print( + f'{e}. Nan detected in fast mode estimation. Fast mode disabled.' + ) + return False + + raise IndexError('Should not reach here') + + @perfcount + @torch.no_grad() + def vae_tile_forward(self, z): + """ + Decode a latent vector z into an image in a tiled manner. + @param z: latent vector + @return: image + """ + device = next(self.net.parameters()).device + net = self.net + tile_size = self.tile_size + is_decoder = self.is_decoder + + z = z.detach() # detach the input to avoid backprop + + N, height, width = z.shape[0], z.shape[2], z.shape[3] + net.last_z_shape = z.shape + + # Split the input into tiles and build a task queue for each tile + print( + f'[Tiled VAE]: input_size: {z.shape}, tile_size: {tile_size}, padding: {self.pad}' + ) + + in_bboxes, out_bboxes = self.split_tiles(height, width) + + # Prepare tiles by split the input latents + tiles = [] + for input_bbox in in_bboxes: + tile = z[:, :, input_bbox[2]:input_bbox[3], + input_bbox[0]:input_bbox[1]].cpu() + tiles.append(tile) + + num_tiles = len(tiles) + num_completed = 0 + + # Build task queues + single_task_queue = build_task_queue(net, is_decoder) + if self.fast_mode: + # Fast mode: downsample the input image to the tile size, + # then estimate the group norm parameters on the downsampled image + scale_factor = tile_size / max(height, width) + z = z.to(device) + downsampled_z = F.interpolate( + z, scale_factor=scale_factor, mode='nearest-exact') + # use nearest-exact to keep statictics as close as possible + print( + f'[Tiled VAE]: Fast mode enabled, estimating group norm parameters on \ + {downsampled_z.shape[3]} x {downsampled_z.shape[2]} image') + + # ======= Special thanks to @Kahsolt for distribution shift issue ======= # + # The downsampling will heavily distort its mean and std, so we need to recover it. + std_old, mean_old = torch.std_mean(z, dim=[0, 2, 3], keepdim=True) + std_new, mean_new = torch.std_mean( + downsampled_z, dim=[0, 2, 3], keepdim=True) + downsampled_z = (downsampled_z + - mean_new) / std_new * std_old + mean_old + del std_old, mean_old, std_new, mean_new + # occasionally the std_new is too small or too large, which exceeds the range of float16 + # so we need to clamp it to max z's range. + downsampled_z = torch.clamp_( + downsampled_z, min=z.min(), max=z.max()) + estimate_task_queue = clone_task_queue(single_task_queue) + if self.estimate_group_norm( + downsampled_z, estimate_task_queue, + color_fix=self.color_fix): + single_task_queue = estimate_task_queue + del downsampled_z + + task_queues = [ + clone_task_queue(single_task_queue) for _ in range(num_tiles) + ] + + # Dummy result + result = None + result_approx = None + # Free memory of input latent tensor + del z + + # Task queue execution + pbar = tqdm( + total=num_tiles * len(task_queues[0]), + desc= + f"[Tiled VAE]: Executing {'Decoder' if is_decoder else 'Encoder'} Task Queue: " + ) + + # execute the task back and forth when switch tiles so that we always + # keep one tile on the GPU to reduce unnecessary data transfer + forward = True + interrupted = False + while True: + + group_norm_param = GroupNormParam() + for i in range(num_tiles) if forward else reversed( + range(num_tiles)): + + tile = tiles[i].to(device) + input_bbox = in_bboxes[i] + task_queue = task_queues[i] + + interrupted = False + while len(task_queue) > 0: + # DEBUG: current task + task = task_queue.pop(0) + if task[0] == 'pre_norm': + group_norm_param.add_tile(tile, task[1]) + break + elif task[0] == 'store_res' or task[0] == 'store_res_cpu': + task_id = 0 + res = task[1](tile) + if not self.fast_mode or task[0] == 'store_res_cpu': + res = res.cpu() + while task_queue[task_id][0] != 'add_res': + task_id += 1 + task_queue[task_id][1] = res + elif task[0] == 'add_res': + tile += task[1].to(device) + task[1] = None + else: + tile = task[1](tile) + pbar.update(1) + + if interrupted: + break + + # check for NaNs in the tile. + # If there are NaNs, we abort the process to save user's time + test_for_nans(tile, 'vae') + + if len(task_queue) == 0: + tiles[i] = None + num_completed += 1 + if result is None: # NOTE: dim C varies from different cases, can only be inited dynamically + result = torch.zeros( + (N, tile.shape[1], + height * 8 if is_decoder else height // 8, + width * 8 if is_decoder else width // 8), + device=device, + requires_grad=False) + result[:, :, out_bboxes[i][2]:out_bboxes[i][3], + out_bboxes[i][0]:out_bboxes[i] + [1]] = crop_valid_region(tile, in_bboxes[i], + out_bboxes[i], is_decoder) + del tile + elif i == num_tiles - 1 and forward: + forward = False + tiles[i] = tile + elif i == 0 and not forward: + forward = True + tiles[i] = tile + else: + tiles[i] = tile.cpu() + del tile + + if interrupted: + break + if num_completed == num_tiles: + break + + # insert the group norm task to the head of each task queue + group_norm_func = group_norm_param.summary() + if group_norm_func is not None: + for i in range(num_tiles): + task_queue = task_queues[i] + task_queue.insert(0, ('apply_norm', group_norm_func)) + + # Done! + pbar.close() + return result if result is not None else result_approx.to(device) diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index 5542bac0..68801e79 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -75,6 +75,7 @@ class CVTasks(object): # image editing skin_retouching = 'skin-retouching' image_super_resolution = 'image-super-resolution' + image_super_resolution_pasd = 'image-super-resolution-pasd' image_debanding = 'image-debanding' image_colorization = 'image-colorization' image_color_enhancement = 'image-color-enhancement' diff --git a/tests/pipelines/test_image_super_resolution_pasd.py b/tests/pipelines/test_image_super_resolution_pasd.py new file mode 100644 index 00000000..6547001a --- /dev/null +++ b/tests/pipelines/test_image_super_resolution_pasd.py @@ -0,0 +1,43 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os.path as osp +import unittest + +import cv2 + +from modelscope.outputs import OutputKeys +from modelscope.pipelines import pipeline +from modelscope.pipelines.base import Pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.test_utils import test_level + + +class ImageSuperResolutionPASDTest(unittest.TestCase): + + def setUp(self) -> None: + self.model_id = 'damo/PASD_image_super_resolutions' + self.img = 'data/test/images/dogs.jpg' + self.input = { + 'image': self.img, + 'prompt': '', + 'upscale': 1, + 'fidelity_scale_fg': 1.5, + 'fidelity_scale_bg': 0.7 + } + self.task = Tasks.image_super_resolution_pasd + + def pipeline_inference(self, pipeline: Pipeline, input: dict): + result = pipeline(input) + if result is not None: + cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG]) + print(f'Output written to {osp.abspath("result.png")}') + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_modelhub(self): + super_resolution = pipeline( + Tasks.image_super_resolution_pasd, model=self.model_id) + + self.pipeline_inference(super_resolution, self.input) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/run_config.yaml b/tests/run_config.yaml index 4c4d732f..d9768a57 100644 --- a/tests/run_config.yaml +++ b/tests/run_config.yaml @@ -7,6 +7,7 @@ isolated: # test cases that may require excessive anmount of GPU memory or run - test_dialog_modeling.py - test_csanmt_translation.py - test_image_super_resolution.py + - test_image_super_resolution_pasd.py - test_easycv_trainer.py - test_segformer.py - test_segmentation_pipeline.py From f05707ff33fd5e8056e8f723ea5e0093d713808c Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Mon, 11 Sep 2023 13:37:10 +0800 Subject: [PATCH 06/13] fix transformers compatible issue of position_ids Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13974608 * fix transformers compatible issue of position_ids --- docker/Dockerfile.ubuntu | 2 +- .../referring_video_object_segmentation/model.py | 4 ++++ modelscope/models/multi_modal/team/utils.py | 5 +++++ modelscope/models/nlp/space_T_en/text_to_sql.py | 5 +++++ .../models/nlp/unite/translation_evaluation.py | 4 ++++ .../nlp/use/user_satisfaction_estimation.py | 8 ++++++-- modelscope/utils/compatible_with_transformers.py | 16 ++++++++++++++++ 7 files changed, 41 insertions(+), 3 deletions(-) create mode 100644 modelscope/utils/compatible_with_transformers.py diff --git a/docker/Dockerfile.ubuntu b/docker/Dockerfile.ubuntu index 0c9c15c4..c37cb950 100644 --- a/docker/Dockerfile.ubuntu +++ b/docker/Dockerfile.ubuntu @@ -48,7 +48,7 @@ ENV SETUPTOOLS_USE_DISTUTILS=stdlib RUN CUDA_HOME=/usr/local/cuda TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6" pip install --no-cache-dir 'git+https://github.com/facebookresearch/detectron2.git' # torchmetrics==0.11.4 for ofa -RUN pip install --no-cache-dir tiktoken torchmetrics==0.11.4 'transformers<4.31.0' transformers_stream_generator 'protobuf<=3.20.0' bitsandbytes basicsr +RUN pip install --no-cache-dir tiktoken torchmetrics==0.11.4 transformers_stream_generator 'protobuf<=3.20.0' bitsandbytes basicsr COPY docker/scripts/install_flash_attension.sh /tmp/install_flash_attension.sh RUN if [ "$USE_GPU" = "True" ] ; then \ bash /tmp/install_flash_attension.sh; \ diff --git a/modelscope/models/cv/referring_video_object_segmentation/model.py b/modelscope/models/cv/referring_video_object_segmentation/model.py index 29e702be..239d0ea5 100644 --- a/modelscope/models/cv/referring_video_object_segmentation/model.py +++ b/modelscope/models/cv/referring_video_object_segmentation/model.py @@ -9,6 +9,8 @@ import torch from modelscope.metainfo import Models from modelscope.models.base.base_torch_model import TorchModel from modelscope.models.builder import MODELS +from modelscope.utils.compatible_with_transformers import \ + compatible_position_ids from modelscope.utils.config import Config from modelscope.utils.constant import ModelFile, Tasks from modelscope.utils.logger import get_logger @@ -40,6 +42,8 @@ class ReferringVideoObjectSegmentation(TorchModel): params_dict = torch.load(model_path, map_location='cpu') if 'model_state_dict' in params_dict.keys(): params_dict = params_dict['model_state_dict'] + compatible_position_ids( + params_dict, 'transformer.text_encoder.embeddings.position_ids') self.model.load_state_dict(params_dict, strict=True) self.set_postprocessor(self.cfg.pipeline.dataset_name) diff --git a/modelscope/models/multi_modal/team/utils.py b/modelscope/models/multi_modal/team/utils.py index 73919179..6ecc2172 100644 --- a/modelscope/models/multi_modal/team/utils.py +++ b/modelscope/models/multi_modal/team/utils.py @@ -15,6 +15,9 @@ import torch.utils.checkpoint as checkpoint from torch import nn from transformers import BertConfig, BertForMaskedLM +from modelscope.utils.compatible_with_transformers import \ + compatible_position_ids + class LayerNorm(nn.LayerNorm): """Subclass torch's LayerNorm to handle fp16.""" @@ -290,6 +293,8 @@ class TEAM(nn.Module): self.text_tensor_fc = nn.Linear(1024, 768) params = torch.load(pretrained, 'cpu') + compatible_position_ids(params, + 'text_model.bert.embeddings.position_ids') self.load_state_dict(params, strict=True) def get_feature(self, text_data=None, text_mask=None, img_tensor=None): diff --git a/modelscope/models/nlp/space_T_en/text_to_sql.py b/modelscope/models/nlp/space_T_en/text_to_sql.py index ca2d2596..d53003bc 100644 --- a/modelscope/models/nlp/space_T_en/text_to_sql.py +++ b/modelscope/models/nlp/space_T_en/text_to_sql.py @@ -12,6 +12,8 @@ from modelscope.metainfo import Models from modelscope.models import TorchModel from modelscope.models.base import Tensor from modelscope.models.builder import MODELS +from modelscope.utils.compatible_with_transformers import \ + compatible_position_ids from modelscope.utils.config import Config from modelscope.utils.constant import ModelFile, Tasks @@ -47,6 +49,9 @@ class StarForTextToSql(TorchModel): open( os.path.join(model_dir, ModelFile.TORCH_MODEL_BIN_FILE), 'rb'), map_location=self.device) + compatible_position_ids( + check_point['model'], + 'encoder.input_layer.plm_model.embeddings.position_ids') self.model.load_state_dict(check_point['model']) def forward(self, input: Dict[str, Tensor]) -> Dict[str, Tensor]: diff --git a/modelscope/models/nlp/unite/translation_evaluation.py b/modelscope/models/nlp/unite/translation_evaluation.py index c7e96027..6fa3264c 100644 --- a/modelscope/models/nlp/unite/translation_evaluation.py +++ b/modelscope/models/nlp/unite/translation_evaluation.py @@ -22,6 +22,8 @@ from modelscope.models.base import TorchModel from modelscope.models.builder import MODELS from modelscope.models.nlp.unite.configuration import InputFormat from modelscope.outputs.nlp_outputs import TranslationEvaluationOutput +from modelscope.utils.compatible_with_transformers import \ + compatible_position_ids from modelscope.utils.constant import Tasks from modelscope.utils.logger import get_logger @@ -305,6 +307,8 @@ class UniTEForTranslationEvaluation(TorchModel): self.encoder.pooler = None else: state_dict = torch.load(path, map_location=device) + compatible_position_ids(state_dict, + 'encoder.embeddings.position_ids') self.load_state_dict(state_dict) logger.info('Loading checkpoint parameters from %s' % path) return diff --git a/modelscope/models/nlp/use/user_satisfaction_estimation.py b/modelscope/models/nlp/use/user_satisfaction_estimation.py index 77b12414..cb136b81 100644 --- a/modelscope/models/nlp/use/user_satisfaction_estimation.py +++ b/modelscope/models/nlp/use/user_satisfaction_estimation.py @@ -15,6 +15,8 @@ from modelscope.models import TorchModel from modelscope.models.base import Tensor from modelscope.models.builder import MODELS from modelscope.outputs import DialogueUserSatisfactionEstimationModelOutput +from modelscope.utils.compatible_with_transformers import \ + compatible_position_ids from modelscope.utils.constant import ModelFile, Tasks from .transformer import TransformerEncoder @@ -45,8 +47,10 @@ class UserSatisfactionEstimation(TorchModel): self.device = device self.model = self.init_model() model_ckpt = os.path.join(model_dir, ModelFile.TORCH_MODEL_FILE) - self.model.load_state_dict( - torch.load(model_ckpt, map_location=torch.device('cpu'))) + stats_dict = torch.load(model_ckpt, map_location=torch.device('cpu')) + compatible_position_ids(stats_dict, + 'private.bert.embeddings.position_ids') + self.model.load_state_dict(stats_dict) def init_model(self): configs = { diff --git a/modelscope/utils/compatible_with_transformers.py b/modelscope/utils/compatible_with_transformers.py new file mode 100644 index 00000000..8b06c6c2 --- /dev/null +++ b/modelscope/utils/compatible_with_transformers.py @@ -0,0 +1,16 @@ +import transformers +from packaging import version + + +def compatible_position_ids(state_dict, position_id_key): + """Transformers no longer expect position_ids after transformers==4.31 + https://github.com/huggingface/transformers/pull/24505 + + Args: + position_id_key (str): position_ids key, + such as(encoder.embeddings.position_ids) + """ + transformer_version = version.parse('.'.join( + transformers.__version__.split('.')[:2])) + if transformer_version >= version.parse('4.31.0'): + del state_dict[position_id_key] From 6e38597d4a013f62bbbb03b6fd0f2b85a3fab743 Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Mon, 11 Sep 2023 13:58:35 +0800 Subject: [PATCH 07/13] modify version to 2099 --- modelscope/version.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modelscope/version.py b/modelscope/version.py index 46db9e93..e1c41d72 100644 --- a/modelscope/version.py +++ b/modelscope/version.py @@ -1,5 +1,5 @@ # Make sure to modify __release_datetime__ to release time when making official release. -__version__ = '1.9.0' +__version__ = '1.9.1' # default release datetime for branches under active development is set # to be a time far-far-away-into-the-future -__release_datetime__ = '2023-09-06 00:00:00' +__release_datetime__ = '2099-10-13 08:56:12' From e297672fa0dfd08225805fef4114f2b7dfebdcfd Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Mon, 11 Sep 2023 16:56:50 +0800 Subject: [PATCH 08/13] skip hf download --- tests/pipelines/test_text_generation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pipelines/test_text_generation.py b/tests/pipelines/test_text_generation.py index 091580dd..2197c4d7 100644 --- a/tests/pipelines/test_text_generation.py +++ b/tests/pipelines/test_text_generation.py @@ -172,7 +172,7 @@ class TextGenerationTest(unittest.TestCase): self.run_pipeline_with_model_id(self.gpt3_large_model_id, self.gpt3_input) - @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_hf_model_stream_generate(self): from transformers import AutoTokenizer, GPT2LMHeadModel tokenizer = AutoTokenizer.from_pretrained('gpt2') From b02f941a02a1dd501090e5625af7bd08b1dd8bc2 Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Fri, 15 Sep 2023 15:10:29 +0800 Subject: [PATCH 09/13] modify version number --- .dev_scripts/build_image.sh | 2 +- docker/Dockerfile.ubuntu | 2 +- modelscope/version.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.dev_scripts/build_image.sh b/.dev_scripts/build_image.sh index 596baeb9..9775d72e 100644 --- a/.dev_scripts/build_image.sh +++ b/.dev_scripts/build_image.sh @@ -150,7 +150,7 @@ echo -e "Building image with:\npython$python_version\npytorch$torch_version\nten docker_file_content=`cat docker/Dockerfile.ubuntu` if [ "$is_ci_test" != "True" ]; then echo "Building ModelScope lib, will install ModelScope lib to image" - docker_file_content="${docker_file_content} \nRUN pip install --no-cache-dir modelscope==$modelscope_version -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html" + docker_file_content="${docker_file_content} \nRUN pip install --no-cache-dir https://modelscope.oss-cn-beijing.aliyuncs.com/releases/build/modelscope-$modelscope_version-py3-none-any.whl " fi echo "$is_dsw" if [ "$is_dsw" == "False" ]; then diff --git a/docker/Dockerfile.ubuntu b/docker/Dockerfile.ubuntu index c37cb950..e26cfd33 100644 --- a/docker/Dockerfile.ubuntu +++ b/docker/Dockerfile.ubuntu @@ -48,7 +48,7 @@ ENV SETUPTOOLS_USE_DISTUTILS=stdlib RUN CUDA_HOME=/usr/local/cuda TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6" pip install --no-cache-dir 'git+https://github.com/facebookresearch/detectron2.git' # torchmetrics==0.11.4 for ofa -RUN pip install --no-cache-dir tiktoken torchmetrics==0.11.4 transformers_stream_generator 'protobuf<=3.20.0' bitsandbytes basicsr +RUN pip install --no-cache-dir tiktoken torchmetrics==0.11.4 https://modelscope.oss-cn-beijing.aliyuncs.com/releases/v/ms_swift-1.1.0-py3-none-any.whl transformers_stream_generator 'protobuf<=3.20.0' bitsandbytes basicsr COPY docker/scripts/install_flash_attension.sh /tmp/install_flash_attension.sh RUN if [ "$USE_GPU" = "True" ] ; then \ bash /tmp/install_flash_attension.sh; \ diff --git a/modelscope/version.py b/modelscope/version.py index e1c41d72..7f3298ab 100644 --- a/modelscope/version.py +++ b/modelscope/version.py @@ -2,4 +2,4 @@ __version__ = '1.9.1' # default release datetime for branches under active development is set # to be a time far-far-away-into-the-future -__release_datetime__ = '2099-10-13 08:56:12' +__release_datetime__ = '2023-09-06 00:00:00' From 0ed8fd51990b6dc6d6e49d6979cd2985b4c0e2e1 Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Fri, 15 Sep 2023 15:26:18 +0800 Subject: [PATCH 10/13] skip huggingface download cases --- tests/preprocessors/test_nlp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/preprocessors/test_nlp.py b/tests/preprocessors/test_nlp.py index 86e127c1..9bb9d10d 100644 --- a/tests/preprocessors/test_nlp.py +++ b/tests/preprocessors/test_nlp.py @@ -11,6 +11,7 @@ from modelscope.utils.logger import get_logger logger = get_logger() +@unittest.skip('skip for huggingface model download failed.') class NLPPreprocessorTest(unittest.TestCase): def setUp(self): From 40ddfd4499d641d80528359efd13a3fa78142d19 Mon Sep 17 00:00:00 2001 From: "yijing.wq" Date: Mon, 18 Sep 2023 21:22:56 +0800 Subject: [PATCH 11/13] fix maybe_allow_in_graph of diffusers update version Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/14077408 * fix_maybe_allow --- modelscope/models/cv/image_super_resolution_pasd/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelscope/models/cv/image_super_resolution_pasd/attention.py b/modelscope/models/cv/image_super_resolution_pasd/attention.py index 1ec566de..825a98a9 100644 --- a/modelscope/models/cv/image_super_resolution_pasd/attention.py +++ b/modelscope/models/cv/image_super_resolution_pasd/attention.py @@ -7,7 +7,7 @@ import torch import torch.nn.functional as F from diffusers.models.attention_processor import Attention from diffusers.models.embeddings import CombinedTimestepLabelEmbeddings -from diffusers.utils import maybe_allow_in_graph +from diffusers.utils.torch_utils import maybe_allow_in_graph from torch import nn From c02ff0151838102d58e75780ce661dde2cdf7b74 Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Tue, 19 Sep 2023 10:43:57 +0800 Subject: [PATCH 12/13] fix image_super_resolution diffusers compatible issue Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/14079826 * fix image_super_resolution diffusers compatible issue * fix lint issue --- .../pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modelscope/pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py b/modelscope/pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py index 7ebd3ce2..fee262b5 100644 --- a/modelscope/pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py +++ b/modelscope/pipelines/multi_modal/diffusers_wrapped/pasd_pipeline.py @@ -21,8 +21,9 @@ from diffusers.pipelines.stable_diffusion.safety_checker import \ StableDiffusionSafetyChecker from diffusers.schedulers import KarrasDiffusionSchedulers from diffusers.utils import (PIL_INTERPOLATION, is_accelerate_available, - is_accelerate_version, is_compiled_module, - logging, randn_tensor, replace_example_docstring) + is_accelerate_version, logging, + replace_example_docstring) +from diffusers.utils.torch_utils import is_compiled_module, randn_tensor from torchvision.utils import save_image from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer From 9de1b1e67428cf3478a76be6188587d54cf9470a Mon Sep 17 00:00:00 2001 From: "mulin.lyh" Date: Tue, 19 Sep 2023 11:42:59 +0800 Subject: [PATCH 13/13] release data to 2099 --- modelscope/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelscope/version.py b/modelscope/version.py index 7f3298ab..23ef0243 100644 --- a/modelscope/version.py +++ b/modelscope/version.py @@ -2,4 +2,4 @@ __version__ = '1.9.1' # default release datetime for branches under active development is set # to be a time far-far-away-into-the-future -__release_datetime__ = '2023-09-06 00:00:00' +__release_datetime__ = '2099-09-06 00:00:00'