(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, use_linear_attn=False,
attn_type="vanilla", num_fuse_block=2, fusion_w=1.0, **ignorekwargs)
| 676 | |
| 677 | class Decoder_Mix(nn.Module): |
| 678 | def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, |
| 679 | attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, |
| 680 | resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False, |
| 681 | attn_type="vanilla", num_fuse_block=2, fusion_w=1.0, **ignorekwargs): |
| 682 | super().__init__() |
| 683 | if use_linear_attn: attn_type = "linear" |
| 684 | self.ch = ch |
| 685 | self.temb_ch = 0 |
| 686 | self.num_resolutions = len(ch_mult) |
| 687 | self.num_res_blocks = num_res_blocks |
| 688 | self.resolution = resolution |
| 689 | self.in_channels = in_channels |
| 690 | self.give_pre_end = give_pre_end |
| 691 | self.tanh_out = tanh_out |
| 692 | self.fusion_w = fusion_w |
| 693 | |
| 694 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 695 | in_ch_mult = (1,)+tuple(ch_mult) |
| 696 | block_in = ch*ch_mult[self.num_resolutions-1] |
| 697 | curr_res = resolution // 2**(self.num_resolutions-1) |
| 698 | self.z_shape = (1,z_channels,curr_res,curr_res) |
| 699 | print("Working with z of shape {} = {} dimensions.".format( |
| 700 | self.z_shape, np.prod(self.z_shape))) |
| 701 | |
| 702 | # z to block_in |
| 703 | self.conv_in = torch.nn.Conv2d(z_channels, |
| 704 | block_in, |
| 705 | kernel_size=3, |
| 706 | stride=1, |
| 707 | padding=1) |
| 708 | |
| 709 | # middle |
| 710 | self.mid = nn.Module() |
| 711 | self.mid.block_1 = ResnetBlock(in_channels=block_in, |
| 712 | out_channels=block_in, |
| 713 | temb_channels=self.temb_ch, |
| 714 | dropout=dropout) |
| 715 | self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) |
| 716 | self.mid.block_2 = ResnetBlock(in_channels=block_in, |
| 717 | out_channels=block_in, |
| 718 | temb_channels=self.temb_ch, |
| 719 | dropout=dropout) |
| 720 | |
| 721 | # upsampling |
| 722 | self.up = nn.ModuleList() |
| 723 | for i_level in reversed(range(self.num_resolutions)): |
| 724 | block = nn.ModuleList() |
| 725 | attn = nn.ModuleList() |
| 726 | block_out = ch*ch_mult[i_level] |
| 727 | |
| 728 | if i_level != self.num_resolutions-1: |
| 729 | if i_level != 0: |
| 730 | fuse_layer = Fuse_sft_block_RRDB(in_ch=block_out, out_ch=block_out, num_block=num_fuse_block) |
| 731 | setattr(self, 'fusion_layer_{}'.format(i_level), fuse_layer) |
| 732 | |
| 733 | for i_block in range(self.num_res_blocks+1): |
| 734 | block.append(ResnetBlock(in_channels=block_in, |
| 735 | out_channels=block_out, |
nothing calls this directly
no test coverage detected