(
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,
double_z=True,
pad_mode="first",
temporal_compress_times=4,
gather_norm=False,
**ignore_kwargs,
)
| 723 | # * ContextParallelEncoder3D has 92.22 M params. |
| 724 | class ContextParallelEncoder3D(nn.Module): |
| 725 | def __init__( |
| 726 | self, |
| 727 | *, |
| 728 | ch, |
| 729 | out_ch, |
| 730 | ch_mult=(1, 2, 4, 8), |
| 731 | num_res_blocks, |
| 732 | attn_resolutions, |
| 733 | dropout=0.0, |
| 734 | resamp_with_conv=True, |
| 735 | in_channels, |
| 736 | resolution, |
| 737 | z_channels, |
| 738 | double_z=True, |
| 739 | pad_mode="first", |
| 740 | temporal_compress_times=4, |
| 741 | gather_norm=False, |
| 742 | **ignore_kwargs, |
| 743 | ): |
| 744 | super().__init__() |
| 745 | self.ch = ch |
| 746 | self.temb_ch = 0 |
| 747 | self.num_resolutions = len(ch_mult) |
| 748 | self.num_res_blocks = num_res_blocks |
| 749 | self.resolution = resolution |
| 750 | self.in_channels = in_channels |
| 751 | |
| 752 | # log2 of temporal_compress_times |
| 753 | self.temporal_compress_level = int(np.log2(temporal_compress_times)) |
| 754 | |
| 755 | self.conv_in = ContextParallelCausalConv3d( |
| 756 | chan_in=in_channels, |
| 757 | chan_out=self.ch, |
| 758 | kernel_size=3, |
| 759 | ) |
| 760 | |
| 761 | curr_res = resolution |
| 762 | in_ch_mult = (1,) + tuple(ch_mult) |
| 763 | self.down = nn.ModuleList() |
| 764 | for i_level in range(self.num_resolutions): |
| 765 | block = nn.ModuleList() |
| 766 | attn = nn.ModuleList() |
| 767 | block_in = ch * in_ch_mult[i_level] |
| 768 | block_out = ch * ch_mult[i_level] |
| 769 | for i_block in range(self.num_res_blocks): |
| 770 | block.append( |
| 771 | ContextParallelResnetBlock3D( |
| 772 | in_channels=block_in, |
| 773 | out_channels=block_out, |
| 774 | dropout=dropout, |
| 775 | temb_channels=self.temb_ch, |
| 776 | gather_norm=gather_norm, |
| 777 | ) |
| 778 | ) |
| 779 | block_in = block_out |
| 780 | down = nn.Module() |
| 781 | down.block = block |
| 782 | down.attn = attn |
nothing calls this directly
no test coverage detected