class

UNet2DConditionModel

extendsModule

UNet2DConditionModel(config: StableDiffusionConfig)

source edit

The denoiser $\epsilon_\theta(z_t, t, \tau_\theta(y))$ .

Parameters

configStableDiffusionConfig

Read for the U-Net fields.

Notes

Reference: Rombach et al., CVPR 2022 (arXiv:2112.10752), §3.3.

Cross-attention runs at every resolution except the deepest down-sampling stage, which is the released arrangement (CrossAttnDownBlock2D three times, then DownBlock2D).

attention_head_dim is read as a head count, not a dimension, despite its name. The released configuration says 8 and the reference builds eight heads of forty channels at the 320-wide stage — not forty heads of eight. Reading it the other way costs nothing in parameters, so the counts still match tensor for tensor; only the activations disagree, which is how it was found.

Examples

>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
...     StableDiffusionConfig, UNet2DConditionModel)
>>> 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)
>>> unet = UNet2DConditionModel(config).eval()
>>> noise = unet(lucid.randn((1, 4, 8, 8)), lucid.tensor([10.0]),
...              lucid.randn((1, 4, 16)))
>>> noise.shape
(1, 4, 8, 8)

Used by 2

Constructors

dunder

init

→None

__init__(config: StableDiffusionConfig)

source edit

Initialise the U-Net. See the class docstring for parameters.

Instance methods

forward

→Tensor

forward(latent: Tensor, timestep: Tensor, context: Tensor)

source edit

Predict the noise in latent at timestep.

Parameters

latentTensor

(B, latent_channels, h, w).

timestepTensor

(B,) or a scalar; broadcast to the batch.

contextTensor

(B, L, cross_attention_dim) —

\tau_\theta(y)

Returns

Tensor

(B, latent_channels, h, w).

>>> import lucid >>> from lucid.models.generative.stable_diffusion import ( ... StableDiffusionConfig, UNet2DConditionModel) >>> 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) >>> unet = UNet2DConditionModel(config).eval() >>> noise = unet(lucid.randn((1, 4, 8, 8)), lucid.tensor([10.0]), ... lucid.randn((1, 4, 16))) >>> noise.shape (1, 4, 8, 8)