The autoencoder that allows us to perform diffusion in the latent space.
| 224 | |
| 225 | |
| 226 | class Autoencoder(nn.Module): |
| 227 | """The autoencoder that allows us to perform diffusion in the latent space.""" |
| 228 | |
| 229 | def __init__(self, config: AutoencoderConfig): |
| 230 | super().__init__() |
| 231 | |
| 232 | self.latent_channels = config.latent_channels_in |
| 233 | self.scaling_factor = config.scaling_factor |
| 234 | self.encoder = Encoder( |
| 235 | config.in_channels, |
| 236 | config.latent_channels_out, |
| 237 | config.block_out_channels, |
| 238 | config.layers_per_block, |
| 239 | resnet_groups=config.norm_num_groups, |
| 240 | ) |
| 241 | self.decoder = Decoder( |
| 242 | config.latent_channels_in, |
| 243 | config.out_channels, |
| 244 | config.block_out_channels, |
| 245 | config.layers_per_block + 1, |
| 246 | resnet_groups=config.norm_num_groups, |
| 247 | ) |
| 248 | |
| 249 | self.quant_proj = nn.Linear( |
| 250 | config.latent_channels_out, config.latent_channels_out |
| 251 | ) |
| 252 | self.post_quant_proj = nn.Linear( |
| 253 | config.latent_channels_in, config.latent_channels_in |
| 254 | ) |
| 255 | |
| 256 | def decode(self, z): |
| 257 | z = z / self.scaling_factor |
| 258 | return self.decoder(self.post_quant_proj(z)) |
| 259 | |
| 260 | def encode(self, x): |
| 261 | x = self.encoder(x) |
| 262 | x = self.quant_proj(x) |
| 263 | mean, logvar = x.split(2, axis=-1) |
| 264 | mean = mean * self.scaling_factor |
| 265 | logvar = logvar + 2 * math.log(self.scaling_factor) |
| 266 | |
| 267 | return mean, logvar |
| 268 | |
| 269 | def __call__(self, x, key=None): |
| 270 | mean, logvar = self.encode(x) |
| 271 | z = mx.random.normal(mean.shape, key=key) * mx.exp(0.5 * logvar) + mean |
| 272 | x_hat = self.decode(z) |
| 273 | |
| 274 | return dict(x_hat=x_hat, z=z, mean=mean, logvar=logvar) |