StableDiffusionModel
PretrainedModelStableDiffusionModel(config: StableDiffusionConfig)Latent diffusion: an autoencoder, a conditional U-Net, a sampler.
Parameters
configStableDiffusionConfigAttributes
schedulerDDIMSchedulerforward uses to noise a
latent. Sampling defaults to PNDM instead — see
StableDiffusionForImageGeneration.generate.Notes
Reference: Rombach et al., "High-Resolution Image Synthesis with Latent Diffusion Models", CVPR, 2022 (arXiv:2112.10752).
The text encoder is not held here. The paper defines
as a domain-specific encoder and evaluates
several, and the released models freeze a CLIP text tower rather
than training one — so conditioning arrives as an already-encoded
(B, L, cross_attention_dim) sequence and the caller chooses what
produced it. At the default width that is
lucid.models.clip_vit_large_14's text tower.
Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
... StableDiffusionConfig, StableDiffusionModel)
>>> config = StableDiffusionConfig(sample_size=32, downsample_factor=4,
... vae_block_out_channels=(32, 64, 64),
... unet_block_out_channels=(32, 64),
... attention_head_dim=32,
... cross_attention_dim=16, context_length=4)
>>> model = StableDiffusionModel(config).eval()
>>> out = model(lucid.randn((1, 3, 32, 32)), lucid.randn((1, 4, 16)))
>>> out.noise_pred.shape
(1, 4, 8, 8)Used by 2
Constructors
1Instance methods
3Scaled latent back to an image.
Parameters
latentTensor(B, latent_channels, h, w).Returns
Tensor(B, out_channels, h*f, w*f).
Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
... StableDiffusionConfig, StableDiffusionModel)
>>> config = StableDiffusionConfig(sample_size=32, downsample_factor=4,
... vae_block_out_channels=(32, 64, 64),
... unet_block_out_channels=(32, 64),
... attention_head_dim=32,
... cross_attention_dim=16, context_length=4)
>>> model = StableDiffusionModel(config).eval()
>>> images = lucid.randn((1, 3, 32, 32))
>>> latent = model.encode_image(images, sample=False)
>>> model.decode_latent(latent).shape
(1, 3, 32, 32)
The scale is divided back out before decoding, so this undoes the
scaling encode_image applied: it matches decoding the
unscaled mode directly.
>>> raw = model.vae.encode(images).mode()
>>> bool(lucid.allclose(model.decode_latent(latent),
... model.vae.decode(raw), atol=1e-5))
TrueImage to scaled latent.
Parameters
imagesTensor(B, in_channels, H, W).samplebool= TrueReturns
Tensor(B, latent_channels, H/f, W/f), already scaled for the
diffusion process.
Examples
>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
... StableDiffusionConfig, StableDiffusionModel)
>>> config = StableDiffusionConfig(sample_size=32, downsample_factor=4,
... vae_block_out_channels=(32, 64, 64),
... unet_block_out_channels=(32, 64),
... attention_head_dim=32,
... cross_attention_dim=16, context_length=4)
>>> model = StableDiffusionModel(config).eval()
>>> images = lucid.randn((1, 3, 32, 32))
>>> latent = model.encode_image(images, sample=False)
>>> latent.shape
(1, 4, 8, 8)
What this adds over the autoencoder is the scaling constant: the
latent is the posterior's mode times 0.18215, the released first
stage's inverse standard deviation.
>>> raw = model.vae.encode(images).mode()
>>> bool(lucid.allclose(latent, raw * 0.18215))
Trueforward(images: Tensor, context: Tensor, timestep: Tensor | None = None, return_loss: bool = False)One training step's worth of the denoising objective.
Parameters
Returns
StableDiffusionOutputPrediction, the noised latent, and optionally the loss.