| 530 | |
| 531 | |
| 532 | class Decoder(nn.Module): |
| 533 | def __init__(self, *, ch, out_ch, ch_mult=(1, 2, 4, 8), num_res_blocks, |
| 534 | attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, |
| 535 | resolution, z_channels, give_pre_end=False, output_key=[], **ignorekwargs): |
| 536 | super().__init__() |
| 537 | self.output_key = output_key |
| 538 | self.ch = ch |
| 539 | self.temb_ch = 0 |
| 540 | self.num_resolutions = len(ch_mult) |
| 541 | self.num_res_blocks = num_res_blocks |
| 542 | self.resolution = resolution |
| 543 | self.in_channels = in_channels |
| 544 | self.give_pre_end = give_pre_end |
| 545 | |
| 546 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 547 | in_ch_mult = (1,)+tuple(ch_mult) |
| 548 | block_in = ch*ch_mult[self.num_resolutions-1] |
| 549 | curr_res = resolution // 2**(self.num_resolutions-1) |
| 550 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 551 | print("Working with z of shape {} = {} dimensions.".format( |
| 552 | self.z_shape, np.prod(self.z_shape))) |
| 553 | |
| 554 | # z to block_in |
| 555 | self.conv_in = torch.nn.Conv2d(z_channels, |
| 556 | block_in, |
| 557 | kernel_size=3, |
| 558 | stride=1, |
| 559 | padding=1) |
| 560 | |
| 561 | # middle |
| 562 | self.mid = nn.Module() |
| 563 | self.mid.block_1 = ResnetBlock(in_channels=block_in, |
| 564 | out_channels=block_in, |
| 565 | temb_channels=self.temb_ch, |
| 566 | dropout=dropout) |
| 567 | self.mid.attn_1 = AttnBlock(block_in) |
| 568 | self.mid.block_2 = ResnetBlock(in_channels=block_in, |
| 569 | out_channels=block_in, |
| 570 | temb_channels=self.temb_ch, |
| 571 | dropout=dropout) |
| 572 | |
| 573 | # upsampling |
| 574 | self.up = nn.ModuleList() |
| 575 | for i_level in reversed(range(self.num_resolutions)): |
| 576 | block = nn.ModuleList() |
| 577 | attn = nn.ModuleList() |
| 578 | block_out = ch*ch_mult[i_level] |
| 579 | for i_block in range(self.num_res_blocks+1): |
| 580 | block.append(ResnetBlock(in_channels=block_in, |
| 581 | out_channels=block_out, |
| 582 | temb_channels=self.temb_ch, |
| 583 | dropout=dropout)) |
| 584 | block_in = block_out |
| 585 | if curr_res in attn_resolutions: |
| 586 | attn.append(AttnBlock(block_in)) |
| 587 | up = nn.Module() |
| 588 | up.block = block |
| 589 | up.attn = attn |
nothing calls this directly
no outgoing calls
no test coverage detected