(
self,
*,
ch,
out_ch,
ch_mult=(1, 2, 4, 8),
num_res_blocks,
attn_resolutions,
dropout=0.0,
resamp_with_conv=True,
in_channels,
resolution,
z_channels,
give_pre_end=False,
zq_ch=None,
add_conv=False,
pad_mode="first",
temporal_compress_times=4,
gather_norm=False,
**ignorekwargs,
)
| 848 | # * ContextParallelDecoder3D has 123.37 M params |
| 849 | class ContextParallelDecoder3D(nn.Module): |
| 850 | def __init__( |
| 851 | self, |
| 852 | *, |
| 853 | ch, |
| 854 | out_ch, |
| 855 | ch_mult=(1, 2, 4, 8), |
| 856 | num_res_blocks, |
| 857 | attn_resolutions, |
| 858 | dropout=0.0, |
| 859 | resamp_with_conv=True, |
| 860 | in_channels, |
| 861 | resolution, |
| 862 | z_channels, |
| 863 | give_pre_end=False, |
| 864 | zq_ch=None, |
| 865 | add_conv=False, |
| 866 | pad_mode="first", |
| 867 | temporal_compress_times=4, |
| 868 | gather_norm=False, |
| 869 | **ignorekwargs, |
| 870 | ): |
| 871 | super().__init__() |
| 872 | self.ch = ch |
| 873 | self.temb_ch = 0 |
| 874 | self.num_resolutions = len(ch_mult) |
| 875 | self.num_res_blocks = num_res_blocks |
| 876 | self.resolution = resolution |
| 877 | self.in_channels = in_channels |
| 878 | self.give_pre_end = give_pre_end |
| 879 | |
| 880 | # log2 of temporal_compress_times |
| 881 | self.temporal_compress_level = int(np.log2(temporal_compress_times)) |
| 882 | |
| 883 | if zq_ch is None: |
| 884 | zq_ch = z_channels |
| 885 | |
| 886 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 887 | in_ch_mult = (1,) + tuple(ch_mult) |
| 888 | block_in = ch * ch_mult[self.num_resolutions - 1] |
| 889 | curr_res = resolution // 2 ** (self.num_resolutions - 1) |
| 890 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 891 | print("Working with z of shape {} = {} dimensions.".format(self.z_shape, np.prod(self.z_shape))) |
| 892 | |
| 893 | self.conv_in = ContextParallelCausalConv3d( |
| 894 | chan_in=z_channels, |
| 895 | chan_out=block_in, |
| 896 | kernel_size=3, |
| 897 | ) |
| 898 | |
| 899 | # middle |
| 900 | self.mid = nn.Module() |
| 901 | self.mid.block_1 = ContextParallelResnetBlock3D( |
| 902 | in_channels=block_in, |
| 903 | out_channels=block_in, |
| 904 | temb_channels=self.temb_ch, |
| 905 | dropout=dropout, |
| 906 | zq_ch=zq_ch, |
| 907 | add_conv=add_conv, |
nothing calls this directly
no test coverage detected