| 15 | |
| 16 | |
| 17 | class TemporalDecoder(nn.Module): |
| 18 | def __init__( |
| 19 | self, |
| 20 | in_channels: int = 4, |
| 21 | out_channels: int = 3, |
| 22 | block_out_channels: Tuple[int] = (128, 256, 512, 512), |
| 23 | layers_per_block: int = 2, |
| 24 | ): |
| 25 | super().__init__() |
| 26 | self.layers_per_block = layers_per_block |
| 27 | |
| 28 | self.conv_in = nn.Conv2d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1) |
| 29 | self.mid_block = MidBlockTemporalDecoder( |
| 30 | num_layers=self.layers_per_block, |
| 31 | in_channels=block_out_channels[-1], |
| 32 | out_channels=block_out_channels[-1], |
| 33 | attention_head_dim=block_out_channels[-1], |
| 34 | ) |
| 35 | |
| 36 | # up |
| 37 | self.up_blocks = nn.ModuleList([]) |
| 38 | reversed_block_out_channels = list(reversed(block_out_channels)) |
| 39 | output_channel = reversed_block_out_channels[0] |
| 40 | for i in range(len(block_out_channels)): |
| 41 | prev_output_channel = output_channel |
| 42 | output_channel = reversed_block_out_channels[i] |
| 43 | |
| 44 | is_final_block = i == len(block_out_channels) - 1 |
| 45 | up_block = UpBlockTemporalDecoder( |
| 46 | num_layers=self.layers_per_block + 1, |
| 47 | in_channels=prev_output_channel, |
| 48 | out_channels=output_channel, |
| 49 | add_upsample=not is_final_block, |
| 50 | ) |
| 51 | self.up_blocks.append(up_block) |
| 52 | prev_output_channel = output_channel |
| 53 | |
| 54 | self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-6) |
| 55 | |
| 56 | self.conv_act = nn.SiLU() |
| 57 | self.conv_out = torch.nn.Conv2d( |
| 58 | in_channels=block_out_channels[0], |
| 59 | out_channels=out_channels, |
| 60 | kernel_size=3, |
| 61 | padding=1, |
| 62 | ) |
| 63 | |
| 64 | conv_out_kernel_size = (3, 1, 1) |
| 65 | padding = [int(k // 2) for k in conv_out_kernel_size] |
| 66 | self.time_conv_out = torch.nn.Conv3d( |
| 67 | in_channels=out_channels, |
| 68 | out_channels=out_channels, |
| 69 | kernel_size=conv_out_kernel_size, |
| 70 | padding=padding, |
| 71 | ) |
| 72 | |
| 73 | self.gradient_checkpointing = False |
| 74 | |