declare-lab/tango2
92
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, Tuple15 16import numpy as np17import torch18from torch import nn19 20from .attention import AdaGroupNorm, AttentionBlock21from .attention_processor import Attention, AttnAddedKVProcessor22from .dual_transformer_2d import DualTransformer2DModel23from .resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D24from .transformer_2d import Transformer2DModel, Transformer2DModelOutput25 26 27def get_down_block(28 down_block_type,29 num_layers,30 in_channels,31 out_channels,32 temb_channels,33 add_downsample,34 resnet_eps,35 resnet_act_fn,36 attn_num_head_channels,37 resnet_groups=None,38 cross_attention_dim=None,39 downsample_padding=None,40 dual_cross_attention=False,41 use_linear_projection=False,42 only_cross_attention=False,43 upcast_attention=False,44 resnet_time_scale_shift="default",45):46 down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type47 if down_block_type == "DownBlock2D":48 return DownBlock2D(49 num_layers=num_layers,50 in_channels=in_channels,51 out_channels=out_channels,52 temb_channels=temb_channels,53 add_downsample=add_downsample,54 resnet_eps=resnet_eps,55 resnet_act_fn=resnet_act_fn,56 resnet_groups=resnet_groups,57 downsample_padding=downsample_padding,58 resnet_time_scale_shift=resnet_time_scale_shift,59 )60 elif down_block_type == "ResnetDownsampleBlock2D":61 return ResnetDownsampleBlock2D(62 num_layers=num_layers,63 in_channels=in_channels,64 out_channels=out_channels,65 temb_channels=temb_channels,66 add_downsample=add_downsample,67 resnet_eps=resnet_eps,68 resnet_act_fn=resnet_act_fn,69 resnet_groups=resnet_groups,70 resnet_time_scale_shift=resnet_time_scale_shift,71 )72 elif down_block_type == "AttnDownBlock2D":73 return AttnDownBlock2D(74 num_layers=num_layers,75 in_channels=in_channels,76 out_channels=out_channels,77 temb_channels=temb_channels,78 add_downsample=add_downsample,79 resnet_eps=resnet_eps,80 resnet_act_fn=resnet_act_fn,81 resnet_groups=resnet_groups,82 downsample_padding=downsample_padding,83 attn_num_head_channels=attn_num_head_channels,84 resnet_time_scale_shift=resnet_time_scale_shift,85 )86 elif down_block_type == "CrossAttnDownBlock2D":87 if cross_attention_dim is None:88 raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D")89 return CrossAttnDownBlock2D(90 num_layers=num_layers,91 in_channels=in_channels,92 out_channels=out_channels,93 temb_channels=temb_channels,94 add_downsample=add_downsample,95 resnet_eps=resnet_eps,96 resnet_act_fn=resnet_act_fn,97 resnet_groups=resnet_groups,98 downsample_padding=downsample_padding,99 cross_attention_dim=cross_attention_dim,100 attn_num_head_channels=attn_num_head_channels,101 dual_cross_attention=dual_cross_attention,102 use_linear_projection=use_linear_projection,103 only_cross_attention=only_cross_attention,104 upcast_attention=upcast_attention,105 resnet_time_scale_shift=resnet_time_scale_shift,106 )107 elif down_block_type == "SimpleCrossAttnDownBlock2D":108 if cross_attention_dim is None:109 raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D")110 return SimpleCrossAttnDownBlock2D(111 num_layers=num_layers,112 in_channels=in_channels,113 out_channels=out_channels,114 temb_channels=temb_channels,115 add_downsample=add_downsample,116 resnet_eps=resnet_eps,117 resnet_act_fn=resnet_act_fn,118 resnet_groups=resnet_groups,119 cross_attention_dim=cross_attention_dim,120 attn_num_head_channels=attn_num_head_channels,121 resnet_time_scale_shift=resnet_time_scale_shift,122 )123 elif down_block_type == "SkipDownBlock2D":124 return SkipDownBlock2D(125 num_layers=num_layers,126 in_channels=in_channels,127 out_channels=out_channels,128 temb_channels=temb_channels,129 add_downsample=add_downsample,130 resnet_eps=resnet_eps,131 resnet_act_fn=resnet_act_fn,132 downsample_padding=downsample_padding,133 resnet_time_scale_shift=resnet_time_scale_shift,134 )135 elif down_block_type == "AttnSkipDownBlock2D":136 return AttnSkipDownBlock2D(137 num_layers=num_layers,138 in_channels=in_channels,139 out_channels=out_channels,140 temb_channels=temb_channels,141 add_downsample=add_downsample,142 resnet_eps=resnet_eps,143 resnet_act_fn=resnet_act_fn,144 downsample_padding=downsample_padding,145 attn_num_head_channels=attn_num_head_channels,146 resnet_time_scale_shift=resnet_time_scale_shift,147 )148 elif down_block_type == "DownEncoderBlock2D":149 return DownEncoderBlock2D(150 num_layers=num_layers,151 in_channels=in_channels,152 out_channels=out_channels,153 add_downsample=add_downsample,154 resnet_eps=resnet_eps,155 resnet_act_fn=resnet_act_fn,156 resnet_groups=resnet_groups,157 downsample_padding=downsample_padding,158 resnet_time_scale_shift=resnet_time_scale_shift,159 )160 elif down_block_type == "AttnDownEncoderBlock2D":161 return AttnDownEncoderBlock2D(162 num_layers=num_layers,163 in_channels=in_channels,164 out_channels=out_channels,165 add_downsample=add_downsample,166 resnet_eps=resnet_eps,167 resnet_act_fn=resnet_act_fn,168 resnet_groups=resnet_groups,169 downsample_padding=downsample_padding,170 attn_num_head_channels=attn_num_head_channels,171 resnet_time_scale_shift=resnet_time_scale_shift,172 )173 elif down_block_type == "KDownBlock2D":174 return KDownBlock2D(175 num_layers=num_layers,176 in_channels=in_channels,177 out_channels=out_channels,178 temb_channels=temb_channels,179 add_downsample=add_downsample,180 resnet_eps=resnet_eps,181 resnet_act_fn=resnet_act_fn,182 )183 elif down_block_type == "KCrossAttnDownBlock2D":184 return KCrossAttnDownBlock2D(185 num_layers=num_layers,186 in_channels=in_channels,187 out_channels=out_channels,188 temb_channels=temb_channels,189 add_downsample=add_downsample,190 resnet_eps=resnet_eps,191 resnet_act_fn=resnet_act_fn,192 cross_attention_dim=cross_attention_dim,193 attn_num_head_channels=attn_num_head_channels,194 add_self_attention=True if not add_downsample else False,195 )196 raise ValueError(f"{down_block_type} does not exist.")197 198 199def get_up_block(200 up_block_type,201 num_layers,202 in_channels,203 out_channels,204 prev_output_channel,205 temb_channels,206 add_upsample,207 resnet_eps,208 resnet_act_fn,209 attn_num_head_channels,210 resnet_groups=None,211 cross_attention_dim=None,212 dual_cross_attention=False,213 use_linear_projection=False,214 only_cross_attention=False,215 upcast_attention=False,216 resnet_time_scale_shift="default",217):218 up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type219 if up_block_type == "UpBlock2D":220 return UpBlock2D(221 num_layers=num_layers,222 in_channels=in_channels,223 out_channels=out_channels,224 prev_output_channel=prev_output_channel,225 temb_channels=temb_channels,226 add_upsample=add_upsample,227 resnet_eps=resnet_eps,228 resnet_act_fn=resnet_act_fn,229 resnet_groups=resnet_groups,230 resnet_time_scale_shift=resnet_time_scale_shift,231 )232 elif up_block_type == "ResnetUpsampleBlock2D":233 return ResnetUpsampleBlock2D(234 num_layers=num_layers,235 in_channels=in_channels,236 out_channels=out_channels,237 prev_output_channel=prev_output_channel,238 temb_channels=temb_channels,239 add_upsample=add_upsample,240 resnet_eps=resnet_eps,241 resnet_act_fn=resnet_act_fn,242 resnet_groups=resnet_groups,243 resnet_time_scale_shift=resnet_time_scale_shift,244 )245 elif up_block_type == "CrossAttnUpBlock2D":246 if cross_attention_dim is None:247 raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")248 return CrossAttnUpBlock2D(249 num_layers=num_layers,250 in_channels=in_channels,251 out_channels=out_channels,252 prev_output_channel=prev_output_channel,253 temb_channels=temb_channels,254 add_upsample=add_upsample,255 resnet_eps=resnet_eps,256 resnet_act_fn=resnet_act_fn,257 resnet_groups=resnet_groups,258 cross_attention_dim=cross_attention_dim,259 attn_num_head_channels=attn_num_head_channels,260 dual_cross_attention=dual_cross_attention,261 use_linear_projection=use_linear_projection,262 only_cross_attention=only_cross_attention,263 upcast_attention=upcast_attention,264 resnet_time_scale_shift=resnet_time_scale_shift,265 )266 elif up_block_type == "SimpleCrossAttnUpBlock2D":267 if cross_attention_dim is None:268 raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D")269 return SimpleCrossAttnUpBlock2D(270 num_layers=num_layers,271 in_channels=in_channels,272 out_channels=out_channels,273 prev_output_channel=prev_output_channel,274 temb_channels=temb_channels,275 add_upsample=add_upsample,276 resnet_eps=resnet_eps,277 resnet_act_fn=resnet_act_fn,278 resnet_groups=resnet_groups,279 cross_attention_dim=cross_attention_dim,280 attn_num_head_channels=attn_num_head_channels,281 resnet_time_scale_shift=resnet_time_scale_shift,282 )283 elif up_block_type == "AttnUpBlock2D":284 return AttnUpBlock2D(285 num_layers=num_layers,286 in_channels=in_channels,287 out_channels=out_channels,288 prev_output_channel=prev_output_channel,289 temb_channels=temb_channels,290 add_upsample=add_upsample,291 resnet_eps=resnet_eps,292 resnet_act_fn=resnet_act_fn,293 resnet_groups=resnet_groups,294 attn_num_head_channels=attn_num_head_channels,295 resnet_time_scale_shift=resnet_time_scale_shift,296 )297 elif up_block_type == "SkipUpBlock2D":298 return SkipUpBlock2D(299 num_layers=num_layers,300 in_channels=in_channels,301 out_channels=out_channels,302 prev_output_channel=prev_output_channel,303 temb_channels=temb_channels,304 add_upsample=add_upsample,305 resnet_eps=resnet_eps,306 resnet_act_fn=resnet_act_fn,307 resnet_time_scale_shift=resnet_time_scale_shift,308 )309 elif up_block_type == "AttnSkipUpBlock2D":310 return AttnSkipUpBlock2D(311 num_layers=num_layers,312 in_channels=in_channels,313 out_channels=out_channels,314 prev_output_channel=prev_output_channel,315 temb_channels=temb_channels,316 add_upsample=add_upsample,317 resnet_eps=resnet_eps,318 resnet_act_fn=resnet_act_fn,319 attn_num_head_channels=attn_num_head_channels,320 resnet_time_scale_shift=resnet_time_scale_shift,321 )322 elif up_block_type == "UpDecoderBlock2D":323 return UpDecoderBlock2D(324 num_layers=num_layers,325 in_channels=in_channels,326 out_channels=out_channels,327 add_upsample=add_upsample,328 resnet_eps=resnet_eps,329 resnet_act_fn=resnet_act_fn,330 resnet_groups=resnet_groups,331 resnet_time_scale_shift=resnet_time_scale_shift,332 )333 elif up_block_type == "AttnUpDecoderBlock2D":334 return AttnUpDecoderBlock2D(335 num_layers=num_layers,336 in_channels=in_channels,337 out_channels=out_channels,338 add_upsample=add_upsample,339 resnet_eps=resnet_eps,340 resnet_act_fn=resnet_act_fn,341 resnet_groups=resnet_groups,342 attn_num_head_channels=attn_num_head_channels,343 resnet_time_scale_shift=resnet_time_scale_shift,344 )345 elif up_block_type == "KUpBlock2D":346 return KUpBlock2D(347 num_layers=num_layers,348 in_channels=in_channels,349 out_channels=out_channels,350 temb_channels=temb_channels,351 add_upsample=add_upsample,352 resnet_eps=resnet_eps,353 resnet_act_fn=resnet_act_fn,354 )355 elif up_block_type == "KCrossAttnUpBlock2D":356 return KCrossAttnUpBlock2D(357 num_layers=num_layers,358 in_channels=in_channels,359 out_channels=out_channels,360 temb_channels=temb_channels,361 add_upsample=add_upsample,362 resnet_eps=resnet_eps,363 resnet_act_fn=resnet_act_fn,364 cross_attention_dim=cross_attention_dim,365 attn_num_head_channels=attn_num_head_channels,366 )367 368 raise ValueError(f"{up_block_type} does not exist.")369 370 371class UNetMidBlock2D(nn.Module):372 def __init__(373 self,374 in_channels: int,375 temb_channels: int,376 dropout: float = 0.0,377 num_layers: int = 1,378 resnet_eps: float = 1e-6,379 resnet_time_scale_shift: str = "default",380 resnet_act_fn: str = "swish",381 resnet_groups: int = 32,382 resnet_pre_norm: bool = True,383 add_attention: bool = True,384 attn_num_head_channels=1,385 output_scale_factor=1.0,386 ):387 super().__init__()388 resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)389 self.add_attention = add_attention390 391 # there is always at least one resnet392 resnets = [393 ResnetBlock2D(394 in_channels=in_channels,395 out_channels=in_channels,396 temb_channels=temb_channels,397 eps=resnet_eps,398 groups=resnet_groups,399 dropout=dropout,400 time_embedding_norm=resnet_time_scale_shift,401 non_linearity=resnet_act_fn,402 output_scale_factor=output_scale_factor,403 pre_norm=resnet_pre_norm,404 )405 ]406 attentions = []407 408 for _ in range(num_layers):409 if self.add_attention:410 attentions.append(411 AttentionBlock(412 in_channels,413 num_head_channels=attn_num_head_channels,414 rescale_output_factor=output_scale_factor,415 eps=resnet_eps,416 norm_num_groups=resnet_groups,417 )418 )419 else:420 attentions.append(None)421 422 resnets.append(423 ResnetBlock2D(424 in_channels=in_channels,425 out_channels=in_channels,426 temb_channels=temb_channels,427 eps=resnet_eps,428 groups=resnet_groups,429 dropout=dropout,430 time_embedding_norm=resnet_time_scale_shift,431 non_linearity=resnet_act_fn,432 output_scale_factor=output_scale_factor,433 pre_norm=resnet_pre_norm,434 )435 )436 437 self.attentions = nn.ModuleList(attentions)438 self.resnets = nn.ModuleList(resnets)439 440 def forward(self, hidden_states, temb=None):441 hidden_states = self.resnets[0](hidden_states, temb)442 for attn, resnet in zip(self.attentions, self.resnets[1:]):443 if attn is not None:444 hidden_states = attn(hidden_states)445 hidden_states = resnet(hidden_states, temb)446 447 return hidden_states448 449 450class UNetMidBlock2DCrossAttn(nn.Module):451 def __init__(452 self,453 in_channels: int,454 temb_channels: int,455 dropout: float = 0.0,456 num_layers: int = 1,457 resnet_eps: float = 1e-6,458 resnet_time_scale_shift: str = "default",459 resnet_act_fn: str = "swish",460 resnet_groups: int = 32,461 resnet_pre_norm: bool = True,462 attn_num_head_channels=1,463 output_scale_factor=1.0,464 cross_attention_dim=1280,465 dual_cross_attention=False,466 use_linear_projection=False,467 upcast_attention=False,468 ):469 super().__init__()470 471 self.has_cross_attention = True472 self.attn_num_head_channels = attn_num_head_channels473 resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)474 475 # there is always at least one resnet476 resnets = [477 ResnetBlock2D(478 in_channels=in_channels,479 out_channels=in_channels,480 temb_channels=temb_channels,481 eps=resnet_eps,482 groups=resnet_groups,483 dropout=dropout,484 time_embedding_norm=resnet_time_scale_shift,485 non_linearity=resnet_act_fn,486 output_scale_factor=output_scale_factor,487 pre_norm=resnet_pre_norm,488 )489 ]490 attentions = []491 492 for _ in range(num_layers):493 if not dual_cross_attention:494 attentions.append(495 Transformer2DModel(496 attn_num_head_channels,497 in_channels // attn_num_head_channels,498 in_channels=in_channels,499 num_layers=1,500 cross_attention_dim=cross_attention_dim,501 norm_num_groups=resnet_groups,502 use_linear_projection=use_linear_projection,503 upcast_attention=upcast_attention,504 )505 )506 else:507 attentions.append(508 DualTransformer2DModel(509 attn_num_head_channels,510 in_channels // attn_num_head_channels,511 in_channels=in_channels,512 num_layers=1,513 cross_attention_dim=cross_attention_dim,514 norm_num_groups=resnet_groups,515 )516 )517 resnets.append(518 ResnetBlock2D(519 in_channels=in_channels,520 out_channels=in_channels,521 temb_channels=temb_channels,522 eps=resnet_eps,523 groups=resnet_groups,524 dropout=dropout,525 time_embedding_norm=resnet_time_scale_shift,526 non_linearity=resnet_act_fn,527 output_scale_factor=output_scale_factor,528 pre_norm=resnet_pre_norm,529 )530 )531 532 self.attentions = nn.ModuleList(attentions)533 self.resnets = nn.ModuleList(resnets)534 535 def forward(536 self,537 hidden_states: torch.FloatTensor,538 temb: Optional[torch.FloatTensor] = None,539 encoder_hidden_states: Optional[torch.FloatTensor] = None,540 attention_mask: Optional[torch.FloatTensor] = None,541 cross_attention_kwargs: Optional[Dict[str, Any]] = None,542 encoder_attention_mask: Optional[torch.FloatTensor] = None,543 ) -> torch.FloatTensor:544 hidden_states = self.resnets[0](hidden_states, temb)545 for attn, resnet in zip(self.attentions, self.resnets[1:]):546 output: Transformer2DModelOutput = attn(547 hidden_states,548 encoder_hidden_states=encoder_hidden_states,549 cross_attention_kwargs=cross_attention_kwargs,550 attention_mask=attention_mask,551 encoder_attention_mask=encoder_attention_mask,552 )553 hidden_states = output.sample554 hidden_states = resnet(hidden_states, temb)555 556 return hidden_states557 558 559class UNetMidBlock2DSimpleCrossAttn(nn.Module):560 def __init__(561 self,562 in_channels: int,563 temb_channels: int,564 dropout: float = 0.0,565 num_layers: int = 1,566 resnet_eps: float = 1e-6,567 resnet_time_scale_shift: str = "default",568 resnet_act_fn: str = "swish",569 resnet_groups: int = 32,570 resnet_pre_norm: bool = True,571 attn_num_head_channels=1,572 output_scale_factor=1.0,573 cross_attention_dim=1280,574 ):575 super().__init__()576 577 self.has_cross_attention = True578 579 self.attn_num_head_channels = attn_num_head_channels580 resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)581 582 self.num_heads = in_channels // self.attn_num_head_channels583 584 # there is always at least one resnet585 resnets = [586 ResnetBlock2D(587 in_channels=in_channels,588 out_channels=in_channels,589 temb_channels=temb_channels,590 eps=resnet_eps,591 groups=resnet_groups,592 dropout=dropout,593 time_embedding_norm=resnet_time_scale_shift,594 non_linearity=resnet_act_fn,595 output_scale_factor=output_scale_factor,596 pre_norm=resnet_pre_norm,597 )598 ]599 attentions = []600 601 for _ in range(num_layers):602 attentions.append(603 Attention(604 query_dim=in_channels,605 cross_attention_dim=in_channels,606 heads=self.num_heads,607 dim_head=attn_num_head_channels,608 added_kv_proj_dim=cross_attention_dim,609 norm_num_groups=resnet_groups,610 bias=True,611 upcast_softmax=True,612 processor=AttnAddedKVProcessor(),613 )614 )615 resnets.append(616 ResnetBlock2D(617 in_channels=in_channels,618 out_channels=in_channels,619 temb_channels=temb_channels,620 eps=resnet_eps,621 groups=resnet_groups,622 dropout=dropout,623 time_embedding_norm=resnet_time_scale_shift,624 non_linearity=resnet_act_fn,625 output_scale_factor=output_scale_factor,626 pre_norm=resnet_pre_norm,627 )628 )629 630 self.attentions = nn.ModuleList(attentions)631 self.resnets = nn.ModuleList(resnets)632 633 def forward(634 self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None, cross_attention_kwargs=None635 ):636 cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}637 hidden_states = self.resnets[0](hidden_states, temb)638 for attn, resnet in zip(self.attentions, self.resnets[1:]):639 # attn640 hidden_states = attn(641 hidden_states,642 encoder_hidden_states=encoder_hidden_states,643 attention_mask=attention_mask,644 **cross_attention_kwargs,645 )646 647 # resnet648 hidden_states = resnet(hidden_states, temb)649 650 return hidden_states651 652 653class AttnDownBlock2D(nn.Module):654 def __init__(655 self,656 in_channels: int,657 out_channels: int,658 temb_channels: int,659 dropout: float = 0.0,660 num_layers: int = 1,661 resnet_eps: float = 1e-6,662 resnet_time_scale_shift: str = "default",663 resnet_act_fn: str = "swish",664 resnet_groups: int = 32,665 resnet_pre_norm: bool = True,666 attn_num_head_channels=1,667 output_scale_factor=1.0,668 downsample_padding=1,669 add_downsample=True,670 ):671 super().__init__()672 resnets = []673 attentions = []674 675 for i in range(num_layers):676 in_channels = in_channels if i == 0 else out_channels677 resnets.append(678 ResnetBlock2D(679 in_channels=in_channels,680 out_channels=out_channels,681 temb_channels=temb_channels,682 eps=resnet_eps,683 groups=resnet_groups,684 dropout=dropout,685 time_embedding_norm=resnet_time_scale_shift,686 non_linearity=resnet_act_fn,687 output_scale_factor=output_scale_factor,688 pre_norm=resnet_pre_norm,689 )690 )691 attentions.append(692 AttentionBlock(693 out_channels,694 num_head_channels=attn_num_head_channels,695 rescale_output_factor=output_scale_factor,696 eps=resnet_eps,697 norm_num_groups=resnet_groups,698 )699 )700 701 self.attentions = nn.ModuleList(attentions)702 self.resnets = nn.ModuleList(resnets)703 704 if add_downsample:705 self.downsamplers = nn.ModuleList(706 [707 Downsample2D(708 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"709 )710 ]711 )712 else:713 self.downsamplers = None714 715 def forward(self, hidden_states, temb=None):716 output_states = ()717 718 for resnet, attn in zip(self.resnets, self.attentions):719 hidden_states = resnet(hidden_states, temb)720 hidden_states = attn(hidden_states)721 output_states += (hidden_states,)722 723 if self.downsamplers is not None:724 for downsampler in self.downsamplers:725 hidden_states = downsampler(hidden_states)726 727 output_states += (hidden_states,)728 729 return hidden_states, output_states730 731 732class CrossAttnDownBlock2D(nn.Module):733 def __init__(734 self,735 in_channels: int,736 out_channels: int,737 temb_channels: int,738 dropout: float = 0.0,739 num_layers: int = 1,740 resnet_eps: float = 1e-6,741 resnet_time_scale_shift: str = "default",742 resnet_act_fn: str = "swish",743 resnet_groups: int = 32,744 resnet_pre_norm: bool = True,745 attn_num_head_channels=1,746 cross_attention_dim=1280,747 output_scale_factor=1.0,748 downsample_padding=1,749 add_downsample=True,750 dual_cross_attention=False,751 use_linear_projection=False,752 only_cross_attention=False,753 upcast_attention=False,754 ):755 super().__init__()756 resnets = []757 attentions = []758 759 self.has_cross_attention = True760 self.attn_num_head_channels = attn_num_head_channels761 762 for i in range(num_layers):763 in_channels = in_channels if i == 0 else out_channels764 resnets.append(765 ResnetBlock2D(766 in_channels=in_channels,767 out_channels=out_channels,768 temb_channels=temb_channels,769 eps=resnet_eps,770 groups=resnet_groups,771 dropout=dropout,772 time_embedding_norm=resnet_time_scale_shift,773 non_linearity=resnet_act_fn,774 output_scale_factor=output_scale_factor,775 pre_norm=resnet_pre_norm,776 )777 )778 if not dual_cross_attention:779 attentions.append(780 Transformer2DModel(781 attn_num_head_channels,782 out_channels // attn_num_head_channels,783 in_channels=out_channels,784 num_layers=1,785 cross_attention_dim=cross_attention_dim,786 norm_num_groups=resnet_groups,787 use_linear_projection=use_linear_projection,788 only_cross_attention=only_cross_attention,789 upcast_attention=upcast_attention,790 )791 )792 else:793 attentions.append(794 DualTransformer2DModel(795 attn_num_head_channels,796 out_channels // attn_num_head_channels,797 in_channels=out_channels,798 num_layers=1,799 cross_attention_dim=cross_attention_dim,800 norm_num_groups=resnet_groups,801 )802 )803 self.attentions = nn.ModuleList(attentions)804 self.resnets = nn.ModuleList(resnets)805 806 if add_downsample:807 self.downsamplers = nn.ModuleList(808 [809 Downsample2D(810 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"811 )812 ]813 )814 else:815 self.downsamplers = None816 817 self.gradient_checkpointing = False818 819 def forward(820 self,821 hidden_states: torch.FloatTensor,822 temb: Optional[torch.FloatTensor] = None,823 encoder_hidden_states: Optional[torch.FloatTensor] = None,824 attention_mask: Optional[torch.FloatTensor] = None,825 cross_attention_kwargs: Optional[Dict[str, Any]] = None,826 encoder_attention_mask: Optional[torch.FloatTensor] = None,827 ):828 output_states = ()829 830 for resnet, attn in zip(self.resnets, self.attentions):831 if self.training and self.gradient_checkpointing:832 833 def create_custom_forward(module, return_dict=None):834 def custom_forward(*inputs):835 if return_dict is not None:836 return module(*inputs, return_dict=return_dict)837 else:838 return module(*inputs)839 840 return custom_forward841 842 hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)843 hidden_states = torch.utils.checkpoint.checkpoint(844 create_custom_forward(attn, return_dict=False),845 hidden_states,846 encoder_hidden_states,847 None, # timestep848 None, # class_labels849 cross_attention_kwargs,850 attention_mask,851 encoder_attention_mask,852 )[0]853 else:854 hidden_states = resnet(hidden_states, temb)855 hidden_states = attn(856 hidden_states,857 encoder_hidden_states=encoder_hidden_states,858 cross_attention_kwargs=cross_attention_kwargs,859 attention_mask=attention_mask,860 encoder_attention_mask=encoder_attention_mask,861 ).sample862 863 output_states += (hidden_states,)864 865 if self.downsamplers is not None:866 for downsampler in self.downsamplers:867 hidden_states = downsampler(hidden_states)868 869 output_states += (hidden_states,)870 871 return hidden_states, output_states872 873 874class DownBlock2D(nn.Module):875 def __init__(876 self,877 in_channels: int,878 out_channels: int,879 temb_channels: int,880 dropout: float = 0.0,881 num_layers: int = 1,882 resnet_eps: float = 1e-6,883 resnet_time_scale_shift: str = "default",884 resnet_act_fn: str = "swish",885 resnet_groups: int = 32,886 resnet_pre_norm: bool = True,887 output_scale_factor=1.0,888 add_downsample=True,889 downsample_padding=1,890 ):891 super().__init__()892 resnets = []893 894 for i in range(num_layers):895 in_channels = in_channels if i == 0 else out_channels896 resnets.append(897 ResnetBlock2D(898 in_channels=in_channels,899 out_channels=out_channels,900 temb_channels=temb_channels,901 eps=resnet_eps,902 groups=resnet_groups,903 dropout=dropout,904 time_embedding_norm=resnet_time_scale_shift,905 non_linearity=resnet_act_fn,906 output_scale_factor=output_scale_factor,907 pre_norm=resnet_pre_norm,908 )909 )910 911 self.resnets = nn.ModuleList(resnets)912 913 if add_downsample:914 self.downsamplers = nn.ModuleList(915 [916 Downsample2D(917 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"918 )919 ]920 )921 else:922 self.downsamplers = None923 924 self.gradient_checkpointing = False925 926 def forward(self, hidden_states, temb=None):927 output_states = ()928 929 for resnet in self.resnets:930 if self.training and self.gradient_checkpointing:931 932 def create_custom_forward(module):933 def custom_forward(*inputs):934 return module(*inputs)935 936 return custom_forward937 938 hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)939 else:940 hidden_states = resnet(hidden_states, temb)941 942 output_states += (hidden_states,)943 944 if self.downsamplers is not None:945 for downsampler in self.downsamplers:946 hidden_states = downsampler(hidden_states)947 948 output_states += (hidden_states,)949 950 return hidden_states, output_states951 952 953class DownEncoderBlock2D(nn.Module):954 def __init__(955 self,956 in_channels: int,957 out_channels: int,958 dropout: float = 0.0,959 num_layers: int = 1,960 resnet_eps: float = 1e-6,961 resnet_time_scale_shift: str = "default",962 resnet_act_fn: str = "swish",963 resnet_groups: int = 32,964 resnet_pre_norm: bool = True,965 output_scale_factor=1.0,966 add_downsample=True,967 downsample_padding=1,968 ):969 super().__init__()970 resnets = []971 972 for i in range(num_layers):973 in_channels = in_channels if i == 0 else out_channels974 resnets.append(975 ResnetBlock2D(976 in_channels=in_channels,977 out_channels=out_channels,978 temb_channels=None,979 eps=resnet_eps,980 groups=resnet_groups,981 dropout=dropout,982 time_embedding_norm=resnet_time_scale_shift,983 non_linearity=resnet_act_fn,984 output_scale_factor=output_scale_factor,985 pre_norm=resnet_pre_norm,986 )987 )988 989 self.resnets = nn.ModuleList(resnets)990 991 if add_downsample:992 self.downsamplers = nn.ModuleList(993 [994 Downsample2D(995 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"996 )997 ]998 )999 else:1000 self.downsamplers = None1001 1002 def forward(self, hidden_states):1003 for resnet in self.resnets:1004 hidden_states = resnet(hidden_states, temb=None)1005 1006 if self.downsamplers is not None:1007 for downsampler in self.downsamplers:1008 hidden_states = downsampler(hidden_states)1009 1010 return hidden_states1011 1012 1013class AttnDownEncoderBlock2D(nn.Module):1014 def __init__(1015 self,1016 in_channels: int,1017 out_channels: int,1018 dropout: float = 0.0,1019 num_layers: int = 1,1020 resnet_eps: float = 1e-6,1021 resnet_time_scale_shift: str = "default",1022 resnet_act_fn: str = "swish",1023 resnet_groups: int = 32,1024 resnet_pre_norm: bool = True,1025 attn_num_head_channels=1,1026 output_scale_factor=1.0,1027 add_downsample=True,1028 downsample_padding=1,1029 ):1030 super().__init__()1031 resnets = []1032 attentions = []1033 1034 for i in range(num_layers):1035 in_channels = in_channels if i == 0 else out_channels1036 resnets.append(1037 ResnetBlock2D(1038 in_channels=in_channels,1039 out_channels=out_channels,1040 temb_channels=None,1041 eps=resnet_eps,1042 groups=resnet_groups,1043 dropout=dropout,1044 time_embedding_norm=resnet_time_scale_shift,1045 non_linearity=resnet_act_fn,1046 output_scale_factor=output_scale_factor,1047 pre_norm=resnet_pre_norm,1048 )1049 )1050 attentions.append(1051 AttentionBlock(1052 out_channels,1053 num_head_channels=attn_num_head_channels,1054 rescale_output_factor=output_scale_factor,1055 eps=resnet_eps,1056 norm_num_groups=resnet_groups,1057 )1058 )1059 1060 self.attentions = nn.ModuleList(attentions)1061 self.resnets = nn.ModuleList(resnets)1062 1063 if add_downsample:1064 self.downsamplers = nn.ModuleList(1065 [1066 Downsample2D(1067 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"1068 )1069 ]1070 )1071 else:1072 self.downsamplers = None1073 1074 def forward(self, hidden_states):1075 for resnet, attn in zip(self.resnets, self.attentions):1076 hidden_states = resnet(hidden_states, temb=None)1077 hidden_states = attn(hidden_states)1078 1079 if self.downsamplers is not None:1080 for downsampler in self.downsamplers:1081 hidden_states = downsampler(hidden_states)1082 1083 return hidden_states1084 1085 1086class AttnSkipDownBlock2D(nn.Module):1087 def __init__(1088 self,1089 in_channels: int,1090 out_channels: int,1091 temb_channels: int,1092 dropout: float = 0.0,1093 num_layers: int = 1,1094 resnet_eps: float = 1e-6,1095 resnet_time_scale_shift: str = "default",1096 resnet_act_fn: str = "swish",1097 resnet_pre_norm: bool = True,1098 attn_num_head_channels=1,1099 output_scale_factor=np.sqrt(2.0),1100 downsample_padding=1,1101 add_downsample=True,1102 ):1103 super().__init__()1104 self.attentions = nn.ModuleList([])1105 self.resnets = nn.ModuleList([])1106 1107 for i in range(num_layers):1108 in_channels = in_channels if i == 0 else out_channels1109 self.resnets.append(1110 ResnetBlock2D(1111 in_channels=in_channels,1112 out_channels=out_channels,1113 temb_channels=temb_channels,1114 eps=resnet_eps,1115 groups=min(in_channels // 4, 32),1116 groups_out=min(out_channels // 4, 32),1117 dropout=dropout,1118 time_embedding_norm=resnet_time_scale_shift,1119 non_linearity=resnet_act_fn,1120 output_scale_factor=output_scale_factor,1121 pre_norm=resnet_pre_norm,1122 )1123 )1124 self.attentions.append(1125 AttentionBlock(1126 out_channels,1127 num_head_channels=attn_num_head_channels,1128 rescale_output_factor=output_scale_factor,1129 eps=resnet_eps,1130 )1131 )1132 1133 if add_downsample:1134 self.resnet_down = ResnetBlock2D(1135 in_channels=out_channels,1136 out_channels=out_channels,1137 temb_channels=temb_channels,1138 eps=resnet_eps,1139 groups=min(out_channels // 4, 32),1140 dropout=dropout,1141 time_embedding_norm=resnet_time_scale_shift,1142 non_linearity=resnet_act_fn,1143 output_scale_factor=output_scale_factor,1144 pre_norm=resnet_pre_norm,1145 use_in_shortcut=True,1146 down=True,1147 kernel="fir",1148 )1149 self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])1150 self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))1151 else:1152 self.resnet_down = None1153 self.downsamplers = None1154 self.skip_conv = None1155 1156 def forward(self, hidden_states, temb=None, skip_sample=None):1157 output_states = ()1158 1159 for resnet, attn in zip(self.resnets, self.attentions):1160 hidden_states = resnet(hidden_states, temb)1161 hidden_states = attn(hidden_states)1162 output_states += (hidden_states,)1163 1164 if self.downsamplers is not None:1165 hidden_states = self.resnet_down(hidden_states, temb)1166 for downsampler in self.downsamplers:1167 skip_sample = downsampler(skip_sample)1168 1169 hidden_states = self.skip_conv(skip_sample) + hidden_states1170 1171 output_states += (hidden_states,)1172 1173 return hidden_states, output_states, skip_sample1174 1175 1176class SkipDownBlock2D(nn.Module):1177 def __init__(1178 self,1179 in_channels: int,1180 out_channels: int,1181 temb_channels: int,1182 dropout: float = 0.0,1183 num_layers: int = 1,1184 resnet_eps: float = 1e-6,1185 resnet_time_scale_shift: str = "default",1186 resnet_act_fn: str = "swish",1187 resnet_pre_norm: bool = True,1188 output_scale_factor=np.sqrt(2.0),1189 add_downsample=True,1190 downsample_padding=1,1191 ):1192 super().__init__()1193 self.resnets = nn.ModuleList([])1194 1195 for i in range(num_layers):1196 in_channels = in_channels if i == 0 else out_channels1197 self.resnets.append(1198 ResnetBlock2D(1199 in_channels=in_channels,1200 out_channels=out_channels,