| 157 | self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) |
| 158 | |
| 159 | def forward(self, x: Tensor) -> Tensor: |
| 160 | # downsampling |
| 161 | hs = [self.conv_in(x)] |
| 162 | for i_level in range(self.num_resolutions): |
| 163 | for i_block in range(self.num_res_blocks): |
| 164 | h = self.down[i_level].block[i_block](hs[-1]) |
| 165 | if len(self.down[i_level].attn) > 0: |
| 166 | h = self.down[i_level].attn[i_block](h) |
| 167 | hs.append(h) |
| 168 | if i_level != self.num_resolutions - 1: |
| 169 | hs.append(self.down[i_level].downsample(hs[-1])) |
| 170 | |
| 171 | # middle |
| 172 | h = hs[-1] |
| 173 | h = self.mid.block_1(h) |
| 174 | h = self.mid.attn_1(h) |
| 175 | h = self.mid.block_2(h) |
| 176 | # end |
| 177 | h = self.norm_out(h) |
| 178 | h = swish(h) |
| 179 | h = self.conv_out(h) |
| 180 | return h |
| 181 | |
| 182 | |
| 183 | class Decoder(nn.Module): |