(
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,
)
| 746 | |
| 747 | class ContextParallelDecoder3D(nn.Module): |
| 748 | def __init__( |
| 749 | self, |
| 750 | *, |
| 751 | ch, |
| 752 | out_ch, |
| 753 | ch_mult=(1, 2, 4, 8), |
| 754 | num_res_blocks, |
| 755 | attn_resolutions, |
| 756 | dropout=0.0, |
| 757 | resamp_with_conv=True, |
| 758 | in_channels, |
| 759 | resolution, |
| 760 | z_channels, |
| 761 | give_pre_end=False, |
| 762 | zq_ch=None, |
| 763 | add_conv=False, |
| 764 | pad_mode="first", |
| 765 | temporal_compress_times=4, |
| 766 | gather_norm=False, |
| 767 | **ignorekwargs, |
| 768 | ): |
| 769 | super().__init__() |
| 770 | self.ch = ch |
| 771 | self.temb_ch = 0 |
| 772 | self.num_resolutions = len(ch_mult) |
| 773 | self.num_res_blocks = num_res_blocks |
| 774 | self.resolution = resolution |
| 775 | self.in_channels = in_channels |
| 776 | self.give_pre_end = give_pre_end |
| 777 | |
| 778 | # log2 of temporal_compress_times |
| 779 | self.temporal_compress_level = int(np.log2(temporal_compress_times)) |
| 780 | |
| 781 | if zq_ch is None: |
| 782 | zq_ch = z_channels |
| 783 | |
| 784 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 785 | in_ch_mult = (1,) + tuple(ch_mult) |
| 786 | block_in = ch * ch_mult[self.num_resolutions - 1] |
| 787 | curr_res = resolution // 2 ** (self.num_resolutions - 1) |
| 788 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 789 | print("Working with z of shape {} = {} dimensions.".format(self.z_shape, np.prod(self.z_shape))) |
| 790 | |
| 791 | self.conv_in = ContextParallelCausalConv3d( |
| 792 | chan_in=z_channels, |
| 793 | chan_out=block_in, |
| 794 | kernel_size=3, |
| 795 | ) |
| 796 | |
| 797 | # middle |
| 798 | self.mid = nn.Module() |
| 799 | self.mid.block_1 = ContextParallelResnetBlock3D( |
| 800 | in_channels=block_in, |
| 801 | out_channels=block_in, |
| 802 | temb_channels=self.temb_ch, |
| 803 | dropout=dropout, |
| 804 | zq_ch=zq_ch, |
| 805 | add_conv=add_conv, |
nothing calls this directly
no test coverage detected