EcoTry/IDM-VTON
0
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from typing import Any, Dict, Optional, Tuple, Union15 16import numpy as np17import torch18import torch.nn.functional as F19from torch import nn20 21from diffusers.utils import is_torch_version, logging22from diffusers.utils.torch_utils import apply_freeu23from diffusers.models.activations import get_activation24from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_025from diffusers.models.dual_transformer_2d import DualTransformer2DModel26from diffusers.models.normalization import AdaGroupNorm27from diffusers.models.resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D28from src.transformerhacked_garmnet import Transformer2DModel29from einops import rearrange30 31logger = logging.get_logger(__name__) # pylint: disable=invalid-name32 33 34def get_down_block(35 down_block_type: str,36 num_layers: int,37 in_channels: int,38 out_channels: int,39 temb_channels: int,40 add_downsample: bool,41 resnet_eps: float,42 resnet_act_fn: str,43 transformer_layers_per_block: int = 1,44 num_attention_heads: Optional[int] = None,45 resnet_groups: Optional[int] = None,46 cross_attention_dim: Optional[int] = None,47 downsample_padding: Optional[int] = None,48 dual_cross_attention: bool = False,49 use_linear_projection: bool = False,50 only_cross_attention: bool = False,51 upcast_attention: bool = False,52 resnet_time_scale_shift: str = "default",53 attention_type: str = "default",54 resnet_skip_time_act: bool = False,55 resnet_out_scale_factor: float = 1.0,56 cross_attention_norm: Optional[str] = None,57 attention_head_dim: Optional[int] = None,58 downsample_type: Optional[str] = None,59 dropout: float = 0.0,60):61 # If attn head dim is not defined, we default it to the number of heads62 if attention_head_dim is None:63 logger.warn(64 f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."65 )66 attention_head_dim = num_attention_heads67 68 down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type69 if down_block_type == "DownBlock2D":70 return DownBlock2D(71 num_layers=num_layers,72 in_channels=in_channels,73 out_channels=out_channels,74 temb_channels=temb_channels,75 dropout=dropout,76 add_downsample=add_downsample,77 resnet_eps=resnet_eps,78 resnet_act_fn=resnet_act_fn,79 resnet_groups=resnet_groups,80 downsample_padding=downsample_padding,81 resnet_time_scale_shift=resnet_time_scale_shift,82 )83 elif down_block_type == "ResnetDownsampleBlock2D":84 return ResnetDownsampleBlock2D(85 num_layers=num_layers,86 in_channels=in_channels,87 out_channels=out_channels,88 temb_channels=temb_channels,89 dropout=dropout,90 add_downsample=add_downsample,91 resnet_eps=resnet_eps,92 resnet_act_fn=resnet_act_fn,93 resnet_groups=resnet_groups,94 resnet_time_scale_shift=resnet_time_scale_shift,95 skip_time_act=resnet_skip_time_act,96 output_scale_factor=resnet_out_scale_factor,97 )98 elif down_block_type == "AttnDownBlock2D":99 if add_downsample is False:100 downsample_type = None101 else:102 downsample_type = downsample_type or "conv" # default to 'conv'103 return AttnDownBlock2D(104 num_layers=num_layers,105 in_channels=in_channels,106 out_channels=out_channels,107 temb_channels=temb_channels,108 dropout=dropout,109 resnet_eps=resnet_eps,110 resnet_act_fn=resnet_act_fn,111 resnet_groups=resnet_groups,112 downsample_padding=downsample_padding,113 attention_head_dim=attention_head_dim,114 resnet_time_scale_shift=resnet_time_scale_shift,115 downsample_type=downsample_type,116 )117 elif down_block_type == "CrossAttnDownBlock2D":118 if cross_attention_dim is None:119 raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D")120 return CrossAttnDownBlock2D(121 num_layers=num_layers,122 transformer_layers_per_block=transformer_layers_per_block,123 in_channels=in_channels,124 out_channels=out_channels,125 temb_channels=temb_channels,126 dropout=dropout,127 add_downsample=add_downsample,128 resnet_eps=resnet_eps,129 resnet_act_fn=resnet_act_fn,130 resnet_groups=resnet_groups,131 downsample_padding=downsample_padding,132 cross_attention_dim=cross_attention_dim,133 num_attention_heads=num_attention_heads,134 dual_cross_attention=dual_cross_attention,135 use_linear_projection=use_linear_projection,136 only_cross_attention=only_cross_attention,137 upcast_attention=upcast_attention,138 resnet_time_scale_shift=resnet_time_scale_shift,139 attention_type=attention_type,140 )141 elif down_block_type == "SimpleCrossAttnDownBlock2D":142 if cross_attention_dim is None:143 raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D")144 return SimpleCrossAttnDownBlock2D(145 num_layers=num_layers,146 in_channels=in_channels,147 out_channels=out_channels,148 temb_channels=temb_channels,149 dropout=dropout,150 add_downsample=add_downsample,151 resnet_eps=resnet_eps,152 resnet_act_fn=resnet_act_fn,153 resnet_groups=resnet_groups,154 cross_attention_dim=cross_attention_dim,155 attention_head_dim=attention_head_dim,156 resnet_time_scale_shift=resnet_time_scale_shift,157 skip_time_act=resnet_skip_time_act,158 output_scale_factor=resnet_out_scale_factor,159 only_cross_attention=only_cross_attention,160 cross_attention_norm=cross_attention_norm,161 )162 elif down_block_type == "SkipDownBlock2D":163 return SkipDownBlock2D(164 num_layers=num_layers,165 in_channels=in_channels,166 out_channels=out_channels,167 temb_channels=temb_channels,168 dropout=dropout,169 add_downsample=add_downsample,170 resnet_eps=resnet_eps,171 resnet_act_fn=resnet_act_fn,172 downsample_padding=downsample_padding,173 resnet_time_scale_shift=resnet_time_scale_shift,174 )175 elif down_block_type == "AttnSkipDownBlock2D":176 return AttnSkipDownBlock2D(177 num_layers=num_layers,178 in_channels=in_channels,179 out_channels=out_channels,180 temb_channels=temb_channels,181 dropout=dropout,182 add_downsample=add_downsample,183 resnet_eps=resnet_eps,184 resnet_act_fn=resnet_act_fn,185 attention_head_dim=attention_head_dim,186 resnet_time_scale_shift=resnet_time_scale_shift,187 )188 elif down_block_type == "DownEncoderBlock2D":189 return DownEncoderBlock2D(190 num_layers=num_layers,191 in_channels=in_channels,192 out_channels=out_channels,193 dropout=dropout,194 add_downsample=add_downsample,195 resnet_eps=resnet_eps,196 resnet_act_fn=resnet_act_fn,197 resnet_groups=resnet_groups,198 downsample_padding=downsample_padding,199 resnet_time_scale_shift=resnet_time_scale_shift,200 )201 elif down_block_type == "AttnDownEncoderBlock2D":202 return AttnDownEncoderBlock2D(203 num_layers=num_layers,204 in_channels=in_channels,205 out_channels=out_channels,206 dropout=dropout,207 add_downsample=add_downsample,208 resnet_eps=resnet_eps,209 resnet_act_fn=resnet_act_fn,210 resnet_groups=resnet_groups,211 downsample_padding=downsample_padding,212 attention_head_dim=attention_head_dim,213 resnet_time_scale_shift=resnet_time_scale_shift,214 )215 elif down_block_type == "KDownBlock2D":216 return KDownBlock2D(217 num_layers=num_layers,218 in_channels=in_channels,219 out_channels=out_channels,220 temb_channels=temb_channels,221 dropout=dropout,222 add_downsample=add_downsample,223 resnet_eps=resnet_eps,224 resnet_act_fn=resnet_act_fn,225 )226 elif down_block_type == "KCrossAttnDownBlock2D":227 return KCrossAttnDownBlock2D(228 num_layers=num_layers,229 in_channels=in_channels,230 out_channels=out_channels,231 temb_channels=temb_channels,232 dropout=dropout,233 add_downsample=add_downsample,234 resnet_eps=resnet_eps,235 resnet_act_fn=resnet_act_fn,236 cross_attention_dim=cross_attention_dim,237 attention_head_dim=attention_head_dim,238 add_self_attention=True if not add_downsample else False,239 )240 raise ValueError(f"{down_block_type} does not exist.")241 242 243def get_up_block(244 up_block_type: str,245 num_layers: int,246 in_channels: int,247 out_channels: int,248 prev_output_channel: int,249 temb_channels: int,250 add_upsample: bool,251 resnet_eps: float,252 resnet_act_fn: str,253 resolution_idx: Optional[int] = None,254 transformer_layers_per_block: int = 1,255 num_attention_heads: Optional[int] = None,256 resnet_groups: Optional[int] = None,257 cross_attention_dim: Optional[int] = None,258 dual_cross_attention: bool = False,259 use_linear_projection: bool = False,260 only_cross_attention: bool = False,261 upcast_attention: bool = False,262 resnet_time_scale_shift: str = "default",263 attention_type: str = "default",264 resnet_skip_time_act: bool = False,265 resnet_out_scale_factor: float = 1.0,266 cross_attention_norm: Optional[str] = None,267 attention_head_dim: Optional[int] = None,268 upsample_type: Optional[str] = None,269 dropout: float = 0.0,270) -> nn.Module:271 # If attn head dim is not defined, we default it to the number of heads272 if attention_head_dim is None:273 logger.warn(274 f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."275 )276 attention_head_dim = num_attention_heads277 278 up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type279 if up_block_type == "UpBlock2D":280 return UpBlock2D(281 num_layers=num_layers,282 in_channels=in_channels,283 out_channels=out_channels,284 prev_output_channel=prev_output_channel,285 temb_channels=temb_channels,286 resolution_idx=resolution_idx,287 dropout=dropout,288 add_upsample=add_upsample,289 resnet_eps=resnet_eps,290 resnet_act_fn=resnet_act_fn,291 resnet_groups=resnet_groups,292 resnet_time_scale_shift=resnet_time_scale_shift,293 )294 elif up_block_type == "ResnetUpsampleBlock2D":295 return ResnetUpsampleBlock2D(296 num_layers=num_layers,297 in_channels=in_channels,298 out_channels=out_channels,299 prev_output_channel=prev_output_channel,300 temb_channels=temb_channels,301 resolution_idx=resolution_idx,302 dropout=dropout,303 add_upsample=add_upsample,304 resnet_eps=resnet_eps,305 resnet_act_fn=resnet_act_fn,306 resnet_groups=resnet_groups,307 resnet_time_scale_shift=resnet_time_scale_shift,308 skip_time_act=resnet_skip_time_act,309 output_scale_factor=resnet_out_scale_factor,310 )311 elif up_block_type == "CrossAttnUpBlock2D":312 if cross_attention_dim is None:313 raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")314 return CrossAttnUpBlock2D(315 num_layers=num_layers,316 transformer_layers_per_block=transformer_layers_per_block,317 in_channels=in_channels,318 out_channels=out_channels,319 prev_output_channel=prev_output_channel,320 temb_channels=temb_channels,321 resolution_idx=resolution_idx,322 dropout=dropout,323 add_upsample=add_upsample,324 resnet_eps=resnet_eps,325 resnet_act_fn=resnet_act_fn,326 resnet_groups=resnet_groups,327 cross_attention_dim=cross_attention_dim,328 num_attention_heads=num_attention_heads,329 dual_cross_attention=dual_cross_attention,330 use_linear_projection=use_linear_projection,331 only_cross_attention=only_cross_attention,332 upcast_attention=upcast_attention,333 resnet_time_scale_shift=resnet_time_scale_shift,334 attention_type=attention_type,335 )336 elif up_block_type == "SimpleCrossAttnUpBlock2D":337 if cross_attention_dim is None:338 raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D")339 return SimpleCrossAttnUpBlock2D(340 num_layers=num_layers,341 in_channels=in_channels,342 out_channels=out_channels,343 prev_output_channel=prev_output_channel,344 temb_channels=temb_channels,345 resolution_idx=resolution_idx,346 dropout=dropout,347 add_upsample=add_upsample,348 resnet_eps=resnet_eps,349 resnet_act_fn=resnet_act_fn,350 resnet_groups=resnet_groups,351 cross_attention_dim=cross_attention_dim,352 attention_head_dim=attention_head_dim,353 resnet_time_scale_shift=resnet_time_scale_shift,354 skip_time_act=resnet_skip_time_act,355 output_scale_factor=resnet_out_scale_factor,356 only_cross_attention=only_cross_attention,357 cross_attention_norm=cross_attention_norm,358 )359 elif up_block_type == "AttnUpBlock2D":360 if add_upsample is False:361 upsample_type = None362 else:363 upsample_type = upsample_type or "conv" # default to 'conv'364 365 return AttnUpBlock2D(366 num_layers=num_layers,367 in_channels=in_channels,368 out_channels=out_channels,369 prev_output_channel=prev_output_channel,370 temb_channels=temb_channels,371 resolution_idx=resolution_idx,372 dropout=dropout,373 resnet_eps=resnet_eps,374 resnet_act_fn=resnet_act_fn,375 resnet_groups=resnet_groups,376 attention_head_dim=attention_head_dim,377 resnet_time_scale_shift=resnet_time_scale_shift,378 upsample_type=upsample_type,379 )380 elif up_block_type == "SkipUpBlock2D":381 return SkipUpBlock2D(382 num_layers=num_layers,383 in_channels=in_channels,384 out_channels=out_channels,385 prev_output_channel=prev_output_channel,386 temb_channels=temb_channels,387 resolution_idx=resolution_idx,388 dropout=dropout,389 add_upsample=add_upsample,390 resnet_eps=resnet_eps,391 resnet_act_fn=resnet_act_fn,392 resnet_time_scale_shift=resnet_time_scale_shift,393 )394 elif up_block_type == "AttnSkipUpBlock2D":395 return AttnSkipUpBlock2D(396 num_layers=num_layers,397 in_channels=in_channels,398 out_channels=out_channels,399 prev_output_channel=prev_output_channel,400 temb_channels=temb_channels,401 resolution_idx=resolution_idx,402 dropout=dropout,403 add_upsample=add_upsample,404 resnet_eps=resnet_eps,405 resnet_act_fn=resnet_act_fn,406 attention_head_dim=attention_head_dim,407 resnet_time_scale_shift=resnet_time_scale_shift,408 )409 elif up_block_type == "UpDecoderBlock2D":410 return UpDecoderBlock2D(411 num_layers=num_layers,412 in_channels=in_channels,413 out_channels=out_channels,414 resolution_idx=resolution_idx,415 dropout=dropout,416 add_upsample=add_upsample,417 resnet_eps=resnet_eps,418 resnet_act_fn=resnet_act_fn,419 resnet_groups=resnet_groups,420 resnet_time_scale_shift=resnet_time_scale_shift,421 temb_channels=temb_channels,422 )423 elif up_block_type == "AttnUpDecoderBlock2D":424 return AttnUpDecoderBlock2D(425 num_layers=num_layers,426 in_channels=in_channels,427 out_channels=out_channels,428 resolution_idx=resolution_idx,429 dropout=dropout,430 add_upsample=add_upsample,431 resnet_eps=resnet_eps,432 resnet_act_fn=resnet_act_fn,433 resnet_groups=resnet_groups,434 attention_head_dim=attention_head_dim,435 resnet_time_scale_shift=resnet_time_scale_shift,436 temb_channels=temb_channels,437 )438 elif up_block_type == "KUpBlock2D":439 return KUpBlock2D(440 num_layers=num_layers,441 in_channels=in_channels,442 out_channels=out_channels,443 temb_channels=temb_channels,444 resolution_idx=resolution_idx,445 dropout=dropout,446 add_upsample=add_upsample,447 resnet_eps=resnet_eps,448 resnet_act_fn=resnet_act_fn,449 )450 elif up_block_type == "KCrossAttnUpBlock2D":451 return KCrossAttnUpBlock2D(452 num_layers=num_layers,453 in_channels=in_channels,454 out_channels=out_channels,455 temb_channels=temb_channels,456 resolution_idx=resolution_idx,457 dropout=dropout,458 add_upsample=add_upsample,459 resnet_eps=resnet_eps,460 resnet_act_fn=resnet_act_fn,461 cross_attention_dim=cross_attention_dim,462 attention_head_dim=attention_head_dim,463 )464 465 raise ValueError(f"{up_block_type} does not exist.")466 467 468class AutoencoderTinyBlock(nn.Module):469 """470 Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU471 blocks.472 473 Args:474 in_channels (`int`): The number of input channels.475 out_channels (`int`): The number of output channels.476 act_fn (`str`):477 ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`.478 479 Returns:480 `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to481 `out_channels`.482 """483 484 def __init__(self, in_channels: int, out_channels: int, act_fn: str):485 super().__init__()486 act_fn = get_activation(act_fn)487 self.conv = nn.Sequential(488 nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),489 act_fn,490 nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),491 act_fn,492 nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),493 )494 self.skip = (495 nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)496 if in_channels != out_channels497 else nn.Identity()498 )499 self.fuse = nn.ReLU()500 501 def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:502 return self.fuse(self.conv(x) + self.skip(x))503 504 505class UNetMidBlock2D(nn.Module):506 """507 A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks.508 509 Args:510 in_channels (`int`): The number of input channels.511 temb_channels (`int`): The number of temporal embedding channels.512 dropout (`float`, *optional*, defaults to 0.0): The dropout rate.513 num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.514 resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.515 resnet_time_scale_shift (`str`, *optional*, defaults to `default`):516 The type of normalization to apply to the time embeddings. This can help to improve the performance of the517 model on tasks with long-range temporal dependencies.518 resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.519 resnet_groups (`int`, *optional*, defaults to 32):520 The number of groups to use in the group normalization layers of the resnet blocks.521 attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.522 resnet_pre_norm (`bool`, *optional*, defaults to `True`):523 Whether to use pre-normalization for the resnet blocks.524 add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.525 attention_head_dim (`int`, *optional*, defaults to 1):526 Dimension of a single attention head. The number of attention heads is determined based on this value and527 the number of input channels.528 output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor.529 530 Returns:531 `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,532 in_channels, height, width)`.533 534 """535 536 def __init__(537 self,538 in_channels: int,539 temb_channels: int,540 dropout: float = 0.0,541 num_layers: int = 1,542 resnet_eps: float = 1e-6,543 resnet_time_scale_shift: str = "default", # default, spatial544 resnet_act_fn: str = "swish",545 resnet_groups: int = 32,546 attn_groups: Optional[int] = None,547 resnet_pre_norm: bool = True,548 add_attention: bool = True,549 attention_head_dim: int = 1,550 output_scale_factor: float = 1.0,551 ):552 super().__init__()553 resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)554 self.add_attention = add_attention555 556 if attn_groups is None:557 attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None558 559 # there is always at least one resnet560 resnets = [561 ResnetBlock2D(562 in_channels=in_channels,563 out_channels=in_channels,564 temb_channels=temb_channels,565 eps=resnet_eps,566 groups=resnet_groups,567 dropout=dropout,568 time_embedding_norm=resnet_time_scale_shift,569 non_linearity=resnet_act_fn,570 output_scale_factor=output_scale_factor,571 pre_norm=resnet_pre_norm,572 )573 ]574 attentions = []575 576 if attention_head_dim is None:577 logger.warn(578 f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."579 )580 attention_head_dim = in_channels581 582 for _ in range(num_layers):583 if self.add_attention:584 attentions.append(585 Attention(586 in_channels,587 heads=in_channels // attention_head_dim,588 dim_head=attention_head_dim,589 rescale_output_factor=output_scale_factor,590 eps=resnet_eps,591 norm_num_groups=attn_groups,592 spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,593 residual_connection=True,594 bias=True,595 upcast_softmax=True,596 _from_deprecated_attn_block=True,597 )598 )599 else:600 attentions.append(None)601 602 resnets.append(603 ResnetBlock2D(604 in_channels=in_channels,605 out_channels=in_channels,606 temb_channels=temb_channels,607 eps=resnet_eps,608 groups=resnet_groups,609 dropout=dropout,610 time_embedding_norm=resnet_time_scale_shift,611 non_linearity=resnet_act_fn,612 output_scale_factor=output_scale_factor,613 pre_norm=resnet_pre_norm,614 )615 )616 617 self.attentions = nn.ModuleList(attentions)618 self.resnets = nn.ModuleList(resnets)619 620 def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:621 hidden_states = self.resnets[0](hidden_states, temb)622 for attn, resnet in zip(self.attentions, self.resnets[1:]):623 if attn is not None:624 hidden_states = attn(hidden_states, temb=temb)625 hidden_states = resnet(hidden_states, temb)626 627 return hidden_states628 629 630class UNetMidBlock2DCrossAttn(nn.Module):631 def __init__(632 self,633 in_channels: int,634 temb_channels: int,635 dropout: float = 0.0,636 num_layers: int = 1,637 transformer_layers_per_block: Union[int, Tuple[int]] = 1,638 resnet_eps: float = 1e-6,639 resnet_time_scale_shift: str = "default",640 resnet_act_fn: str = "swish",641 resnet_groups: int = 32,642 resnet_pre_norm: bool = True,643 num_attention_heads: int = 1,644 output_scale_factor: float = 1.0,645 cross_attention_dim: int = 1280,646 dual_cross_attention: bool = False,647 use_linear_projection: bool = False,648 upcast_attention: bool = False,649 attention_type: str = "default",650 ):651 super().__init__()652 653 self.has_cross_attention = True654 self.num_attention_heads = num_attention_heads655 resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)656 657 # support for variable transformer layers per block658 if isinstance(transformer_layers_per_block, int):659 transformer_layers_per_block = [transformer_layers_per_block] * num_layers660 661 # there is always at least one resnet662 resnets = [663 ResnetBlock2D(664 in_channels=in_channels,665 out_channels=in_channels,666 temb_channels=temb_channels,667 eps=resnet_eps,668 groups=resnet_groups,669 dropout=dropout,670 time_embedding_norm=resnet_time_scale_shift,671 non_linearity=resnet_act_fn,672 output_scale_factor=output_scale_factor,673 pre_norm=resnet_pre_norm,674 )675 ]676 attentions = []677 678 for i in range(num_layers):679 if not dual_cross_attention:680 attentions.append(681 Transformer2DModel(682 num_attention_heads,683 in_channels // num_attention_heads,684 in_channels=in_channels,685 num_layers=transformer_layers_per_block[i],686 cross_attention_dim=cross_attention_dim,687 norm_num_groups=resnet_groups,688 use_linear_projection=use_linear_projection,689 upcast_attention=upcast_attention,690 attention_type=attention_type,691 )692 )693 else:694 attentions.append(695 DualTransformer2DModel(696 num_attention_heads,697 in_channels // num_attention_heads,698 in_channels=in_channels,699 num_layers=1,700 cross_attention_dim=cross_attention_dim,701 norm_num_groups=resnet_groups,702 )703 )704 resnets.append(705 ResnetBlock2D(706 in_channels=in_channels,707 out_channels=in_channels,708 temb_channels=temb_channels,709 eps=resnet_eps,710 groups=resnet_groups,711 dropout=dropout,712 time_embedding_norm=resnet_time_scale_shift,713 non_linearity=resnet_act_fn,714 output_scale_factor=output_scale_factor,715 pre_norm=resnet_pre_norm,716 )717 )718 719 self.attentions = nn.ModuleList(attentions)720 self.resnets = nn.ModuleList(resnets)721 722 self.gradient_checkpointing = False723 724 def forward(725 self,726 hidden_states: torch.FloatTensor,727 temb: Optional[torch.FloatTensor] = None,728 encoder_hidden_states: Optional[torch.FloatTensor] = None,729 attention_mask: Optional[torch.FloatTensor] = None,730 cross_attention_kwargs: Optional[Dict[str, Any]] = None,731 encoder_attention_mask: Optional[torch.FloatTensor] = None,732 ) -> torch.FloatTensor:733 lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0734 hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)735 garment_features = []736 for attn, resnet in zip(self.attentions, self.resnets[1:]):737 if self.training and self.gradient_checkpointing:738 739 def create_custom_forward(module, return_dict=None):740 def custom_forward(*inputs):741 if return_dict is not None:742 return module(*inputs, return_dict=return_dict)743 else:744 return module(*inputs)745 746 return custom_forward747 748 ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}749 # hidden_states = attn(750 hidden_states,out_garment_feat = attn(751 hidden_states,752 encoder_hidden_states=encoder_hidden_states,753 cross_attention_kwargs=cross_attention_kwargs,754 attention_mask=attention_mask,755 encoder_attention_mask=encoder_attention_mask,756 return_dict=False,757 )758 hidden_states=hidden_states[0]759 hidden_states = torch.utils.checkpoint.checkpoint(760 create_custom_forward(resnet),761 hidden_states,762 temb,763 **ckpt_kwargs,764 )765 else:766 # hidden_states= attn(767 hidden_states,out_garment_feat = attn(768 hidden_states,769 encoder_hidden_states=encoder_hidden_states,770 cross_attention_kwargs=cross_attention_kwargs,771 attention_mask=attention_mask,772 encoder_attention_mask=encoder_attention_mask,773 return_dict=False,774 )775 hidden_states=hidden_states[0]776 hidden_states = resnet(hidden_states, temb, scale=lora_scale)777 garment_features += out_garment_feat778 return hidden_states,garment_features779 # return hidden_states 780 781 782class UNetMidBlock2DSimpleCrossAttn(nn.Module):783 def __init__(784 self,785 in_channels: int,786 temb_channels: int,787 dropout: float = 0.0,788 num_layers: int = 1,789 resnet_eps: float = 1e-6,790 resnet_time_scale_shift: str = "default",791 resnet_act_fn: str = "swish",792 resnet_groups: int = 32,793 resnet_pre_norm: bool = True,794 attention_head_dim: int = 1,795 output_scale_factor: float = 1.0,796 cross_attention_dim: int = 1280,797 skip_time_act: bool = False,798 only_cross_attention: bool = False,799 cross_attention_norm: Optional[str] = None,800 ):801 super().__init__()802 803 self.has_cross_attention = True804 805 self.attention_head_dim = attention_head_dim806 resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)807 808 self.num_heads = in_channels // self.attention_head_dim809 810 # there is always at least one resnet811 resnets = [812 ResnetBlock2D(813 in_channels=in_channels,814 out_channels=in_channels,815 temb_channels=temb_channels,816 eps=resnet_eps,817 groups=resnet_groups,818 dropout=dropout,819 time_embedding_norm=resnet_time_scale_shift,820 non_linearity=resnet_act_fn,821 output_scale_factor=output_scale_factor,822 pre_norm=resnet_pre_norm,823 skip_time_act=skip_time_act,824 )825 ]826 attentions = []827 828 for _ in range(num_layers):829 processor = (830 AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()831 )832 833 attentions.append(834 Attention(835 query_dim=in_channels,836 cross_attention_dim=in_channels,837 heads=self.num_heads,838 dim_head=self.attention_head_dim,839 added_kv_proj_dim=cross_attention_dim,840 norm_num_groups=resnet_groups,841 bias=True,842 upcast_softmax=True,843 only_cross_attention=only_cross_attention,844 cross_attention_norm=cross_attention_norm,845 processor=processor,846 )847 )848 resnets.append(849 ResnetBlock2D(850 in_channels=in_channels,851 out_channels=in_channels,852 temb_channels=temb_channels,853 eps=resnet_eps,854 groups=resnet_groups,855 dropout=dropout,856 time_embedding_norm=resnet_time_scale_shift,857 non_linearity=resnet_act_fn,858 output_scale_factor=output_scale_factor,859 pre_norm=resnet_pre_norm,860 skip_time_act=skip_time_act,861 )862 )863 864 self.attentions = nn.ModuleList(attentions)865 self.resnets = nn.ModuleList(resnets)866 867 def forward(868 self,869 hidden_states: torch.FloatTensor,870 temb: Optional[torch.FloatTensor] = None,871 encoder_hidden_states: Optional[torch.FloatTensor] = None,872 attention_mask: Optional[torch.FloatTensor] = None,873 cross_attention_kwargs: Optional[Dict[str, Any]] = None,874 encoder_attention_mask: Optional[torch.FloatTensor] = None,875 ) -> torch.FloatTensor:876 cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}877 lora_scale = cross_attention_kwargs.get("scale", 1.0)878 879 if attention_mask is None:880 # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask.881 mask = None if encoder_hidden_states is None else encoder_attention_mask882 else:883 # when attention_mask is defined: we don't even check for encoder_attention_mask.884 # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks.885 # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask.886 # then we can simplify this whole if/else block to:887 # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask888 mask = attention_mask889 890 hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)891 for attn, resnet in zip(self.attentions, self.resnets[1:]):892 # attn893 hidden_states = attn(894 hidden_states,895 encoder_hidden_states=encoder_hidden_states,896 attention_mask=mask,897 **cross_attention_kwargs,898 )899 900 # resnet901 hidden_states = resnet(hidden_states, temb, scale=lora_scale)902 903 return hidden_states904 905 906class AttnDownBlock2D(nn.Module):907 def __init__(908 self,909 in_channels: int,910 out_channels: int,911 temb_channels: int,912 dropout: float = 0.0,913 num_layers: int = 1,914 resnet_eps: float = 1e-6,915 resnet_time_scale_shift: str = "default",916 resnet_act_fn: str = "swish",917 resnet_groups: int = 32,918 resnet_pre_norm: bool = True,919 attention_head_dim: int = 1,920 output_scale_factor: float = 1.0,921 downsample_padding: int = 1,922 downsample_type: str = "conv",923 ):924 super().__init__()925 resnets = []926 attentions = []927 self.downsample_type = downsample_type928 929 if attention_head_dim is None:930 logger.warn(931 f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."932 )933 attention_head_dim = out_channels934 935 for i in range(num_layers):936 in_channels = in_channels if i == 0 else out_channels937 resnets.append(938 ResnetBlock2D(939 in_channels=in_channels,940 out_channels=out_channels,941 temb_channels=temb_channels,942 eps=resnet_eps,943 groups=resnet_groups,944 dropout=dropout,945 time_embedding_norm=resnet_time_scale_shift,946 non_linearity=resnet_act_fn,947 output_scale_factor=output_scale_factor,948 pre_norm=resnet_pre_norm,949 )950 )951 attentions.append(952 Attention(953 out_channels,954 heads=out_channels // attention_head_dim,955 dim_head=attention_head_dim,956 rescale_output_factor=output_scale_factor,957 eps=resnet_eps,958 norm_num_groups=resnet_groups,959 residual_connection=True,960 bias=True,961 upcast_softmax=True,962 _from_deprecated_attn_block=True,963 )964 )965 966 self.attentions = nn.ModuleList(attentions)967 self.resnets = nn.ModuleList(resnets)968 969 if downsample_type == "conv":970 self.downsamplers = nn.ModuleList(971 [972 Downsample2D(973 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"974 )975 ]976 )977 elif downsample_type == "resnet":978 self.downsamplers = nn.ModuleList(979 [980 ResnetBlock2D(981 in_channels=out_channels,982 out_channels=out_channels,983 temb_channels=temb_channels,984 eps=resnet_eps,985 groups=resnet_groups,986 dropout=dropout,987 time_embedding_norm=resnet_time_scale_shift,988 non_linearity=resnet_act_fn,989 output_scale_factor=output_scale_factor,990 pre_norm=resnet_pre_norm,991 down=True,992 )993 ]994 )995 else:996 self.downsamplers = None997 998 def forward(999 self,1000 hidden_states: torch.FloatTensor,1001 temb: Optional[torch.FloatTensor] = None,1002 upsample_size: Optional[int] = None,1003 cross_attention_kwargs: Optional[Dict[str, Any]] = None,1004 ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:1005 cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}1006 1007 lora_scale = cross_attention_kwargs.get("scale", 1.0)1008 1009 output_states = ()1010 1011 for resnet, attn in zip(self.resnets, self.attentions):1012 cross_attention_kwargs.update({"scale": lora_scale})1013 hidden_states = resnet(hidden_states, temb, scale=lora_scale)1014 hidden_states = attn(hidden_states, **cross_attention_kwargs)1015 output_states = output_states + (hidden_states,)1016 1017 if self.downsamplers is not None:1018 for downsampler in self.downsamplers:1019 if self.downsample_type == "resnet":1020 hidden_states = downsampler(hidden_states, temb=temb, scale=lora_scale)1021 else:1022 hidden_states = downsampler(hidden_states, scale=lora_scale)1023 1024 output_states += (hidden_states,)1025 1026 return hidden_states, output_states1027 1028 1029class CrossAttnDownBlock2D(nn.Module):1030 def __init__(1031 self,1032 in_channels: int,1033 out_channels: int,1034 temb_channels: int,1035 dropout: float = 0.0,1036 num_layers: int = 1,1037 transformer_layers_per_block: Union[int, Tuple[int]] = 1,1038 resnet_eps: float = 1e-6,1039 resnet_time_scale_shift: str = "default",1040 resnet_act_fn: str = "swish",1041 resnet_groups: int = 32,1042 resnet_pre_norm: bool = True,1043 num_attention_heads: int = 1,1044 cross_attention_dim: int = 1280,1045 output_scale_factor: float = 1.0,1046 downsample_padding: int = 1,1047 add_downsample: bool = True,1048 dual_cross_attention: bool = False,1049 use_linear_projection: bool = False,1050 only_cross_attention: bool = False,1051 upcast_attention: bool = False,1052 attention_type: str = "default",1053 ):1054 super().__init__()1055 resnets = []1056 attentions = []1057 1058 self.has_cross_attention = True1059 self.num_attention_heads = num_attention_heads1060 if isinstance(transformer_layers_per_block, int):1061 transformer_layers_per_block = [transformer_layers_per_block] * num_layers1062 1063 for i in range(num_layers):1064 in_channels = in_channels if i == 0 else out_channels1065 resnets.append(1066 ResnetBlock2D(1067 in_channels=in_channels,1068 out_channels=out_channels,1069 temb_channels=temb_channels,1070 eps=resnet_eps,1071 groups=resnet_groups,1072 dropout=dropout,1073 time_embedding_norm=resnet_time_scale_shift,1074 non_linearity=resnet_act_fn,1075 output_scale_factor=output_scale_factor,1076 pre_norm=resnet_pre_norm,1077 )1078 )1079 if not dual_cross_attention:1080 attentions.append(1081 Transformer2DModel(1082 num_attention_heads,1083 out_channels // num_attention_heads,1084 in_channels=out_channels,1085 num_layers=transformer_layers_per_block[i],1086 cross_attention_dim=cross_attention_dim,1087 norm_num_groups=resnet_groups,1088 use_linear_projection=use_linear_projection,1089 only_cross_attention=only_cross_attention,1090 upcast_attention=upcast_attention,1091 attention_type=attention_type,1092 )1093 )1094 else:1095 attentions.append(1096 DualTransformer2DModel(1097 num_attention_heads,1098 out_channels // num_attention_heads,1099 in_channels=out_channels,1100 num_layers=1,1101 cross_attention_dim=cross_attention_dim,1102 norm_num_groups=resnet_groups,1103 )1104 )1105 self.attentions = nn.ModuleList(attentions)1106 self.resnets = nn.ModuleList(resnets)1107 1108 if add_downsample:1109 self.downsamplers = nn.ModuleList(1110 [1111 Downsample2D(1112 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"1113 )1114 ]1115 )1116 else:1117 self.downsamplers = None1118 1119 self.gradient_checkpointing = False1120 1121 def forward(1122 self,1123 hidden_states: torch.FloatTensor,1124 temb: Optional[torch.FloatTensor] = None,1125 encoder_hidden_states: Optional[torch.FloatTensor] = None,1126 attention_mask: Optional[torch.FloatTensor] = None,1127 cross_attention_kwargs: Optional[Dict[str, Any]] = None,1128 encoder_attention_mask: Optional[torch.FloatTensor] = None,1129 additional_residuals: Optional[torch.FloatTensor] = None,1130 ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:1131 output_states = ()1132 1133 lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.01134 1135 blocks = list(zip(self.resnets, self.attentions))1136 garment_features = []1137 for i, (resnet, attn) in enumerate(blocks):1138 if self.training and self.gradient_checkpointing:1139 1140 def create_custom_forward(module, return_dict=None):1141 def custom_forward(*inputs):1142 if return_dict is not None:1143 return module(*inputs, return_dict=return_dict)1144 else:1145 return module(*inputs)1146 1147 return custom_forward1148 1149 ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}1150 hidden_states = torch.utils.checkpoint.checkpoint(1151 create_custom_forward(resnet),1152 hidden_states,1153 temb,1154 **ckpt_kwargs,1155 )1156 hidden_states,out_garment_feat = attn(1157 hidden_states,1158 encoder_hidden_states=encoder_hidden_states,1159 cross_attention_kwargs=cross_attention_kwargs,1160 attention_mask=attention_mask,1161 encoder_attention_mask=encoder_attention_mask,1162 return_dict=False,1163 )1164 hidden_states=hidden_states[0]1165 else:1166 hidden_states = resnet(hidden_states, temb, scale=lora_scale)1167 hidden_states,out_garment_feat = attn(1168 hidden_states,1169 encoder_hidden_states=encoder_hidden_states,1170 cross_attention_kwargs=cross_attention_kwargs,1171 attention_mask=attention_mask,1172 encoder_attention_mask=encoder_attention_mask,1173 return_dict=False,1174 )1175 hidden_states=hidden_states[0]1176 garment_features += out_garment_feat1177 # apply additional residuals to the output of the last pair of resnet and attention blocks1178 if i == len(blocks) - 1 and additional_residuals is not None:1179 hidden_states = hidden_states + additional_residuals1180 1181 output_states = output_states + (hidden_states,)1182 1183 if self.downsamplers is not None:1184 for downsampler in self.downsamplers:1185 hidden_states = downsampler(hidden_states, scale=lora_scale)1186 1187 output_states = output_states + (hidden_states,)1188 1189 return hidden_states, output_states,garment_features1190 1191 1192class DownBlock2D(nn.Module):1193 def __init__(1194 self,1195 in_channels: int,1196 out_channels: int,1197 temb_channels: int,1198 dropout: float = 0.0,1199 num_layers: int = 1,1200 resnet_eps: float = 1e-6,