| 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 |