multimodalart/EchoMimic-zero
8
1# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_blocks.py2from typing import Any, Dict, Optional, Tuple, Union3 4import numpy as np5import torch6import torch.nn.functional as F7from diffusers.models.activations import get_activation8from diffusers.models.attention_processor import Attention9from diffusers.models.dual_transformer_2d import DualTransformer2DModel10from diffusers.models.resnet import Downsample2D, ResnetBlock2D, Upsample2D11from diffusers.utils import is_torch_version, logging12from diffusers.utils.torch_utils import apply_freeu13from torch import nn14 15from .transformer_2d import Transformer2DModel16 17logger = logging.get_logger(__name__) # pylint: disable=invalid-name18 19 20def get_down_block(21 down_block_type: str,22 num_layers: int,23 in_channels: int,24 out_channels: int,25 temb_channels: int,26 add_downsample: bool,27 resnet_eps: float,28 resnet_act_fn: str,29 transformer_layers_per_block: int = 1,30 num_attention_heads: Optional[int] = None,31 resnet_groups: Optional[int] = None,32 cross_attention_dim: Optional[int] = None,33 downsample_padding: Optional[int] = None,34 dual_cross_attention: bool = False,35 use_linear_projection: bool = False,36 only_cross_attention: bool = False,37 upcast_attention: bool = False,38 resnet_time_scale_shift: str = "default",39 attention_type: str = "default",40 resnet_skip_time_act: bool = False,41 resnet_out_scale_factor: float = 1.0,42 cross_attention_norm: Optional[str] = None,43 attention_head_dim: Optional[int] = None,44 downsample_type: Optional[str] = None,45 dropout: float = 0.0,46):47 # If attn head dim is not defined, we default it to the number of heads48 if attention_head_dim is None:49 logger.warn(50 f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."51 )52 attention_head_dim = num_attention_heads53 54 down_block_type = (55 down_block_type[7:]56 if down_block_type.startswith("UNetRes")57 else down_block_type58 )59 if down_block_type == "DownBlock2D":60 return DownBlock2D(61 num_layers=num_layers,62 in_channels=in_channels,63 out_channels=out_channels,64 temb_channels=temb_channels,65 dropout=dropout,66 add_downsample=add_downsample,67 resnet_eps=resnet_eps,68 resnet_act_fn=resnet_act_fn,69 resnet_groups=resnet_groups,70 downsample_padding=downsample_padding,71 resnet_time_scale_shift=resnet_time_scale_shift,72 )73 elif down_block_type == "CrossAttnDownBlock2D":74 if cross_attention_dim is None:75 raise ValueError(76 "cross_attention_dim must be specified for CrossAttnDownBlock2D"77 )78 return CrossAttnDownBlock2D(79 num_layers=num_layers,80 transformer_layers_per_block=transformer_layers_per_block,81 in_channels=in_channels,82 out_channels=out_channels,83 temb_channels=temb_channels,84 dropout=dropout,85 add_downsample=add_downsample,86 resnet_eps=resnet_eps,87 resnet_act_fn=resnet_act_fn,88 resnet_groups=resnet_groups,89 downsample_padding=downsample_padding,90 cross_attention_dim=cross_attention_dim,91 num_attention_heads=num_attention_heads,92 dual_cross_attention=dual_cross_attention,93 use_linear_projection=use_linear_projection,94 only_cross_attention=only_cross_attention,95 upcast_attention=upcast_attention,96 resnet_time_scale_shift=resnet_time_scale_shift,97 attention_type=attention_type,98 )99 raise ValueError(f"{down_block_type} does not exist.")100 101 102def get_up_block(103 up_block_type: str,104 num_layers: int,105 in_channels: int,106 out_channels: int,107 prev_output_channel: int,108 temb_channels: int,109 add_upsample: bool,110 resnet_eps: float,111 resnet_act_fn: str,112 resolution_idx: Optional[int] = None,113 transformer_layers_per_block: int = 1,114 num_attention_heads: Optional[int] = None,115 resnet_groups: Optional[int] = None,116 cross_attention_dim: Optional[int] = None,117 dual_cross_attention: bool = False,118 use_linear_projection: bool = False,119 only_cross_attention: bool = False,120 upcast_attention: bool = False,121 resnet_time_scale_shift: str = "default",122 attention_type: str = "default",123 resnet_skip_time_act: bool = False,124 resnet_out_scale_factor: float = 1.0,125 cross_attention_norm: Optional[str] = None,126 attention_head_dim: Optional[int] = None,127 upsample_type: Optional[str] = None,128 dropout: float = 0.0,129) -> nn.Module:130 # If attn head dim is not defined, we default it to the number of heads131 if attention_head_dim is None:132 logger.warn(133 f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."134 )135 attention_head_dim = num_attention_heads136 137 up_block_type = (138 up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type139 )140 if up_block_type == "UpBlock2D":141 return UpBlock2D(142 num_layers=num_layers,143 in_channels=in_channels,144 out_channels=out_channels,145 prev_output_channel=prev_output_channel,146 temb_channels=temb_channels,147 resolution_idx=resolution_idx,148 dropout=dropout,149 add_upsample=add_upsample,150 resnet_eps=resnet_eps,151 resnet_act_fn=resnet_act_fn,152 resnet_groups=resnet_groups,153 resnet_time_scale_shift=resnet_time_scale_shift,154 )155 elif up_block_type == "CrossAttnUpBlock2D":156 if cross_attention_dim is None:157 raise ValueError(158 "cross_attention_dim must be specified for CrossAttnUpBlock2D"159 )160 return CrossAttnUpBlock2D(161 num_layers=num_layers,162 transformer_layers_per_block=transformer_layers_per_block,163 in_channels=in_channels,164 out_channels=out_channels,165 prev_output_channel=prev_output_channel,166 temb_channels=temb_channels,167 resolution_idx=resolution_idx,168 dropout=dropout,169 add_upsample=add_upsample,170 resnet_eps=resnet_eps,171 resnet_act_fn=resnet_act_fn,172 resnet_groups=resnet_groups,173 cross_attention_dim=cross_attention_dim,174 num_attention_heads=num_attention_heads,175 dual_cross_attention=dual_cross_attention,176 use_linear_projection=use_linear_projection,177 only_cross_attention=only_cross_attention,178 upcast_attention=upcast_attention,179 resnet_time_scale_shift=resnet_time_scale_shift,180 attention_type=attention_type,181 )182 183 raise ValueError(f"{up_block_type} does not exist.")184 185 186class AutoencoderTinyBlock(nn.Module):187 """188 Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU189 blocks.190 191 Args:192 in_channels (`int`): The number of input channels.193 out_channels (`int`): The number of output channels.194 act_fn (`str`):195 ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`.196 197 Returns:198 `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to199 `out_channels`.200 """201 202 def __init__(self, in_channels: int, out_channels: int, act_fn: str):203 super().__init__()204 act_fn = get_activation(act_fn)205 self.conv = nn.Sequential(206 nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),207 act_fn,208 nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),209 act_fn,210 nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),211 )212 self.skip = (213 nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)214 if in_channels != out_channels215 else nn.Identity()216 )217 self.fuse = nn.ReLU()218 219 def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:220 return self.fuse(self.conv(x) + self.skip(x))221 222 223class UNetMidBlock2D(nn.Module):224 """225 A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks.226 227 Args:228 in_channels (`int`): The number of input channels.229 temb_channels (`int`): The number of temporal embedding channels.230 dropout (`float`, *optional*, defaults to 0.0): The dropout rate.231 num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.232 resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.233 resnet_time_scale_shift (`str`, *optional*, defaults to `default`):234 The type of normalization to apply to the time embeddings. This can help to improve the performance of the235 model on tasks with long-range temporal dependencies.236 resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.237 resnet_groups (`int`, *optional*, defaults to 32):238 The number of groups to use in the group normalization layers of the resnet blocks.239 attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.240 resnet_pre_norm (`bool`, *optional*, defaults to `True`):241 Whether to use pre-normalization for the resnet blocks.242 add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.243 attention_head_dim (`int`, *optional*, defaults to 1):244 Dimension of a single attention head. The number of attention heads is determined based on this value and245 the number of input channels.246 output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor.247 248 Returns:249 `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,250 in_channels, height, width)`.251 252 """253 254 def __init__(255 self,256 in_channels: int,257 temb_channels: int,258 dropout: float = 0.0,259 num_layers: int = 1,260 resnet_eps: float = 1e-6,261 resnet_time_scale_shift: str = "default", # default, spatial262 resnet_act_fn: str = "swish",263 resnet_groups: int = 32,264 attn_groups: Optional[int] = None,265 resnet_pre_norm: bool = True,266 add_attention: bool = True,267 attention_head_dim: int = 1,268 output_scale_factor: float = 1.0,269 ):270 super().__init__()271 resnet_groups = (272 resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)273 )274 self.add_attention = add_attention275 276 if attn_groups is None:277 attn_groups = (278 resnet_groups if resnet_time_scale_shift == "default" else None279 )280 281 # there is always at least one resnet282 resnets = [283 ResnetBlock2D(284 in_channels=in_channels,285 out_channels=in_channels,286 temb_channels=temb_channels,287 eps=resnet_eps,288 groups=resnet_groups,289 dropout=dropout,290 time_embedding_norm=resnet_time_scale_shift,291 non_linearity=resnet_act_fn,292 output_scale_factor=output_scale_factor,293 pre_norm=resnet_pre_norm,294 )295 ]296 attentions = []297 298 if attention_head_dim is None:299 logger.warn(300 f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."301 )302 attention_head_dim = in_channels303 304 for _ in range(num_layers):305 if self.add_attention:306 attentions.append(307 Attention(308 in_channels,309 heads=in_channels // attention_head_dim,310 dim_head=attention_head_dim,311 rescale_output_factor=output_scale_factor,312 eps=resnet_eps,313 norm_num_groups=attn_groups,314 spatial_norm_dim=temb_channels315 if resnet_time_scale_shift == "spatial"316 else None,317 residual_connection=True,318 bias=True,319 upcast_softmax=True,320 _from_deprecated_attn_block=True,321 )322 )323 else:324 attentions.append(None)325 326 resnets.append(327 ResnetBlock2D(328 in_channels=in_channels,329 out_channels=in_channels,330 temb_channels=temb_channels,331 eps=resnet_eps,332 groups=resnet_groups,333 dropout=dropout,334 time_embedding_norm=resnet_time_scale_shift,335 non_linearity=resnet_act_fn,336 output_scale_factor=output_scale_factor,337 pre_norm=resnet_pre_norm,338 )339 )340 341 self.attentions = nn.ModuleList(attentions)342 self.resnets = nn.ModuleList(resnets)343 344 def forward(345 self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None346 ) -> torch.FloatTensor:347 hidden_states = self.resnets[0](hidden_states, temb)348 for attn, resnet in zip(self.attentions, self.resnets[1:]):349 if attn is not None:350 hidden_states = attn(hidden_states, temb=temb)351 hidden_states = resnet(hidden_states, temb)352 353 return hidden_states354 355 356class UNetMidBlock2DCrossAttn(nn.Module):357 def __init__(358 self,359 in_channels: int,360 temb_channels: int,361 dropout: float = 0.0,362 num_layers: int = 1,363 transformer_layers_per_block: Union[int, Tuple[int]] = 1,364 resnet_eps: float = 1e-6,365 resnet_time_scale_shift: str = "default",366 resnet_act_fn: str = "swish",367 resnet_groups: int = 32,368 resnet_pre_norm: bool = True,369 num_attention_heads: int = 1,370 output_scale_factor: float = 1.0,371 cross_attention_dim: int = 1280,372 dual_cross_attention: bool = False,373 use_linear_projection: bool = False,374 upcast_attention: bool = False,375 attention_type: str = "default",376 ):377 super().__init__()378 379 self.has_cross_attention = True380 self.num_attention_heads = num_attention_heads381 resnet_groups = (382 resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)383 )384 385 # support for variable transformer layers per block386 if isinstance(transformer_layers_per_block, int):387 transformer_layers_per_block = [transformer_layers_per_block] * num_layers388 389 # there is always at least one resnet390 resnets = [391 ResnetBlock2D(392 in_channels=in_channels,393 out_channels=in_channels,394 temb_channels=temb_channels,395 eps=resnet_eps,396 groups=resnet_groups,397 dropout=dropout,398 time_embedding_norm=resnet_time_scale_shift,399 non_linearity=resnet_act_fn,400 output_scale_factor=output_scale_factor,401 pre_norm=resnet_pre_norm,402 )403 ]404 attentions = []405 406 for i in range(num_layers):407 if not dual_cross_attention:408 attentions.append(409 Transformer2DModel(410 num_attention_heads,411 in_channels // num_attention_heads,412 in_channels=in_channels,413 num_layers=transformer_layers_per_block[i],414 cross_attention_dim=cross_attention_dim,415 norm_num_groups=resnet_groups,416 use_linear_projection=use_linear_projection,417 upcast_attention=upcast_attention,418 attention_type=attention_type,419 )420 )421 else:422 attentions.append(423 DualTransformer2DModel(424 num_attention_heads,425 in_channels // num_attention_heads,426 in_channels=in_channels,427 num_layers=1,428 cross_attention_dim=cross_attention_dim,429 norm_num_groups=resnet_groups,430 )431 )432 resnets.append(433 ResnetBlock2D(434 in_channels=in_channels,435 out_channels=in_channels,436 temb_channels=temb_channels,437 eps=resnet_eps,438 groups=resnet_groups,439 dropout=dropout,440 time_embedding_norm=resnet_time_scale_shift,441 non_linearity=resnet_act_fn,442 output_scale_factor=output_scale_factor,443 pre_norm=resnet_pre_norm,444 )445 )446 447 self.attentions = nn.ModuleList(attentions)448 self.resnets = nn.ModuleList(resnets)449 450 self.gradient_checkpointing = False451 452 def forward(453 self,454 hidden_states: torch.FloatTensor,455 temb: Optional[torch.FloatTensor] = None,456 encoder_hidden_states: Optional[torch.FloatTensor] = None,457 attention_mask: Optional[torch.FloatTensor] = None,458 cross_attention_kwargs: Optional[Dict[str, Any]] = None,459 encoder_attention_mask: Optional[torch.FloatTensor] = None,460 ) -> torch.FloatTensor:461 lora_scale = (462 cross_attention_kwargs.get("scale", 1.0)463 if cross_attention_kwargs is not None464 else 1.0465 )466 hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)467 for attn, resnet in zip(self.attentions, self.resnets[1:]):468 if self.training and self.gradient_checkpointing:469 470 def create_custom_forward(module, return_dict=None):471 def custom_forward(*inputs):472 if return_dict is not None:473 return module(*inputs, return_dict=return_dict)474 else:475 return module(*inputs)476 477 return custom_forward478 479 ckpt_kwargs: Dict[str, Any] = (480 {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}481 )482 hidden_states, ref_feature = attn(483 hidden_states,484 encoder_hidden_states=encoder_hidden_states,485 cross_attention_kwargs=cross_attention_kwargs,486 attention_mask=attention_mask,487 encoder_attention_mask=encoder_attention_mask,488 return_dict=False,489 )490 hidden_states = torch.utils.checkpoint.checkpoint(491 create_custom_forward(resnet),492 hidden_states,493 temb,494 **ckpt_kwargs,495 )496 else:497 hidden_states, ref_feature = attn(498 hidden_states,499 encoder_hidden_states=encoder_hidden_states,500 cross_attention_kwargs=cross_attention_kwargs,501 attention_mask=attention_mask,502 encoder_attention_mask=encoder_attention_mask,503 return_dict=False,504 )505 hidden_states = resnet(hidden_states, temb, scale=lora_scale)506 507 return hidden_states508 509 510class CrossAttnDownBlock2D(nn.Module):511 def __init__(512 self,513 in_channels: int,514 out_channels: int,515 temb_channels: int,516 dropout: float = 0.0,517 num_layers: int = 1,518 transformer_layers_per_block: Union[int, Tuple[int]] = 1,519 resnet_eps: float = 1e-6,520 resnet_time_scale_shift: str = "default",521 resnet_act_fn: str = "swish",522 resnet_groups: int = 32,523 resnet_pre_norm: bool = True,524 num_attention_heads: int = 1,525 cross_attention_dim: int = 1280,526 output_scale_factor: float = 1.0,527 downsample_padding: int = 1,528 add_downsample: bool = True,529 dual_cross_attention: bool = False,530 use_linear_projection: bool = False,531 only_cross_attention: bool = False,532 upcast_attention: bool = False,533 attention_type: str = "default",534 ):535 super().__init__()536 resnets = []537 attentions = []538 539 self.has_cross_attention = True540 self.num_attention_heads = num_attention_heads541 if isinstance(transformer_layers_per_block, int):542 transformer_layers_per_block = [transformer_layers_per_block] * num_layers543 544 for i in range(num_layers):545 in_channels = in_channels if i == 0 else out_channels546 resnets.append(547 ResnetBlock2D(548 in_channels=in_channels,549 out_channels=out_channels,550 temb_channels=temb_channels,551 eps=resnet_eps,552 groups=resnet_groups,553 dropout=dropout,554 time_embedding_norm=resnet_time_scale_shift,555 non_linearity=resnet_act_fn,556 output_scale_factor=output_scale_factor,557 pre_norm=resnet_pre_norm,558 )559 )560 if not dual_cross_attention:561 attentions.append(562 Transformer2DModel(563 num_attention_heads,564 out_channels // num_attention_heads,565 in_channels=out_channels,566 num_layers=transformer_layers_per_block[i],567 cross_attention_dim=cross_attention_dim,568 norm_num_groups=resnet_groups,569 use_linear_projection=use_linear_projection,570 only_cross_attention=only_cross_attention,571 upcast_attention=upcast_attention,572 attention_type=attention_type,573 )574 )575 else:576 attentions.append(577 DualTransformer2DModel(578 num_attention_heads,579 out_channels // num_attention_heads,580 in_channels=out_channels,581 num_layers=1,582 cross_attention_dim=cross_attention_dim,583 norm_num_groups=resnet_groups,584 )585 )586 self.attentions = nn.ModuleList(attentions)587 self.resnets = nn.ModuleList(resnets)588 589 if add_downsample:590 self.downsamplers = nn.ModuleList(591 [592 Downsample2D(593 out_channels,594 use_conv=True,595 out_channels=out_channels,596 padding=downsample_padding,597 name="op",598 )599 ]600 )601 else:602 self.downsamplers = None603 604 self.gradient_checkpointing = False605 606 def forward(607 self,608 hidden_states: torch.FloatTensor,609 temb: Optional[torch.FloatTensor] = None,610 encoder_hidden_states: Optional[torch.FloatTensor] = None,611 attention_mask: Optional[torch.FloatTensor] = None,612 cross_attention_kwargs: Optional[Dict[str, Any]] = None,613 encoder_attention_mask: Optional[torch.FloatTensor] = None,614 additional_residuals: Optional[torch.FloatTensor] = None,615 ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:616 output_states = ()617 618 lora_scale = (619 cross_attention_kwargs.get("scale", 1.0)620 if cross_attention_kwargs is not None621 else 1.0622 )623 624 blocks = list(zip(self.resnets, self.attentions))625 626 for i, (resnet, attn) in enumerate(blocks):627 if self.training and self.gradient_checkpointing:628 629 def create_custom_forward(module, return_dict=None):630 def custom_forward(*inputs):631 if return_dict is not None:632 return module(*inputs, return_dict=return_dict)633 else:634 return module(*inputs)635 636 return custom_forward637 638 ckpt_kwargs: Dict[str, Any] = (639 {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}640 )641 hidden_states = torch.utils.checkpoint.checkpoint(642 create_custom_forward(resnet),643 hidden_states,644 temb,645 **ckpt_kwargs,646 )647 hidden_states, ref_feature = attn(648 hidden_states,649 encoder_hidden_states=encoder_hidden_states,650 cross_attention_kwargs=cross_attention_kwargs,651 attention_mask=attention_mask,652 encoder_attention_mask=encoder_attention_mask,653 return_dict=False,654 )655 else:656 hidden_states = resnet(hidden_states, temb, scale=lora_scale)657 hidden_states, ref_feature = attn(658 hidden_states,659 encoder_hidden_states=encoder_hidden_states,660 cross_attention_kwargs=cross_attention_kwargs,661 attention_mask=attention_mask,662 encoder_attention_mask=encoder_attention_mask,663 return_dict=False,664 )665 666 # apply additional residuals to the output of the last pair of resnet and attention blocks667 if i == len(blocks) - 1 and additional_residuals is not None:668 hidden_states = hidden_states + additional_residuals669 670 output_states = output_states + (hidden_states,)671 672 if self.downsamplers is not None:673 for downsampler in self.downsamplers:674 hidden_states = downsampler(hidden_states, scale=lora_scale)675 676 output_states = output_states + (hidden_states,)677 678 return hidden_states, output_states679 680 681class DownBlock2D(nn.Module):682 def __init__(683 self,684 in_channels: int,685 out_channels: int,686 temb_channels: int,687 dropout: float = 0.0,688 num_layers: int = 1,689 resnet_eps: float = 1e-6,690 resnet_time_scale_shift: str = "default",691 resnet_act_fn: str = "swish",692 resnet_groups: int = 32,693 resnet_pre_norm: bool = True,694 output_scale_factor: float = 1.0,695 add_downsample: bool = True,696 downsample_padding: int = 1,697 ):698 super().__init__()699 resnets = []700 701 for i in range(num_layers):702 in_channels = in_channels if i == 0 else out_channels703 resnets.append(704 ResnetBlock2D(705 in_channels=in_channels,706 out_channels=out_channels,707 temb_channels=temb_channels,708 eps=resnet_eps,709 groups=resnet_groups,710 dropout=dropout,711 time_embedding_norm=resnet_time_scale_shift,712 non_linearity=resnet_act_fn,713 output_scale_factor=output_scale_factor,714 pre_norm=resnet_pre_norm,715 )716 )717 718 self.resnets = nn.ModuleList(resnets)719 720 if add_downsample:721 self.downsamplers = nn.ModuleList(722 [723 Downsample2D(724 out_channels,725 use_conv=True,726 out_channels=out_channels,727 padding=downsample_padding,728 name="op",729 )730 ]731 )732 else:733 self.downsamplers = None734 735 self.gradient_checkpointing = False736 737 def forward(738 self,739 hidden_states: torch.FloatTensor,740 temb: Optional[torch.FloatTensor] = None,741 scale: float = 1.0,742 ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:743 output_states = ()744 745 for resnet in self.resnets:746 if self.training and self.gradient_checkpointing:747 748 def create_custom_forward(module):749 def custom_forward(*inputs):750 return module(*inputs)751 752 return custom_forward753 754 if is_torch_version(">=", "1.11.0"):755 hidden_states = torch.utils.checkpoint.checkpoint(756 create_custom_forward(resnet),757 hidden_states,758 temb,759 use_reentrant=False,760 )761 else:762 hidden_states = torch.utils.checkpoint.checkpoint(763 create_custom_forward(resnet), hidden_states, temb764 )765 else:766 hidden_states = resnet(hidden_states, temb, scale=scale)767 768 output_states = output_states + (hidden_states,)769 770 if self.downsamplers is not None:771 for downsampler in self.downsamplers:772 hidden_states = downsampler(hidden_states, scale=scale)773 774 output_states = output_states + (hidden_states,)775 776 return hidden_states, output_states777 778 779class CrossAttnUpBlock2D(nn.Module):780 def __init__(781 self,782 in_channels: int,783 out_channels: int,784 prev_output_channel: int,785 temb_channels: int,786 resolution_idx: Optional[int] = None,787 dropout: float = 0.0,788 num_layers: int = 1,789 transformer_layers_per_block: Union[int, Tuple[int]] = 1,790 resnet_eps: float = 1e-6,791 resnet_time_scale_shift: str = "default",792 resnet_act_fn: str = "swish",793 resnet_groups: int = 32,794 resnet_pre_norm: bool = True,795 num_attention_heads: int = 1,796 cross_attention_dim: int = 1280,797 output_scale_factor: float = 1.0,798 add_upsample: bool = True,799 dual_cross_attention: bool = False,800 use_linear_projection: bool = False,801 only_cross_attention: bool = False,802 upcast_attention: bool = False,803 attention_type: str = "default",804 ):805 super().__init__()806 resnets = []807 attentions = []808 809 self.has_cross_attention = True810 self.num_attention_heads = num_attention_heads811 812 if isinstance(transformer_layers_per_block, int):813 transformer_layers_per_block = [transformer_layers_per_block] * num_layers814 815 for i in range(num_layers):816 res_skip_channels = in_channels if (i == num_layers - 1) else out_channels817 resnet_in_channels = prev_output_channel if i == 0 else out_channels818 819 resnets.append(820 ResnetBlock2D(821 in_channels=resnet_in_channels + res_skip_channels,822 out_channels=out_channels,823 temb_channels=temb_channels,824 eps=resnet_eps,825 groups=resnet_groups,826 dropout=dropout,827 time_embedding_norm=resnet_time_scale_shift,828 non_linearity=resnet_act_fn,829 output_scale_factor=output_scale_factor,830 pre_norm=resnet_pre_norm,831 )832 )833 if not dual_cross_attention:834 attentions.append(835 Transformer2DModel(836 num_attention_heads,837 out_channels // num_attention_heads,838 in_channels=out_channels,839 num_layers=transformer_layers_per_block[i],840 cross_attention_dim=cross_attention_dim,841 norm_num_groups=resnet_groups,842 use_linear_projection=use_linear_projection,843 only_cross_attention=only_cross_attention,844 upcast_attention=upcast_attention,845 attention_type=attention_type,846 )847 )848 else:849 attentions.append(850 DualTransformer2DModel(851 num_attention_heads,852 out_channels // num_attention_heads,853 in_channels=out_channels,854 num_layers=1,855 cross_attention_dim=cross_attention_dim,856 norm_num_groups=resnet_groups,857 )858 )859 self.attentions = nn.ModuleList(attentions)860 self.resnets = nn.ModuleList(resnets)861 862 if add_upsample:863 self.upsamplers = nn.ModuleList(864 [Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]865 )866 else:867 self.upsamplers = None868 869 self.gradient_checkpointing = False870 self.resolution_idx = resolution_idx871 872 def forward(873 self,874 hidden_states: torch.FloatTensor,875 res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],876 temb: Optional[torch.FloatTensor] = None,877 encoder_hidden_states: Optional[torch.FloatTensor] = None,878 cross_attention_kwargs: Optional[Dict[str, Any]] = None,879 upsample_size: Optional[int] = None,880 attention_mask: Optional[torch.FloatTensor] = None,881 encoder_attention_mask: Optional[torch.FloatTensor] = None,882 ) -> torch.FloatTensor:883 lora_scale = (884 cross_attention_kwargs.get("scale", 1.0)885 if cross_attention_kwargs is not None886 else 1.0887 )888 is_freeu_enabled = (889 getattr(self, "s1", None)890 and getattr(self, "s2", None)891 and getattr(self, "b1", None)892 and getattr(self, "b2", None)893 )894 895 for resnet, attn in zip(self.resnets, self.attentions):896 # pop res hidden states897 res_hidden_states = res_hidden_states_tuple[-1]898 res_hidden_states_tuple = res_hidden_states_tuple[:-1]899 900 # FreeU: Only operate on the first two stages901 if is_freeu_enabled:902 hidden_states, res_hidden_states = apply_freeu(903 self.resolution_idx,904 hidden_states,905 res_hidden_states,906 s1=self.s1,907 s2=self.s2,908 b1=self.b1,909 b2=self.b2,910 )911 912 hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)913 914 if self.training and self.gradient_checkpointing:915 916 def create_custom_forward(module, return_dict=None):917 def custom_forward(*inputs):918 if return_dict is not None:919 return module(*inputs, return_dict=return_dict)920 else:921 return module(*inputs)922 923 return custom_forward924 925 ckpt_kwargs: Dict[str, Any] = (926 {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}927 )928 hidden_states = torch.utils.checkpoint.checkpoint(929 create_custom_forward(resnet),930 hidden_states,931 temb,932 **ckpt_kwargs,933 )934 hidden_states, ref_feature = attn(935 hidden_states,936 encoder_hidden_states=encoder_hidden_states,937 cross_attention_kwargs=cross_attention_kwargs,938 attention_mask=attention_mask,939 encoder_attention_mask=encoder_attention_mask,940 return_dict=False,941 )942 else:943 hidden_states = resnet(hidden_states, temb, scale=lora_scale)944 hidden_states, ref_feature = attn(945 hidden_states,946 encoder_hidden_states=encoder_hidden_states,947 cross_attention_kwargs=cross_attention_kwargs,948 attention_mask=attention_mask,949 encoder_attention_mask=encoder_attention_mask,950 return_dict=False,951 )952 953 if self.upsamplers is not None:954 for upsampler in self.upsamplers:955 hidden_states = upsampler(956 hidden_states, upsample_size, scale=lora_scale957 )958 959 return hidden_states960 961 962class UpBlock2D(nn.Module):963 def __init__(964 self,965 in_channels: int,966 prev_output_channel: int,967 out_channels: int,968 temb_channels: int,969 resolution_idx: Optional[int] = None,970 dropout: float = 0.0,971 num_layers: int = 1,972 resnet_eps: float = 1e-6,973 resnet_time_scale_shift: str = "default",974 resnet_act_fn: str = "swish",975 resnet_groups: int = 32,976 resnet_pre_norm: bool = True,977 output_scale_factor: float = 1.0,978 add_upsample: bool = True,979 ):980 super().__init__()981 resnets = []982 983 for i in range(num_layers):984 res_skip_channels = in_channels if (i == num_layers - 1) else out_channels985 resnet_in_channels = prev_output_channel if i == 0 else out_channels986 987 resnets.append(988 ResnetBlock2D(989 in_channels=resnet_in_channels + res_skip_channels,990 out_channels=out_channels,991 temb_channels=temb_channels,992 eps=resnet_eps,993 groups=resnet_groups,994 dropout=dropout,995 time_embedding_norm=resnet_time_scale_shift,996 non_linearity=resnet_act_fn,997 output_scale_factor=output_scale_factor,998 pre_norm=resnet_pre_norm,999 )1000 )1001 1002 self.resnets = nn.ModuleList(resnets)1003 1004 if add_upsample:1005 self.upsamplers = nn.ModuleList(1006 [Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]1007 )1008 else:1009 self.upsamplers = None1010 1011 self.gradient_checkpointing = False1012 self.resolution_idx = resolution_idx1013 1014 def forward(1015 self,1016 hidden_states: torch.FloatTensor,1017 res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],1018 temb: Optional[torch.FloatTensor] = None,1019 upsample_size: Optional[int] = None,1020 scale: float = 1.0,1021 ) -> torch.FloatTensor:1022 is_freeu_enabled = (1023 getattr(self, "s1", None)1024 and getattr(self, "s2", None)1025 and getattr(self, "b1", None)1026 and getattr(self, "b2", None)1027 )1028 1029 for resnet in self.resnets:1030 # pop res hidden states1031 res_hidden_states = res_hidden_states_tuple[-1]1032 res_hidden_states_tuple = res_hidden_states_tuple[:-1]1033 1034 # FreeU: Only operate on the first two stages1035 if is_freeu_enabled:1036 hidden_states, res_hidden_states = apply_freeu(1037 self.resolution_idx,1038 hidden_states,1039 res_hidden_states,1040 s1=self.s1,1041 s2=self.s2,1042 b1=self.b1,1043 b2=self.b2,1044 )1045 1046 hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)1047 1048 if self.training and self.gradient_checkpointing:1049 1050 def create_custom_forward(module):1051 def custom_forward(*inputs):1052 return module(*inputs)1053 1054 return custom_forward1055 1056 if is_torch_version(">=", "1.11.0"):1057 hidden_states = torch.utils.checkpoint.checkpoint(1058 create_custom_forward(resnet),1059 hidden_states,1060 temb,1061 use_reentrant=False,1062 )1063 else:1064 hidden_states = torch.utils.checkpoint.checkpoint(1065 create_custom_forward(resnet), hidden_states, temb1066 )1067 else:1068 hidden_states = resnet(hidden_states, temb, scale=scale)1069 1070 if self.upsamplers is not None:1071 for upsampler in self.upsamplers:1072 hidden_states = upsampler(hidden_states, upsample_size, scale=scale)1073 1074 return hidden_states1075