| 62 | |
| 63 | |
| 64 | class Encoder(nn.Module): |
| 65 | def __init__(self, in_channels=3, ch=128, ch_mult=(1,1,2,2,4), num_res_blocks=2, |
| 66 | norm_type='group', dropout=0.0, resamp_with_conv=True, z_channels=256): |
| 67 | super().__init__() |
| 68 | self.num_resolutions = len(ch_mult) |
| 69 | self.num_res_blocks = num_res_blocks |
| 70 | self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1) |
| 71 | |
| 72 | # downsampling |
| 73 | in_ch_mult = (1,) + tuple(ch_mult) |
| 74 | self.conv_blocks = nn.ModuleList() |
| 75 | for i_level in range(self.num_resolutions): |
| 76 | conv_block = nn.Module() |
| 77 | # res & attn |
| 78 | res_block = nn.ModuleList() |
| 79 | attn_block = nn.ModuleList() |
| 80 | block_in = ch*in_ch_mult[i_level] |
| 81 | block_out = ch*ch_mult[i_level] |
| 82 | for _ in range(self.num_res_blocks): |
| 83 | res_block.append(ResnetBlock(block_in, block_out, dropout=dropout, norm_type=norm_type)) |
| 84 | block_in = block_out |
| 85 | if i_level == self.num_resolutions - 1: |
| 86 | attn_block.append(AttnBlock(block_in, norm_type)) |
| 87 | conv_block.res = res_block |
| 88 | conv_block.attn = attn_block |
| 89 | # downsample |
| 90 | if i_level != self.num_resolutions-1: |
| 91 | conv_block.downsample = Downsample(block_in, resamp_with_conv) |
| 92 | self.conv_blocks.append(conv_block) |
| 93 | |
| 94 | # middle |
| 95 | self.mid = nn.ModuleList() |
| 96 | self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type)) |
| 97 | self.mid.append(AttnBlock(block_in, norm_type=norm_type)) |
| 98 | self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type)) |
| 99 | |
| 100 | # end |
| 101 | self.norm_out = Normalize(block_in, norm_type) |
| 102 | self.conv_out = nn.Conv2d(block_in, z_channels, kernel_size=3, stride=1, padding=1) |
| 103 | |
| 104 | |
| 105 | def forward(self, x): |
| 106 | h = self.conv_in(x) |
| 107 | # downsampling |
| 108 | for i_level, block in enumerate(self.conv_blocks): |
| 109 | for i_block in range(self.num_res_blocks): |
| 110 | h = block.res[i_block](h) |
| 111 | if len(block.attn) > 0: |
| 112 | h = block.attn[i_block](h) |
| 113 | if i_level != self.num_resolutions - 1: |
| 114 | h = block.downsample(h) |
| 115 | |
| 116 | # middle |
| 117 | for mid_block in self.mid: |
| 118 | h = mid_block(h) |
| 119 | |
| 120 | # end |
| 121 | h = self.norm_out(h) |