The KL-regularised first stage of a latent diffusion model.
Parameters
configStableDiffusionConfigRead for the autoencoder fields.
Notes
Reference: Rombach et al., CVPR 2022 (arXiv:2112.10752), §3.1.
The encoder ends in a 2c-channel convolution whose halves are
the posterior's mean and log-variance, and a 1 x 1
quant_conv after it; the decoder begins with the mirroring
post_quant_conv. Those two look redundant and are not — they
are where the released VQ- and KL-regularised first stages differ,
so keeping them makes the two interchangeable behind one interface.
Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
... AutoencoderKL, StableDiffusionConfig)
>>> config = StableDiffusionConfig(sample_size=32, downsample_factor=4,
... vae_block_out_channels=(32, 64, 64),
... unet_block_out_channels=(32, 64),
... norm_num_groups=32)
>>> vae = AutoencoderKL(config).eval()
>>> out = vae(lucid.randn((1, 3, 32, 32)))
>>> out.latent.shape, out.reconstruction.shape
((1, 4, 8, 8), (1, 3, 32, 32))Used by 2
Constructors
1Instance methods
3Map a latent back to an image.
Parameters
(B, latent_channels, h, w).Returns
Tensor(B, out_channels, h*f, w*f).
Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
... AutoencoderKL, StableDiffusionConfig)
>>> config = StableDiffusionConfig(sample_size=32, downsample_factor=4,
... vae_block_out_channels=(32, 64, 64),
... unet_block_out_channels=(32, 64),
... norm_num_groups=32)
>>> vae = AutoencoderKL(config).eval()
>>> images = lucid.randn((1, 3, 32, 32))
>>> vae.decode(vae.encode(images).mode()).shape
(1, 3, 32, 32)
Nothing in the decoder is tied to a size — it is convolutions and
one attention over positions — so any latent decodes, each side
scaled by f (4 here) on its own.
>>> vae.decode(lucid.randn((1, 4, 4, 6))).shape
(1, 3, 16, 24)encode(x: Tensor)Map an image to its posterior over latents.
Parameters
(B, in_channels, H, W).Returns
DiagonalGaussianMean and log-variance, each (B, latent_channels, H/f, W/f).
Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
... AutoencoderKL, StableDiffusionConfig)
>>> config = StableDiffusionConfig(sample_size=32, downsample_factor=4,
... vae_block_out_channels=(32, 64, 64),
... unet_block_out_channels=(32, 64),
... norm_num_groups=32)
>>> vae = AutoencoderKL(config).eval()
>>> posterior = vae.encode(lucid.randn((2, 3, 32, 32)))
>>> posterior.mean.shape, posterior.logvar.shape
((2, 4, 8, 8), (2, 4, 8, 8))
A distribution rather than a tensor, so the caller chooses: the mode
is the mean and repeats, while a sample draws fresh noise.
>>> bool((posterior.mode() == posterior.mean).all())
True
>>> bool((posterior.sample() == posterior.mode()).all())
Falseforward(x: Tensor, sample: bool = True)Encode, take a latent, and decode it back.
Parameters
(B, in_channels, H, W).samplebool= TrueDraw from the posterior, or take its mode. Training draws;
a deterministic reconstruction does not.
Returns
AutoencoderKLOutputReconstruction, latent and KL.