| 135 | return self.conv_out(self.norm_out(x).swish()) |
| 136 | |
| 137 | class AutoencoderKL: |
| 138 | def __init__(self): |
| 139 | self.encoder = Encoder() |
| 140 | self.decoder = Decoder() |
| 141 | self.quant_conv = Conv2d(8, 8, 1) |
| 142 | self.post_quant_conv = Conv2d(4, 4, 1) |
| 143 | |
| 144 | def __call__(self, x): |
| 145 | latent = self.encoder(x) |
| 146 | latent = self.quant_conv(latent) |
| 147 | latent = latent[:, 0:4] # only the means |
| 148 | latent = self.post_quant_conv(latent) |
| 149 | return self.decoder(latent) |
| 150 | |
| 151 | # not to be confused with ResnetBlock |
| 152 | class ResBlock: |