declare-lab/tango2
92
1# Copyright 2023 The HuggingFace Team. All rights reserved.2# `TemporalConvLayer` Copyright 2023 Alibaba DAMO-VILAB, The ModelScope Team and The HuggingFace Team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16from functools import partial17from typing import Optional18 19import torch20import torch.nn as nn21import torch.nn.functional as F22 23from .attention import AdaGroupNorm24 25 26class Upsample1D(nn.Module):27 """28 An upsampling layer with an optional convolution.29 30 Parameters:31 channels: channels in the inputs and outputs.32 use_conv: a bool determining if a convolution is applied.33 use_conv_transpose:34 out_channels:35 """36 37 def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"):38 super().__init__()39 self.channels = channels40 self.out_channels = out_channels or channels41 self.use_conv = use_conv42 self.use_conv_transpose = use_conv_transpose43 self.name = name44 45 self.conv = None46 if use_conv_transpose:47 self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1)48 elif use_conv:49 self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1)50 51 def forward(self, x):52 assert x.shape[1] == self.channels53 if self.use_conv_transpose:54 return self.conv(x)55 56 x = F.interpolate(x, scale_factor=2.0, mode="nearest")57 58 if self.use_conv:59 x = self.conv(x)60 61 return x62 63 64class Downsample1D(nn.Module):65 """66 A downsampling layer with an optional convolution.67 68 Parameters:69 channels: channels in the inputs and outputs.70 use_conv: a bool determining if a convolution is applied.71 out_channels:72 padding:73 """74 75 def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"):76 super().__init__()77 self.channels = channels78 self.out_channels = out_channels or channels79 self.use_conv = use_conv80 self.padding = padding81 stride = 282 self.name = name83 84 if use_conv:85 self.conv = nn.Conv1d(self.channels, self.out_channels, 3, stride=stride, padding=padding)86 else:87 assert self.channels == self.out_channels88 self.conv = nn.AvgPool1d(kernel_size=stride, stride=stride)89 90 def forward(self, x):91 assert x.shape[1] == self.channels92 return self.conv(x)93 94 95class Upsample2D(nn.Module):96 """97 An upsampling layer with an optional convolution.98 99 Parameters:100 channels: channels in the inputs and outputs.101 use_conv: a bool determining if a convolution is applied.102 use_conv_transpose:103 out_channels:104 """105 106 def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"):107 super().__init__()108 self.channels = channels109 self.out_channels = out_channels or channels110 self.use_conv = use_conv111 self.use_conv_transpose = use_conv_transpose112 self.name = name113 114 conv = None115 if use_conv_transpose:116 conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1)117 elif use_conv:118 conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1)119 120 # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed121 if name == "conv":122 self.conv = conv123 else:124 self.Conv2d_0 = conv125 126 def forward(self, hidden_states, output_size=None):127 assert hidden_states.shape[1] == self.channels128 129 if self.use_conv_transpose:130 return self.conv(hidden_states)131 132 # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16133 # TODO(Suraj): Remove this cast once the issue is fixed in PyTorch134 # https://github.com/pytorch/pytorch/issues/86679135 dtype = hidden_states.dtype136 if dtype == torch.bfloat16:137 hidden_states = hidden_states.to(torch.float32)138 139 # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984140 if hidden_states.shape[0] >= 64:141 hidden_states = hidden_states.contiguous()142 143 # if `output_size` is passed we force the interpolation output144 # size and do not make use of `scale_factor=2`145 if output_size is None:146 hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")147 else:148 hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")149 150 # If the input is bfloat16, we cast back to bfloat16151 if dtype == torch.bfloat16:152 hidden_states = hidden_states.to(dtype)153 154 # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed155 if self.use_conv:156 if self.name == "conv":157 hidden_states = self.conv(hidden_states)158 else:159 hidden_states = self.Conv2d_0(hidden_states)160 161 return hidden_states162 163 164class Downsample2D(nn.Module):165 """166 A downsampling layer with an optional convolution.167 168 Parameters:169 channels: channels in the inputs and outputs.170 use_conv: a bool determining if a convolution is applied.171 out_channels:172 padding:173 """174 175 def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"):176 super().__init__()177 self.channels = channels178 self.out_channels = out_channels or channels179 self.use_conv = use_conv180 self.padding = padding181 stride = 2182 self.name = name183 184 if use_conv:185 conv = nn.Conv2d(self.channels, self.out_channels, 3, stride=stride, padding=padding)186 else:187 assert self.channels == self.out_channels188 conv = nn.AvgPool2d(kernel_size=stride, stride=stride)189 190 # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed191 if name == "conv":192 self.Conv2d_0 = conv193 self.conv = conv194 elif name == "Conv2d_0":195 self.conv = conv196 else:197 self.conv = conv198 199 def forward(self, hidden_states):200 assert hidden_states.shape[1] == self.channels201 if self.use_conv and self.padding == 0:202 pad = (0, 1, 0, 1)203 hidden_states = F.pad(hidden_states, pad, mode="constant", value=0)204 205 assert hidden_states.shape[1] == self.channels206 hidden_states = self.conv(hidden_states)207 208 return hidden_states209 210 211class FirUpsample2D(nn.Module):212 def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel=(1, 3, 3, 1)):213 super().__init__()214 out_channels = out_channels if out_channels else channels215 if use_conv:216 self.Conv2d_0 = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)217 self.use_conv = use_conv218 self.fir_kernel = fir_kernel219 self.out_channels = out_channels220 221 def _upsample_2d(self, hidden_states, weight=None, kernel=None, factor=2, gain=1):222 """Fused `upsample_2d()` followed by `Conv2d()`.223 224 Padding is performed only once at the beginning, not between the operations. The fused op is considerably more225 efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of226 arbitrary order.227 228 Args:229 hidden_states: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.230 weight: Weight tensor of the shape `[filterH, filterW, inChannels,231 outChannels]`. Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`.232 kernel: FIR filter of the shape `[firH, firW]` or `[firN]`233 (separable). The default is `[1] * factor`, which corresponds to nearest-neighbor upsampling.234 factor: Integer upsampling factor (default: 2).235 gain: Scaling factor for signal magnitude (default: 1.0).236 237 Returns:238 output: Tensor of the shape `[N, C, H * factor, W * factor]` or `[N, H * factor, W * factor, C]`, and same239 datatype as `hidden_states`.240 """241 242 assert isinstance(factor, int) and factor >= 1243 244 # Setup filter kernel.245 if kernel is None:246 kernel = [1] * factor247 248 # setup kernel249 kernel = torch.tensor(kernel, dtype=torch.float32)250 if kernel.ndim == 1:251 kernel = torch.outer(kernel, kernel)252 kernel /= torch.sum(kernel)253 254 kernel = kernel * (gain * (factor**2))255 256 if self.use_conv:257 convH = weight.shape[2]258 convW = weight.shape[3]259 inC = weight.shape[1]260 261 pad_value = (kernel.shape[0] - factor) - (convW - 1)262 263 stride = (factor, factor)264 # Determine data dimensions.265 output_shape = (266 (hidden_states.shape[2] - 1) * factor + convH,267 (hidden_states.shape[3] - 1) * factor + convW,268 )269 output_padding = (270 output_shape[0] - (hidden_states.shape[2] - 1) * stride[0] - convH,271 output_shape[1] - (hidden_states.shape[3] - 1) * stride[1] - convW,272 )273 assert output_padding[0] >= 0 and output_padding[1] >= 0274 num_groups = hidden_states.shape[1] // inC275 276 # Transpose weights.277 weight = torch.reshape(weight, (num_groups, -1, inC, convH, convW))278 weight = torch.flip(weight, dims=[3, 4]).permute(0, 2, 1, 3, 4)279 weight = torch.reshape(weight, (num_groups * inC, -1, convH, convW))280 281 inverse_conv = F.conv_transpose2d(282 hidden_states, weight, stride=stride, output_padding=output_padding, padding=0283 )284 285 output = upfirdn2d_native(286 inverse_conv,287 torch.tensor(kernel, device=inverse_conv.device),288 pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2 + 1),289 )290 else:291 pad_value = kernel.shape[0] - factor292 output = upfirdn2d_native(293 hidden_states,294 torch.tensor(kernel, device=hidden_states.device),295 up=factor,296 pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2),297 )298 299 return output300 301 def forward(self, hidden_states):302 if self.use_conv:303 height = self._upsample_2d(hidden_states, self.Conv2d_0.weight, kernel=self.fir_kernel)304 height = height + self.Conv2d_0.bias.reshape(1, -1, 1, 1)305 else:306 height = self._upsample_2d(hidden_states, kernel=self.fir_kernel, factor=2)307 308 return height309 310 311class FirDownsample2D(nn.Module):312 def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel=(1, 3, 3, 1)):313 super().__init__()314 out_channels = out_channels if out_channels else channels315 if use_conv:316 self.Conv2d_0 = nn.Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)317 self.fir_kernel = fir_kernel318 self.use_conv = use_conv319 self.out_channels = out_channels320 321 def _downsample_2d(self, hidden_states, weight=None, kernel=None, factor=2, gain=1):322 """Fused `Conv2d()` followed by `downsample_2d()`.323 Padding is performed only once at the beginning, not between the operations. The fused op is considerably more324 efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of325 arbitrary order.326 327 Args:328 hidden_states: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.329 weight:330 Weight tensor of the shape `[filterH, filterW, inChannels, outChannels]`. Grouped convolution can be331 performed by `inChannels = x.shape[0] // numGroups`.332 kernel: FIR filter of the shape `[firH, firW]` or `[firN]` (separable). The default is `[1] *333 factor`, which corresponds to average pooling.334 factor: Integer downsampling factor (default: 2).335 gain: Scaling factor for signal magnitude (default: 1.0).336 337 Returns:338 output: Tensor of the shape `[N, C, H // factor, W // factor]` or `[N, H // factor, W // factor, C]`, and339 same datatype as `x`.340 """341 342 assert isinstance(factor, int) and factor >= 1343 if kernel is None:344 kernel = [1] * factor345 346 # setup kernel347 kernel = torch.tensor(kernel, dtype=torch.float32)348 if kernel.ndim == 1:349 kernel = torch.outer(kernel, kernel)350 kernel /= torch.sum(kernel)351 352 kernel = kernel * gain353 354 if self.use_conv:355 _, _, convH, convW = weight.shape356 pad_value = (kernel.shape[0] - factor) + (convW - 1)357 stride_value = [factor, factor]358 upfirdn_input = upfirdn2d_native(359 hidden_states,360 torch.tensor(kernel, device=hidden_states.device),361 pad=((pad_value + 1) // 2, pad_value // 2),362 )363 output = F.conv2d(upfirdn_input, weight, stride=stride_value, padding=0)364 else:365 pad_value = kernel.shape[0] - factor366 output = upfirdn2d_native(367 hidden_states,368 torch.tensor(kernel, device=hidden_states.device),369 down=factor,370 pad=((pad_value + 1) // 2, pad_value // 2),371 )372 373 return output374 375 def forward(self, hidden_states):376 if self.use_conv:377 downsample_input = self._downsample_2d(hidden_states, weight=self.Conv2d_0.weight, kernel=self.fir_kernel)378 hidden_states = downsample_input + self.Conv2d_0.bias.reshape(1, -1, 1, 1)379 else:380 hidden_states = self._downsample_2d(hidden_states, kernel=self.fir_kernel, factor=2)381 382 return hidden_states383 384 385# downsample/upsample layer used in k-upscaler, might be able to use FirDownsample2D/DirUpsample2D instead386class KDownsample2D(nn.Module):387 def __init__(self, pad_mode="reflect"):388 super().__init__()389 self.pad_mode = pad_mode390 kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]])391 self.pad = kernel_1d.shape[1] // 2 - 1392 self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False)393 394 def forward(self, x):395 x = F.pad(x, (self.pad,) * 4, self.pad_mode)396 weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0], self.kernel.shape[1]])397 indices = torch.arange(x.shape[1], device=x.device)398 weight[indices, indices] = self.kernel.to(weight)399 return F.conv2d(x, weight, stride=2)400 401 402class KUpsample2D(nn.Module):403 def __init__(self, pad_mode="reflect"):404 super().__init__()405 self.pad_mode = pad_mode406 kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]]) * 2407 self.pad = kernel_1d.shape[1] // 2 - 1408 self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False)409 410 def forward(self, x):411 x = F.pad(x, ((self.pad + 1) // 2,) * 4, self.pad_mode)412 weight = x.new_zeros([x.shape[1], x.shape[1], self.kernel.shape[0], self.kernel.shape[1]])413 indices = torch.arange(x.shape[1], device=x.device)414 weight[indices, indices] = self.kernel.to(weight)415 return F.conv_transpose2d(x, weight, stride=2, padding=self.pad * 2 + 1)416 417 418class ResnetBlock2D(nn.Module):419 r"""420 A Resnet block.421 422 Parameters:423 in_channels (`int`): The number of channels in the input.424 out_channels (`int`, *optional*, default to be `None`):425 The number of output channels for the first conv2d layer. If None, same as `in_channels`.426 dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use.427 temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding.428 groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer.429 groups_out (`int`, *optional*, default to None):430 The number of groups to use for the second normalization layer. if set to None, same as `groups`.431 eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization.432 non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use.433 time_embedding_norm (`str`, *optional*, default to `"default"` ): Time scale shift config.434 By default, apply timestep embedding conditioning with a simple shift mechanism. Choose "scale_shift" or435 "ada_group" for a stronger conditioning with scale and shift.436 kernel (`torch.FloatTensor`, optional, default to None): FIR filter, see437 [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`].438 output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output.439 use_in_shortcut (`bool`, *optional*, default to `True`):440 If `True`, add a 1x1 nn.conv2d layer for skip-connection.441 up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer.442 down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer.443 conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the444 `conv_shortcut` output.445 conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output.446 If None, same as `out_channels`.447 """448 449 def __init__(450 self,451 *,452 in_channels,453 out_channels=None,454 conv_shortcut=False,455 dropout=0.0,456 temb_channels=512,457 groups=32,458 groups_out=None,459 pre_norm=True,460 eps=1e-6,461 non_linearity="swish",462 time_embedding_norm="default", # default, scale_shift, ada_group463 kernel=None,464 output_scale_factor=1.0,465 use_in_shortcut=None,466 up=False,467 down=False,468 conv_shortcut_bias: bool = True,469 conv_2d_out_channels: Optional[int] = None,470 ):471 super().__init__()472 self.pre_norm = pre_norm473 self.pre_norm = True474 self.in_channels = in_channels475 out_channels = in_channels if out_channels is None else out_channels476 self.out_channels = out_channels477 self.use_conv_shortcut = conv_shortcut478 self.up = up479 self.down = down480 self.output_scale_factor = output_scale_factor481 self.time_embedding_norm = time_embedding_norm482 483 if groups_out is None:484 groups_out = groups485 486 if self.time_embedding_norm == "ada_group":487 self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps)488 else:489 self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)490 491 self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)492 493 if temb_channels is not None:494 if self.time_embedding_norm == "default":495 self.time_emb_proj = torch.nn.Linear(temb_channels, out_channels)496 elif self.time_embedding_norm == "scale_shift":497 self.time_emb_proj = torch.nn.Linear(temb_channels, 2 * out_channels)498 elif self.time_embedding_norm == "ada_group":499 self.time_emb_proj = None500 else:501 raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")502 else:503 self.time_emb_proj = None504 505 if self.time_embedding_norm == "ada_group":506 self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps)507 else:508 self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)509 510 self.dropout = torch.nn.Dropout(dropout)511 conv_2d_out_channels = conv_2d_out_channels or out_channels512 self.conv2 = torch.nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1)513 514 if non_linearity == "swish":515 self.nonlinearity = lambda x: F.silu(x)516 elif non_linearity == "mish":517 self.nonlinearity = nn.Mish()518 elif non_linearity == "silu":519 self.nonlinearity = nn.SiLU()520 elif non_linearity == "gelu":521 self.nonlinearity = nn.GELU()522 523 self.upsample = self.downsample = None524 if self.up:525 if kernel == "fir":526 fir_kernel = (1, 3, 3, 1)527 self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel)528 elif kernel == "sde_vp":529 self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest")530 else:531 self.upsample = Upsample2D(in_channels, use_conv=False)532 elif self.down:533 if kernel == "fir":534 fir_kernel = (1, 3, 3, 1)535 self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel)536 elif kernel == "sde_vp":537 self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2)538 else:539 self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op")540 541 self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut542 543 self.conv_shortcut = None544 if self.use_in_shortcut:545 self.conv_shortcut = torch.nn.Conv2d(546 in_channels, conv_2d_out_channels, kernel_size=1, stride=1, padding=0, bias=conv_shortcut_bias547 )548 549 def forward(self, input_tensor, temb):550 hidden_states = input_tensor551 552 if self.time_embedding_norm == "ada_group":553 hidden_states = self.norm1(hidden_states, temb)554 else:555 hidden_states = self.norm1(hidden_states)556 557 hidden_states = self.nonlinearity(hidden_states)558 559 if self.upsample is not None:560 # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984561 if hidden_states.shape[0] >= 64:562 input_tensor = input_tensor.contiguous()563 hidden_states = hidden_states.contiguous()564 input_tensor = self.upsample(input_tensor)565 hidden_states = self.upsample(hidden_states)566 elif self.downsample is not None:567 input_tensor = self.downsample(input_tensor)568 hidden_states = self.downsample(hidden_states)569 570 hidden_states = self.conv1(hidden_states)571 572 if self.time_emb_proj is not None:573 temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None]574 575 if temb is not None and self.time_embedding_norm == "default":576 hidden_states = hidden_states + temb577 578 if self.time_embedding_norm == "ada_group":579 hidden_states = self.norm2(hidden_states, temb)580 else:581 hidden_states = self.norm2(hidden_states)582 583 if temb is not None and self.time_embedding_norm == "scale_shift":584 scale, shift = torch.chunk(temb, 2, dim=1)585 hidden_states = hidden_states * (1 + scale) + shift586 587 hidden_states = self.nonlinearity(hidden_states)588 589 hidden_states = self.dropout(hidden_states)590 hidden_states = self.conv2(hidden_states)591 592 if self.conv_shortcut is not None:593 input_tensor = self.conv_shortcut(input_tensor)594 595 output_tensor = (input_tensor + hidden_states) / self.output_scale_factor596 597 return output_tensor598 599 600class Mish(torch.nn.Module):601 def forward(self, hidden_states):602 return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))603 604 605# unet_rl.py606def rearrange_dims(tensor):607 if len(tensor.shape) == 2:608 return tensor[:, :, None]609 if len(tensor.shape) == 3:610 return tensor[:, :, None, :]611 elif len(tensor.shape) == 4:612 return tensor[:, :, 0, :]613 else:614 raise ValueError(f"`len(tensor)`: {len(tensor)} has to be 2, 3 or 4.")615 616 617class Conv1dBlock(nn.Module):618 """619 Conv1d --> GroupNorm --> Mish620 """621 622 def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8):623 super().__init__()624 625 self.conv1d = nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2)626 self.group_norm = nn.GroupNorm(n_groups, out_channels)627 self.mish = nn.Mish()628 629 def forward(self, x):630 x = self.conv1d(x)631 x = rearrange_dims(x)632 x = self.group_norm(x)633 x = rearrange_dims(x)634 x = self.mish(x)635 return x636 637 638# unet_rl.py639class ResidualTemporalBlock1D(nn.Module):640 def __init__(self, inp_channels, out_channels, embed_dim, kernel_size=5):641 super().__init__()642 self.conv_in = Conv1dBlock(inp_channels, out_channels, kernel_size)643 self.conv_out = Conv1dBlock(out_channels, out_channels, kernel_size)644 645 self.time_emb_act = nn.Mish()646 self.time_emb = nn.Linear(embed_dim, out_channels)647 648 self.residual_conv = (649 nn.Conv1d(inp_channels, out_channels, 1) if inp_channels != out_channels else nn.Identity()650 )651 652 def forward(self, x, t):653 """654 Args:655 x : [ batch_size x inp_channels x horizon ]656 t : [ batch_size x embed_dim ]657 658 returns:659 out : [ batch_size x out_channels x horizon ]660 """661 t = self.time_emb_act(t)662 t = self.time_emb(t)663 out = self.conv_in(x) + rearrange_dims(t)664 out = self.conv_out(out)665 return out + self.residual_conv(x)666 667 668def upsample_2d(hidden_states, kernel=None, factor=2, gain=1):669 r"""Upsample2D a batch of 2D images with the given filter.670 Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and upsamples each image with the given671 filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the specified672 `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its shape is673 a: multiple of the upsampling factor.674 675 Args:676 hidden_states: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.677 kernel: FIR filter of the shape `[firH, firW]` or `[firN]`678 (separable). The default is `[1] * factor`, which corresponds to nearest-neighbor upsampling.679 factor: Integer upsampling factor (default: 2).680 gain: Scaling factor for signal magnitude (default: 1.0).681 682 Returns:683 output: Tensor of the shape `[N, C, H * factor, W * factor]`684 """685 assert isinstance(factor, int) and factor >= 1686 if kernel is None:687 kernel = [1] * factor688 689 kernel = torch.tensor(kernel, dtype=torch.float32)690 if kernel.ndim == 1:691 kernel = torch.outer(kernel, kernel)692 kernel /= torch.sum(kernel)693 694 kernel = kernel * (gain * (factor**2))695 pad_value = kernel.shape[0] - factor696 output = upfirdn2d_native(697 hidden_states,698 kernel.to(device=hidden_states.device),699 up=factor,700 pad=((pad_value + 1) // 2 + factor - 1, pad_value // 2),701 )702 return output703 704 705def downsample_2d(hidden_states, kernel=None, factor=2, gain=1):706 r"""Downsample2D a batch of 2D images with the given filter.707 Accepts a batch of 2D images of the shape `[N, C, H, W]` or `[N, H, W, C]` and downsamples each image with the708 given filter. The filter is normalized so that if the input pixels are constant, they will be scaled by the709 specified `gain`. Pixels outside the image are assumed to be zero, and the filter is padded with zeros so that its710 shape is a multiple of the downsampling factor.711 712 Args:713 hidden_states: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.714 kernel: FIR filter of the shape `[firH, firW]` or `[firN]`715 (separable). The default is `[1] * factor`, which corresponds to average pooling.716 factor: Integer downsampling factor (default: 2).717 gain: Scaling factor for signal magnitude (default: 1.0).718 719 Returns:720 output: Tensor of the shape `[N, C, H // factor, W // factor]`721 """722 723 assert isinstance(factor, int) and factor >= 1724 if kernel is None:725 kernel = [1] * factor726 727 kernel = torch.tensor(kernel, dtype=torch.float32)728 if kernel.ndim == 1:729 kernel = torch.outer(kernel, kernel)730 kernel /= torch.sum(kernel)731 732 kernel = kernel * gain733 pad_value = kernel.shape[0] - factor734 output = upfirdn2d_native(735 hidden_states, kernel.to(device=hidden_states.device), down=factor, pad=((pad_value + 1) // 2, pad_value // 2)736 )737 return output738 739 740def upfirdn2d_native(tensor, kernel, up=1, down=1, pad=(0, 0)):741 up_x = up_y = up742 down_x = down_y = down743 pad_x0 = pad_y0 = pad[0]744 pad_x1 = pad_y1 = pad[1]745 746 _, channel, in_h, in_w = tensor.shape747 tensor = tensor.reshape(-1, in_h, in_w, 1)748 749 _, in_h, in_w, minor = tensor.shape750 kernel_h, kernel_w = kernel.shape751 752 out = tensor.view(-1, in_h, 1, in_w, 1, minor)753 out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])754 out = out.view(-1, in_h * up_y, in_w * up_x, minor)755 756 out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])757 out = out.to(tensor.device) # Move back to mps if necessary758 out = out[759 :,760 max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),761 max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),762 :,763 ]764 765 out = out.permute(0, 3, 1, 2)766 out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])767 w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)768 out = F.conv2d(out, w)769 out = out.reshape(770 -1,771 minor,772 in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,773 in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,774 )775 out = out.permute(0, 2, 3, 1)776 out = out[:, ::down_y, ::down_x, :]777 778 out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1779 out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1780 781 return out.view(-1, channel, out_h, out_w)782 783 784class TemporalConvLayer(nn.Module):785 """786 Temporal convolutional layer that can be used for video (sequence of images) input Code mostly copied from:787 https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/models/multi_modal/video_synthesis/unet_sd.py#L1016788 """789 790 def __init__(self, in_dim, out_dim=None, dropout=0.0):791 super().__init__()792 out_dim = out_dim or in_dim793 self.in_dim = in_dim794 self.out_dim = out_dim795 796 # conv layers797 self.conv1 = nn.Sequential(798 nn.GroupNorm(32, in_dim), nn.SiLU(), nn.Conv3d(in_dim, out_dim, (3, 1, 1), padding=(1, 0, 0))799 )800 self.conv2 = nn.Sequential(801 nn.GroupNorm(32, out_dim),802 nn.SiLU(),803 nn.Dropout(dropout),804 nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)),805 )806 self.conv3 = nn.Sequential(807 nn.GroupNorm(32, out_dim),808 nn.SiLU(),809 nn.Dropout(dropout),810 nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)),811 )812 self.conv4 = nn.Sequential(813 nn.GroupNorm(32, out_dim),814 nn.SiLU(),815 nn.Dropout(dropout),816 nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)),817 )818 819 # zero out the last layer params,so the conv block is identity820 nn.init.zeros_(self.conv4[-1].weight)821 nn.init.zeros_(self.conv4[-1].bias)822 823 def forward(self, hidden_states, num_frames=1):824 hidden_states = (825 hidden_states[None, :].reshape((-1, num_frames) + hidden_states.shape[1:]).permute(0, 2, 1, 3, 4)826 )827 828 identity = hidden_states829 hidden_states = self.conv1(hidden_states)830 hidden_states = self.conv2(hidden_states)831 hidden_states = self.conv3(hidden_states)832 hidden_states = self.conv4(hidden_states)833 834 hidden_states = identity + hidden_states835 836 hidden_states = hidden_states.permute(0, 2, 1, 3, 4).reshape(837 (hidden_states.shape[0] * hidden_states.shape[2], -1) + hidden_states.shape[3:]838 )839 return hidden_states840