(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, attn_type="vanilla",
natten_kernel_size=-1, use_null_attention=False,
**ignore_kwargs)
| 249 | |
| 250 | class Encoder(nn.Module): |
| 251 | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, |
| 252 | attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, |
| 253 | resolution, z_channels, double_z=True, attn_type="vanilla", |
| 254 | natten_kernel_size=-1, use_null_attention=False, |
| 255 | **ignore_kwargs): |
| 256 | super().__init__() |
| 257 | self.ch = ch |
| 258 | self.temb_ch = 0 |
| 259 | self.num_resolutions = len(ch_mult) |
| 260 | self.num_res_blocks = num_res_blocks |
| 261 | self.resolution = resolution |
| 262 | self.in_channels = in_channels |
| 263 | |
| 264 | # downsampling |
| 265 | self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1) |
| 266 | |
| 267 | curr_res = resolution |
| 268 | in_ch_mult = (1,)+tuple(ch_mult) |
| 269 | self.in_ch_mult = in_ch_mult |
| 270 | self.down = nn.ModuleList() |
| 271 | for i_level in range(self.num_resolutions): |
| 272 | block = nn.ModuleList() |
| 273 | attn = nn.ModuleList() |
| 274 | block_in = ch*in_ch_mult[i_level] |
| 275 | block_out = ch*ch_mult[i_level] |
| 276 | for i_block in range(self.num_res_blocks): |
| 277 | block.append(ResnetBlock(in_channels=block_in, |
| 278 | out_channels=block_out, |
| 279 | temb_channels=self.temb_ch, |
| 280 | dropout=dropout)) |
| 281 | block_in = block_out |
| 282 | if curr_res in attn_resolutions: |
| 283 | attn.append(make_attn(block_in, attn_type=attn_type, natten_kernel_size=natten_kernel_size, use_null_attention=use_null_attention)) |
| 284 | down = nn.Module() |
| 285 | down.block = block |
| 286 | down.attn = attn |
| 287 | if i_level != self.num_resolutions-1: |
| 288 | down.downsample = Downsample(block_in, resamp_with_conv) |
| 289 | curr_res = curr_res // 2 |
| 290 | self.down.append(down) |
| 291 | |
| 292 | # middle |
| 293 | self.mid = nn.Module() |
| 294 | self.mid.block_1 = ResnetBlock(in_channels=block_in, |
| 295 | out_channels=block_in, |
| 296 | temb_channels=self.temb_ch, |
| 297 | dropout=dropout) |
| 298 | self.mid.attn_1 = make_attn(block_in, attn_type=attn_type, natten_kernel_size=natten_kernel_size, use_null_attention=use_null_attention) |
| 299 | self.mid.block_2 = ResnetBlock(in_channels=block_in, |
| 300 | out_channels=block_in, |
| 301 | temb_channels=self.temb_ch, |
| 302 | dropout=dropout) |
| 303 | |
| 304 | # end |
| 305 | self.norm_out = Normalize(block_in) |
| 306 | self.conv_out = nn.Conv2d(block_in, |
| 307 | 2*z_channels if double_z else z_channels, |
| 308 | kernel_size=3, |
no test coverage detected