| 865 | |
| 866 | |
| 867 | class Decoder(nn.Module): |
| 868 | def __init__( |
| 869 | self, |
| 870 | *, |
| 871 | ch, |
| 872 | out_ch, |
| 873 | ch_mult=(1, 2, 4, 8), |
| 874 | num_res_blocks, |
| 875 | attn_resolutions, |
| 876 | dropout=0.0, |
| 877 | resamp_with_conv=True, |
| 878 | in_channels, |
| 879 | z_channels, |
| 880 | resolution=256, |
| 881 | give_pre_end=False, |
| 882 | tanh_out=False, |
| 883 | use_linear_attn=False, |
| 884 | attn_type="vanilla", |
| 885 | use_3d_conv=True, |
| 886 | half_3d=True, |
| 887 | causal=True, |
| 888 | half_t_mult=True, |
| 889 | gradient_checkpointing=True, |
| 890 | **ignorekwargs, |
| 891 | ): |
| 892 | super().__init__() |
| 893 | if use_linear_attn: |
| 894 | attn_type = "linear" |
| 895 | self.ch = ch |
| 896 | self.temb_ch = 0 |
| 897 | self.num_resolutions = len(ch_mult) |
| 898 | self.num_res_blocks = num_res_blocks |
| 899 | self.resolution = resolution |
| 900 | self.in_channels = in_channels |
| 901 | self.give_pre_end = give_pre_end |
| 902 | self.tanh_out = tanh_out |
| 903 | self.gradient_checkpointing = gradient_checkpointing |
| 904 | |
| 905 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 906 | in_ch_mult = (1,) + tuple(ch_mult) |
| 907 | block_in = ch * ch_mult[self.num_resolutions - 1] |
| 908 | curr_res = resolution // 2 ** (self.num_resolutions - 1) |
| 909 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 910 | logpy.info( |
| 911 | "Working with z of shape {} = {} dimensions.".format( |
| 912 | self.z_shape, np.prod(self.z_shape) |
| 913 | ) |
| 914 | ) |
| 915 | self.causal = causal |
| 916 | self.use_3d_conv = use_3d_conv |
| 917 | |
| 918 | make_attn_cls = self._make_attn() |
| 919 | make_resblock_cls = self._make_resblock() |
| 920 | make_conv_cls = self._make_conv() |
| 921 | # z to block_in |
| 922 | self.conv_in = make_conv_cls( |
| 923 | z_channels, block_in, kernel_size=3, stride=1, padding=1 |
| 924 | ) |