(self,
ch=128,
z_channels=16,
out_channels=3,
ch_mult=(1, 2, 4, 4),
num_res_blocks=2,
temporal_up_layers=[2, 3],
temporal_downsample=4,
resamp_with_conv=True,
version=1,
)
| 756 | |
| 757 | class VideoDecoder(nn.Module): |
| 758 | def __init__(self, |
| 759 | ch=128, |
| 760 | z_channels=16, |
| 761 | out_channels=3, |
| 762 | ch_mult=(1, 2, 4, 4), |
| 763 | num_res_blocks=2, |
| 764 | temporal_up_layers=[2, 3], |
| 765 | temporal_downsample=4, |
| 766 | resamp_with_conv=True, |
| 767 | version=1, |
| 768 | ): |
| 769 | super().__init__() |
| 770 | |
| 771 | temb_ch = 0 |
| 772 | |
| 773 | self.num_resolutions = len(ch_mult) |
| 774 | self.num_res_blocks = num_res_blocks |
| 775 | self.temporal_downsample = temporal_downsample |
| 776 | |
| 777 | block_in = ch * ch_mult[self.num_resolutions - 1] |
| 778 | self.version = version |
| 779 | if version == 2: |
| 780 | channels = 4 * z_channels * 2 ** 3 |
| 781 | self.conv_in = CausalConv(z_channels, channels, kernel_size=3) |
| 782 | self.shortcut_in = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(z_channels, channels, 1) |
| 783 | self.conv_unpatchify = ConvPixelShuffleUpSampleLayer3D(channels, block_in, kernel_size=3, factor=2) |
| 784 | self.shortcut_unpathify = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(channels, block_in, 2) |
| 785 | else: |
| 786 | self.conv_in = CausalConv(z_channels, block_in, kernel_size=3) |
| 787 | |
| 788 | # middle |
| 789 | self.mid = nn.Module() |
| 790 | self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) |
| 791 | self.mid.attn_1 = AttnBlock(block_in) |
| 792 | self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) |
| 793 | |
| 794 | # upsampling |
| 795 | self.up_id = len(temporal_up_layers) |
| 796 | self.video_frame_num = 1 |
| 797 | self.cur_video_frame_num = self.video_frame_num // 2 ** self.up_id + 1 |
| 798 | self.up = nn.ModuleList() |
| 799 | for i_level in reversed(range(self.num_resolutions)): |
| 800 | block = nn.ModuleList() |
| 801 | attn = nn.ModuleList() |
| 802 | block_out = ch * ch_mult[i_level] |
| 803 | for i_block in range(self.num_res_blocks + 1): |
| 804 | block.append( |
| 805 | Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch)) |
| 806 | block_in = block_out |
| 807 | up = nn.Module() |
| 808 | up.block = block |
| 809 | up.attn = attn |
| 810 | if i_level != 0: |
| 811 | if i_level in temporal_up_layers: |
| 812 | up.upsample = Upsample3D(block_in) |
| 813 | self.cur_video_frame_num = self.cur_video_frame_num * 2 |
| 814 | else: |
| 815 | up.upsample = Upsample2D(block_in, resamp_with_conv) |
nothing calls this directly
no test coverage detected