| 297 | |
| 298 | class Decoder(nn.Module): |
| 299 | def __init__( |
| 300 | self, |
| 301 | *, |
| 302 | ch, |
| 303 | out_ch, |
| 304 | ch_mult=(1, 2, 4, 8), |
| 305 | num_res_blocks, |
| 306 | attn_resolutions, |
| 307 | dropout=0.0, |
| 308 | resamp_with_conv=True, |
| 309 | in_channels, |
| 310 | resolution, |
| 311 | z_channels, |
| 312 | give_pre_end=False, |
| 313 | **ignorekwargs, |
| 314 | ): |
| 315 | super().__init__() |
| 316 | self.ch = ch |
| 317 | self.temb_ch = 0 |
| 318 | self.num_resolutions = len(ch_mult) |
| 319 | self.num_res_blocks = num_res_blocks |
| 320 | self.resolution = resolution |
| 321 | self.in_channels = in_channels |
| 322 | self.give_pre_end = give_pre_end |
| 323 | |
| 324 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 325 | in_ch_mult = (1,) + tuple(ch_mult) |
| 326 | block_in = ch * ch_mult[self.num_resolutions - 1] |
| 327 | curr_res = resolution // 2 ** (self.num_resolutions - 1) |
| 328 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 329 | print("Working with z of shape {} = {} dimensions.".format(self.z_shape, np.prod(self.z_shape))) |
| 330 | |
| 331 | # z to block_in |
| 332 | self.conv_in = torch.nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) |
| 333 | |
| 334 | # middle |
| 335 | self.mid = nn.Module() |
| 336 | self.mid.block_1 = ResnetBlock( |
| 337 | in_channels=block_in, out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout |
| 338 | ) |
| 339 | self.mid.attn_1 = AttnBlock(block_in) |
| 340 | self.mid.block_2 = ResnetBlock( |
| 341 | in_channels=block_in, out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout |
| 342 | ) |
| 343 | |
| 344 | # upsampling |
| 345 | self.up = nn.ModuleList() |
| 346 | for i_level in reversed(range(self.num_resolutions)): |
| 347 | block = nn.ModuleList() |
| 348 | attn = nn.ModuleList() |
| 349 | block_out = ch * ch_mult[i_level] |
| 350 | for i_block in range(self.num_res_blocks + 1): |
| 351 | block.append( |
| 352 | ResnetBlock( |
| 353 | in_channels=block_in, out_channels=block_out, temb_channels=self.temb_ch, dropout=dropout |
| 354 | ) |
| 355 | ) |
| 356 | block_in = block_out |