| 718 | |
| 719 | |
| 720 | class Encoder(nn.Module): |
| 721 | def __init__( |
| 722 | self, |
| 723 | *, |
| 724 | ch, |
| 725 | out_ch, |
| 726 | ch_mult=(1, 2, 4, 8), |
| 727 | num_res_blocks, |
| 728 | attn_resolutions, |
| 729 | dropout=0.0, |
| 730 | resamp_with_conv=True, |
| 731 | in_channels, |
| 732 | z_channels, |
| 733 | resolution=256, |
| 734 | double_z=True, |
| 735 | use_linear_attn=False, |
| 736 | attn_type="vanilla", |
| 737 | use_3d_conv=True, |
| 738 | half_3d=True, |
| 739 | causal=True, |
| 740 | half_t_mult=True, |
| 741 | gradient_checkpointing=True, |
| 742 | **ignore_kwargs, |
| 743 | ): |
| 744 | super().__init__() |
| 745 | if use_linear_attn: |
| 746 | attn_type = "linear" |
| 747 | self.ch = ch |
| 748 | self.temb_ch = 0 |
| 749 | self.num_resolutions = len(ch_mult) |
| 750 | self.num_res_blocks = num_res_blocks |
| 751 | self.resolution = resolution |
| 752 | self.in_channels = in_channels |
| 753 | self.gradient_checkpointing = gradient_checkpointing |
| 754 | |
| 755 | conv_cls = CausalConv3d if causal else nn.Conv3d |
| 756 | conv_cls = conv_cls if use_3d_conv else Conv2dWithExtraDim |
| 757 | |
| 758 | # downsampling |
| 759 | self.conv_in = conv_cls( |
| 760 | in_channels, self.ch, kernel_size=3, stride=1, padding=1 |
| 761 | ) |
| 762 | |
| 763 | curr_res = resolution |
| 764 | in_ch_mult = (1,) + tuple(ch_mult) |
| 765 | self.in_ch_mult = in_ch_mult |
| 766 | self.down = nn.ModuleList() |
| 767 | for i_level in range(self.num_resolutions): |
| 768 | block = nn.ModuleList() |
| 769 | attn = nn.ModuleList() |
| 770 | block_in = ch * in_ch_mult[i_level] |
| 771 | block_out = ch * ch_mult[i_level] |
| 772 | for i_block in range(self.num_res_blocks): |
| 773 | block.append( |
| 774 | ResnetBlock3D( |
| 775 | in_channels=block_in, |
| 776 | out_channels=block_out, |
| 777 | temb_channels=self.temb_ch, |