(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, give_pre_end=False, tanh_out=False,
natten_kernel_size=-1, use_null_attention=False,
attn_type="vanilla", **ignorekwargs)
| 339 | |
| 340 | class Decoder(nn.Module): |
| 341 | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, |
| 342 | attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, |
| 343 | resolution, z_channels, give_pre_end=False, tanh_out=False, |
| 344 | natten_kernel_size=-1, use_null_attention=False, |
| 345 | attn_type="vanilla", **ignorekwargs): |
| 346 | super().__init__() |
| 347 | self.ch = ch |
| 348 | self.temb_ch = 0 |
| 349 | self.num_resolutions = len(ch_mult) |
| 350 | self.num_res_blocks = num_res_blocks |
| 351 | self.resolution = resolution |
| 352 | self.in_channels = in_channels |
| 353 | self.give_pre_end = give_pre_end |
| 354 | self.tanh_out = tanh_out |
| 355 | self.natten_kernel_size = natten_kernel_size |
| 356 | self.use_null_attention = use_null_attention |
| 357 | |
| 358 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 359 | in_ch_mult = (1,)+tuple(ch_mult) |
| 360 | block_in = ch*ch_mult[self.num_resolutions-1] |
| 361 | curr_res = resolution // 2**(self.num_resolutions-1) |
| 362 | self.z_shape = (1,z_channels,curr_res,curr_res) |
| 363 | |
| 364 | # z to block_in |
| 365 | self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) |
| 366 | |
| 367 | # middle |
| 368 | self.mid = nn.Module() |
| 369 | self.mid.block_1 = ResnetBlock(in_channels=block_in, |
| 370 | out_channels=block_in, |
| 371 | temb_channels=self.temb_ch, |
| 372 | dropout=dropout) |
| 373 | self.mid.attn_1 = make_attn(block_in, attn_type=attn_type, natten_kernel_size=natten_kernel_size, use_null_attention=use_null_attention) |
| 374 | self.mid.block_2 = ResnetBlock(in_channels=block_in, |
| 375 | out_channels=block_in, |
| 376 | temb_channels=self.temb_ch, |
| 377 | dropout=dropout) |
| 378 | |
| 379 | # upsampling |
| 380 | self.up = nn.ModuleList() |
| 381 | for i_level in reversed(range(self.num_resolutions)): |
| 382 | block = nn.ModuleList() |
| 383 | attn = nn.ModuleList() |
| 384 | block_out = ch*ch_mult[i_level] |
| 385 | for i_block in range(self.num_res_blocks+1): |
| 386 | block.append(ResnetBlock(in_channels=block_in, |
| 387 | out_channels=block_out, |
| 388 | temb_channels=self.temb_ch, |
| 389 | dropout=dropout)) |
| 390 | block_in = block_out |
| 391 | if curr_res in attn_resolutions: |
| 392 | attn.append(make_attn(block_in, attn_type=attn_type, natten_kernel_size=natten_kernel_size, use_null_attention=use_null_attention)) |
| 393 | up = nn.Module() |
| 394 | up.block = block |
| 395 | up.attn = attn |
| 396 | if i_level != 0: |
| 397 | up.upsample = Upsample(block_in, resamp_with_conv) |
| 398 | curr_res = curr_res * 2 |
nothing calls this directly
no test coverage detected