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.14 15import torch16from torch import nn17 18from .resnet import Downsample2D, ResnetBlock2D, TemporalConvLayer, Upsample2D19from .transformer_2d import Transformer2DModel20from .transformer_temporal import TransformerTemporalModel21 22 23def get_down_block(24 down_block_type,25 num_layers,26 in_channels,27 out_channels,28 temb_channels,29 add_downsample,30 resnet_eps,31 resnet_act_fn,32 attn_num_head_channels,33 resnet_groups=None,34 cross_attention_dim=None,35 downsample_padding=None,36 dual_cross_attention=False,37 use_linear_projection=True,38 only_cross_attention=False,39 upcast_attention=False,40 resnet_time_scale_shift="default",41):42 if down_block_type == "DownBlock3D":43 return DownBlock3D(44 num_layers=num_layers,45 in_channels=in_channels,46 out_channels=out_channels,47 temb_channels=temb_channels,48 add_downsample=add_downsample,49 resnet_eps=resnet_eps,50 resnet_act_fn=resnet_act_fn,51 resnet_groups=resnet_groups,52 downsample_padding=downsample_padding,53 resnet_time_scale_shift=resnet_time_scale_shift,54 )55 elif down_block_type == "CrossAttnDownBlock3D":56 if cross_attention_dim is None:57 raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock3D")58 return CrossAttnDownBlock3D(59 num_layers=num_layers,60 in_channels=in_channels,61 out_channels=out_channels,62 temb_channels=temb_channels,63 add_downsample=add_downsample,64 resnet_eps=resnet_eps,65 resnet_act_fn=resnet_act_fn,66 resnet_groups=resnet_groups,67 downsample_padding=downsample_padding,68 cross_attention_dim=cross_attention_dim,69 attn_num_head_channels=attn_num_head_channels,70 dual_cross_attention=dual_cross_attention,71 use_linear_projection=use_linear_projection,72 only_cross_attention=only_cross_attention,73 upcast_attention=upcast_attention,74 resnet_time_scale_shift=resnet_time_scale_shift,75 )76 raise ValueError(f"{down_block_type} does not exist.")77 78 79def get_up_block(80 up_block_type,81 num_layers,82 in_channels,83 out_channels,84 prev_output_channel,85 temb_channels,86 add_upsample,87 resnet_eps,88 resnet_act_fn,89 attn_num_head_channels,90 resnet_groups=None,91 cross_attention_dim=None,92 dual_cross_attention=False,93 use_linear_projection=True,94 only_cross_attention=False,95 upcast_attention=False,96 resnet_time_scale_shift="default",97):98 if up_block_type == "UpBlock3D":99 return UpBlock3D(100 num_layers=num_layers,101 in_channels=in_channels,102 out_channels=out_channels,103 prev_output_channel=prev_output_channel,104 temb_channels=temb_channels,105 add_upsample=add_upsample,106 resnet_eps=resnet_eps,107 resnet_act_fn=resnet_act_fn,108 resnet_groups=resnet_groups,109 resnet_time_scale_shift=resnet_time_scale_shift,110 )111 elif up_block_type == "CrossAttnUpBlock3D":112 if cross_attention_dim is None:113 raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock3D")114 return CrossAttnUpBlock3D(115 num_layers=num_layers,116 in_channels=in_channels,117 out_channels=out_channels,118 prev_output_channel=prev_output_channel,119 temb_channels=temb_channels,120 add_upsample=add_upsample,121 resnet_eps=resnet_eps,122 resnet_act_fn=resnet_act_fn,123 resnet_groups=resnet_groups,124 cross_attention_dim=cross_attention_dim,125 attn_num_head_channels=attn_num_head_channels,126 dual_cross_attention=dual_cross_attention,127 use_linear_projection=use_linear_projection,128 only_cross_attention=only_cross_attention,129 upcast_attention=upcast_attention,130 resnet_time_scale_shift=resnet_time_scale_shift,131 )132 raise ValueError(f"{up_block_type} does not exist.")133 134 135class UNetMidBlock3DCrossAttn(nn.Module):136 def __init__(137 self,138 in_channels: int,139 temb_channels: int,140 dropout: float = 0.0,141 num_layers: int = 1,142 resnet_eps: float = 1e-6,143 resnet_time_scale_shift: str = "default",144 resnet_act_fn: str = "swish",145 resnet_groups: int = 32,146 resnet_pre_norm: bool = True,147 attn_num_head_channels=1,148 output_scale_factor=1.0,149 cross_attention_dim=1280,150 dual_cross_attention=False,151 use_linear_projection=True,152 upcast_attention=False,153 ):154 super().__init__()155 156 self.has_cross_attention = True157 self.attn_num_head_channels = attn_num_head_channels158 resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)159 160 # there is always at least one resnet161 resnets = [162 ResnetBlock2D(163 in_channels=in_channels,164 out_channels=in_channels,165 temb_channels=temb_channels,166 eps=resnet_eps,167 groups=resnet_groups,168 dropout=dropout,169 time_embedding_norm=resnet_time_scale_shift,170 non_linearity=resnet_act_fn,171 output_scale_factor=output_scale_factor,172 pre_norm=resnet_pre_norm,173 )174 ]175 temp_convs = [176 TemporalConvLayer(177 in_channels,178 in_channels,179 dropout=0.1,180 )181 ]182 attentions = []183 temp_attentions = []184 185 for _ in range(num_layers):186 attentions.append(187 Transformer2DModel(188 in_channels // attn_num_head_channels,189 attn_num_head_channels,190 in_channels=in_channels,191 num_layers=1,192 cross_attention_dim=cross_attention_dim,193 norm_num_groups=resnet_groups,194 use_linear_projection=use_linear_projection,195 upcast_attention=upcast_attention,196 )197 )198 temp_attentions.append(199 TransformerTemporalModel(200 in_channels // attn_num_head_channels,201 attn_num_head_channels,202 in_channels=in_channels,203 num_layers=1,204 cross_attention_dim=cross_attention_dim,205 norm_num_groups=resnet_groups,206 )207 )208 resnets.append(209 ResnetBlock2D(210 in_channels=in_channels,211 out_channels=in_channels,212 temb_channels=temb_channels,213 eps=resnet_eps,214 groups=resnet_groups,215 dropout=dropout,216 time_embedding_norm=resnet_time_scale_shift,217 non_linearity=resnet_act_fn,218 output_scale_factor=output_scale_factor,219 pre_norm=resnet_pre_norm,220 )221 )222 temp_convs.append(223 TemporalConvLayer(224 in_channels,225 in_channels,226 dropout=0.1,227 )228 )229 230 self.resnets = nn.ModuleList(resnets)231 self.temp_convs = nn.ModuleList(temp_convs)232 self.attentions = nn.ModuleList(attentions)233 self.temp_attentions = nn.ModuleList(temp_attentions)234 235 def forward(236 self,237 hidden_states,238 temb=None,239 encoder_hidden_states=None,240 attention_mask=None,241 num_frames=1,242 cross_attention_kwargs=None,243 ):244 hidden_states = self.resnets[0](hidden_states, temb)245 hidden_states = self.temp_convs[0](hidden_states, num_frames=num_frames)246 for attn, temp_attn, resnet, temp_conv in zip(247 self.attentions, self.temp_attentions, self.resnets[1:], self.temp_convs[1:]248 ):249 hidden_states = attn(250 hidden_states,251 encoder_hidden_states=encoder_hidden_states,252 cross_attention_kwargs=cross_attention_kwargs,253 ).sample254 hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample255 hidden_states = resnet(hidden_states, temb)256 hidden_states = temp_conv(hidden_states, num_frames=num_frames)257 258 return hidden_states259 260 261class CrossAttnDownBlock3D(nn.Module):262 def __init__(263 self,264 in_channels: int,265 out_channels: int,266 temb_channels: int,267 dropout: float = 0.0,268 num_layers: int = 1,269 resnet_eps: float = 1e-6,270 resnet_time_scale_shift: str = "default",271 resnet_act_fn: str = "swish",272 resnet_groups: int = 32,273 resnet_pre_norm: bool = True,274 attn_num_head_channels=1,275 cross_attention_dim=1280,276 output_scale_factor=1.0,277 downsample_padding=1,278 add_downsample=True,279 dual_cross_attention=False,280 use_linear_projection=False,281 only_cross_attention=False,282 upcast_attention=False,283 ):284 super().__init__()285 resnets = []286 attentions = []287 temp_attentions = []288 temp_convs = []289 290 self.has_cross_attention = True291 self.attn_num_head_channels = attn_num_head_channels292 293 for i in range(num_layers):294 in_channels = in_channels if i == 0 else out_channels295 resnets.append(296 ResnetBlock2D(297 in_channels=in_channels,298 out_channels=out_channels,299 temb_channels=temb_channels,300 eps=resnet_eps,301 groups=resnet_groups,302 dropout=dropout,303 time_embedding_norm=resnet_time_scale_shift,304 non_linearity=resnet_act_fn,305 output_scale_factor=output_scale_factor,306 pre_norm=resnet_pre_norm,307 )308 )309 temp_convs.append(310 TemporalConvLayer(311 out_channels,312 out_channels,313 dropout=0.1,314 )315 )316 attentions.append(317 Transformer2DModel(318 out_channels // attn_num_head_channels,319 attn_num_head_channels,320 in_channels=out_channels,321 num_layers=1,322 cross_attention_dim=cross_attention_dim,323 norm_num_groups=resnet_groups,324 use_linear_projection=use_linear_projection,325 only_cross_attention=only_cross_attention,326 upcast_attention=upcast_attention,327 )328 )329 temp_attentions.append(330 TransformerTemporalModel(331 out_channels // attn_num_head_channels,332 attn_num_head_channels,333 in_channels=out_channels,334 num_layers=1,335 cross_attention_dim=cross_attention_dim,336 norm_num_groups=resnet_groups,337 )338 )339 self.resnets = nn.ModuleList(resnets)340 self.temp_convs = nn.ModuleList(temp_convs)341 self.attentions = nn.ModuleList(attentions)342 self.temp_attentions = nn.ModuleList(temp_attentions)343 344 if add_downsample:345 self.downsamplers = nn.ModuleList(346 [347 Downsample2D(348 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"349 )350 ]351 )352 else:353 self.downsamplers = None354 355 self.gradient_checkpointing = False356 357 def forward(358 self,359 hidden_states,360 temb=None,361 encoder_hidden_states=None,362 attention_mask=None,363 num_frames=1,364 cross_attention_kwargs=None,365 ):366 # TODO(Patrick, William) - attention mask is not used367 output_states = ()368 369 for resnet, temp_conv, attn, temp_attn in zip(370 self.resnets, self.temp_convs, self.attentions, self.temp_attentions371 ):372 hidden_states = resnet(hidden_states, temb)373 hidden_states = temp_conv(hidden_states, num_frames=num_frames)374 hidden_states = attn(375 hidden_states,376 encoder_hidden_states=encoder_hidden_states,377 cross_attention_kwargs=cross_attention_kwargs,378 ).sample379 hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample380 381 output_states += (hidden_states,)382 383 if self.downsamplers is not None:384 for downsampler in self.downsamplers:385 hidden_states = downsampler(hidden_states)386 387 output_states += (hidden_states,)388 389 return hidden_states, output_states390 391 392class DownBlock3D(nn.Module):393 def __init__(394 self,395 in_channels: int,396 out_channels: int,397 temb_channels: int,398 dropout: float = 0.0,399 num_layers: int = 1,400 resnet_eps: float = 1e-6,401 resnet_time_scale_shift: str = "default",402 resnet_act_fn: str = "swish",403 resnet_groups: int = 32,404 resnet_pre_norm: bool = True,405 output_scale_factor=1.0,406 add_downsample=True,407 downsample_padding=1,408 ):409 super().__init__()410 resnets = []411 temp_convs = []412 413 for i in range(num_layers):414 in_channels = in_channels if i == 0 else out_channels415 resnets.append(416 ResnetBlock2D(417 in_channels=in_channels,418 out_channels=out_channels,419 temb_channels=temb_channels,420 eps=resnet_eps,421 groups=resnet_groups,422 dropout=dropout,423 time_embedding_norm=resnet_time_scale_shift,424 non_linearity=resnet_act_fn,425 output_scale_factor=output_scale_factor,426 pre_norm=resnet_pre_norm,427 )428 )429 temp_convs.append(430 TemporalConvLayer(431 out_channels,432 out_channels,433 dropout=0.1,434 )435 )436 437 self.resnets = nn.ModuleList(resnets)438 self.temp_convs = nn.ModuleList(temp_convs)439 440 if add_downsample:441 self.downsamplers = nn.ModuleList(442 [443 Downsample2D(444 out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"445 )446 ]447 )448 else:449 self.downsamplers = None450 451 self.gradient_checkpointing = False452 453 def forward(self, hidden_states, temb=None, num_frames=1):454 output_states = ()455 456 for resnet, temp_conv in zip(self.resnets, self.temp_convs):457 hidden_states = resnet(hidden_states, temb)458 hidden_states = temp_conv(hidden_states, num_frames=num_frames)459 460 output_states += (hidden_states,)461 462 if self.downsamplers is not None:463 for downsampler in self.downsamplers:464 hidden_states = downsampler(hidden_states)465 466 output_states += (hidden_states,)467 468 return hidden_states, output_states469 470 471class CrossAttnUpBlock3D(nn.Module):472 def __init__(473 self,474 in_channels: int,475 out_channels: int,476 prev_output_channel: int,477 temb_channels: int,478 dropout: float = 0.0,479 num_layers: int = 1,480 resnet_eps: float = 1e-6,481 resnet_time_scale_shift: str = "default",482 resnet_act_fn: str = "swish",483 resnet_groups: int = 32,484 resnet_pre_norm: bool = True,485 attn_num_head_channels=1,486 cross_attention_dim=1280,487 output_scale_factor=1.0,488 add_upsample=True,489 dual_cross_attention=False,490 use_linear_projection=False,491 only_cross_attention=False,492 upcast_attention=False,493 ):494 super().__init__()495 resnets = []496 temp_convs = []497 attentions = []498 temp_attentions = []499 500 self.has_cross_attention = True501 self.attn_num_head_channels = attn_num_head_channels502 503 for i in range(num_layers):504 res_skip_channels = in_channels if (i == num_layers - 1) else out_channels505 resnet_in_channels = prev_output_channel if i == 0 else out_channels506 507 resnets.append(508 ResnetBlock2D(509 in_channels=resnet_in_channels + res_skip_channels,510 out_channels=out_channels,511 temb_channels=temb_channels,512 eps=resnet_eps,513 groups=resnet_groups,514 dropout=dropout,515 time_embedding_norm=resnet_time_scale_shift,516 non_linearity=resnet_act_fn,517 output_scale_factor=output_scale_factor,518 pre_norm=resnet_pre_norm,519 )520 )521 temp_convs.append(522 TemporalConvLayer(523 out_channels,524 out_channels,525 dropout=0.1,526 )527 )528 attentions.append(529 Transformer2DModel(530 out_channels // attn_num_head_channels,531 attn_num_head_channels,532 in_channels=out_channels,533 num_layers=1,534 cross_attention_dim=cross_attention_dim,535 norm_num_groups=resnet_groups,536 use_linear_projection=use_linear_projection,537 only_cross_attention=only_cross_attention,538 upcast_attention=upcast_attention,539 )540 )541 temp_attentions.append(542 TransformerTemporalModel(543 out_channels // attn_num_head_channels,544 attn_num_head_channels,545 in_channels=out_channels,546 num_layers=1,547 cross_attention_dim=cross_attention_dim,548 norm_num_groups=resnet_groups,549 )550 )551 self.resnets = nn.ModuleList(resnets)552 self.temp_convs = nn.ModuleList(temp_convs)553 self.attentions = nn.ModuleList(attentions)554 self.temp_attentions = nn.ModuleList(temp_attentions)555 556 if add_upsample:557 self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])558 else:559 self.upsamplers = None560 561 self.gradient_checkpointing = False562 563 def forward(564 self,565 hidden_states,566 res_hidden_states_tuple,567 temb=None,568 encoder_hidden_states=None,569 upsample_size=None,570 attention_mask=None,571 num_frames=1,572 cross_attention_kwargs=None,573 ):574 # TODO(Patrick, William) - attention mask is not used575 for resnet, temp_conv, attn, temp_attn in zip(576 self.resnets, self.temp_convs, self.attentions, self.temp_attentions577 ):578 # pop res hidden states579 res_hidden_states = res_hidden_states_tuple[-1]580 res_hidden_states_tuple = res_hidden_states_tuple[:-1]581 hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)582 583 hidden_states = resnet(hidden_states, temb)584 hidden_states = temp_conv(hidden_states, num_frames=num_frames)585 hidden_states = attn(586 hidden_states,587 encoder_hidden_states=encoder_hidden_states,588 cross_attention_kwargs=cross_attention_kwargs,589 ).sample590 hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample591 592 if self.upsamplers is not None:593 for upsampler in self.upsamplers:594 hidden_states = upsampler(hidden_states, upsample_size)595 596 return hidden_states597 598 599class UpBlock3D(nn.Module):600 def __init__(601 self,602 in_channels: int,603 prev_output_channel: int,604 out_channels: int,605 temb_channels: int,606 dropout: float = 0.0,607 num_layers: int = 1,608 resnet_eps: float = 1e-6,609 resnet_time_scale_shift: str = "default",610 resnet_act_fn: str = "swish",611 resnet_groups: int = 32,612 resnet_pre_norm: bool = True,613 output_scale_factor=1.0,614 add_upsample=True,615 ):616 super().__init__()617 resnets = []618 temp_convs = []619 620 for i in range(num_layers):621 res_skip_channels = in_channels if (i == num_layers - 1) else out_channels622 resnet_in_channels = prev_output_channel if i == 0 else out_channels623 624 resnets.append(625 ResnetBlock2D(626 in_channels=resnet_in_channels + res_skip_channels,627 out_channels=out_channels,628 temb_channels=temb_channels,629 eps=resnet_eps,630 groups=resnet_groups,631 dropout=dropout,632 time_embedding_norm=resnet_time_scale_shift,633 non_linearity=resnet_act_fn,634 output_scale_factor=output_scale_factor,635 pre_norm=resnet_pre_norm,636 )637 )638 temp_convs.append(639 TemporalConvLayer(640 out_channels,641 out_channels,642 dropout=0.1,643 )644 )645 646 self.resnets = nn.ModuleList(resnets)647 self.temp_convs = nn.ModuleList(temp_convs)648 649 if add_upsample:650 self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])651 else:652 self.upsamplers = None653 654 self.gradient_checkpointing = False655 656 def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, num_frames=1):657 for resnet, temp_conv in zip(self.resnets, self.temp_convs):658 # pop res hidden states659 res_hidden_states = res_hidden_states_tuple[-1]660 res_hidden_states_tuple = res_hidden_states_tuple[:-1]661 hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)662 663 hidden_states = resnet(hidden_states, temb)664 hidden_states = temp_conv(hidden_states, num_frames=num_frames)665 666 if self.upsamplers is not None:667 for upsampler in self.upsamplers:668 hidden_states = upsampler(hidden_states, upsample_size)669 670 return hidden_states671