| 142 | |
| 143 | |
| 144 | class Encoder(nn.Module): |
| 145 | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, |
| 146 | attn_resolutions, dropout=0.0, resamp_with_conv=False, in_channels, |
| 147 | resolution, z_channels, double_z=True, **ignore_kwargs): |
| 148 | super().__init__() |
| 149 | self.ch = ch |
| 150 | self.temb_ch = 0 |
| 151 | self.num_resolutions = len(ch_mult) |
| 152 | self.num_res_blocks = num_res_blocks |
| 153 | self.resolution = resolution |
| 154 | self.in_channels = in_channels |
| 155 | |
| 156 | # downsampling |
| 157 | self.conv_in = torch.nn.Conv2d(in_channels, |
| 158 | self.ch, |
| 159 | kernel_size=3, |
| 160 | stride=1, |
| 161 | padding=1, |
| 162 | bias=False) |
| 163 | |
| 164 | curr_res = resolution |
| 165 | in_ch_mult = (1,)+tuple(ch_mult) |
| 166 | self.down = nn.ModuleList() |
| 167 | for i_level in range(self.num_resolutions): |
| 168 | block = nn.ModuleList() |
| 169 | block_in = ch*in_ch_mult[i_level] |
| 170 | block_out = ch*ch_mult[i_level] |
| 171 | for i_block in range(self.num_res_blocks): |
| 172 | block.append(ResnetBlock(in_channels=block_in, |
| 173 | out_channels=block_out, |
| 174 | temb_channels=self.temb_ch, |
| 175 | dropout=dropout)) |
| 176 | block_in = block_out |
| 177 | down = nn.Module() |
| 178 | down.block = block |
| 179 | if i_level != self.num_resolutions-1: |
| 180 | down.downsample = Downsample(block_in, resamp_with_conv) |
| 181 | curr_res = curr_res // 2 |
| 182 | self.down.append(down) |
| 183 | |
| 184 | # middle |
| 185 | self.mid = nn.Module() |
| 186 | self.mid.block_1 = ResnetBlock(in_channels=block_in, |
| 187 | out_channels=block_in, |
| 188 | temb_channels=self.temb_ch, |
| 189 | dropout=dropout) |
| 190 | self.mid.block_2 = ResnetBlock(in_channels=block_in, |
| 191 | out_channels=block_in, |
| 192 | temb_channels=self.temb_ch, |
| 193 | dropout=dropout) |
| 194 | |
| 195 | # end |
| 196 | self.norm_out = Normalize(block_in) |
| 197 | self.conv_out = torch.nn.Conv2d(block_in, |
| 198 | 2*z_channels if double_z else z_channels, |
| 199 | kernel_size=1, |
| 200 | stride=1, |
| 201 | padding=0) |