class

DDPMUNet

extendsModule

DDPMUNet(config: DDPMConfig)

source edit

U-Net denoiser shared by every DDPM variant.

Implements the architecture introduced by Ho, Jain, and Abbeel, 2020: a symmetric U-Net augmented with sinusoidal-timestep conditioning, GroupNorm-equipped residual blocks, and self-attention at low-resolution feature maps. The encoder reduces spatial resolution by a factor of $2^{L-1}$ over $L = \mathrm{len}(\text{channel\_mult})$ stages; the middle block applies ResBlock → Attn → ResBlock at the bottleneck; the decoder mirrors the encoder with concatenated skip connections. The output convolution emits out_channels_effective channels — equal to in_channels (or 2 * in_channels when learn_sigma is enabled for variance prediction).

Spatial layout per stage $l \in [0, L)$ : * num_res_blocks ResBlocks (with optional attention at attention_resolutions). * Downsample → next stage (skipped for l == L - 1).

Decoder mirrors the encoder with concatenated skip features; the extra ResBlock per stage (num_res_blocks + 1 total on the decoder side) consumes the post-downsample skip.

Parameters

configDDPMConfig

Hyperparameters controlling base_channels, channel_mult, num_res_blocks, attention_resolutions, num_heads, dropout, resnet_groups, and learn_sigma.

Attributes

time_mlpnn.TimestepEmbedding

Sinusoidal timestep embedding + 2-layer MLP producing the per- timestep conditioning vector of width 4 * base_channels.

conv_innn.Conv2d

Stem convolution lifting input channels to base_channels.

down_res, down_attn, down_samplenn.ModuleList

Encoder ResBlocks, optional AttentionBlocks, and Downsamples.

mid_block1, mid_attn, mid_block2nn.Module

Bottleneck stack at the lowest spatial resolution.

up_res, up_attn, up_samplenn.ModuleList

Decoder ResBlocks, optional AttentionBlocks, and Upsamples.

norm_outnn.GroupNorm

Output GroupNorm before the final convolution.

conv_outnn.Conv2d

Final 3x3 convolution projecting to out_channels_effective.

Notes

Reference: Ho, Jain, and Abbeel, "Denoising Diffusion Probabilistic Models", NeurIPS, 2020 (arXiv:2006.11239); architecture details in Appendix B. Improved-DDPM additions follow Nichol and Dhariwal, * "Improved Denoising Diffusion Probabilistic Models"*, ICML, 2021 (arXiv:2102.09672).

ResBlock conditioning recipe:

h = \mathrm{Conv}_2\!\big( \mathrm{SiLU}( \mathrm{Norm}_2( \mathrm{SiLU}(t \cdot W_t) + \mathrm{Conv}_1( \mathrm{SiLU}(\mathrm{Norm}_1(x)) ) ) ) \big) + \mathrm{Skip}(x).

Examples

>>> import lucid
>>> from lucid.models.generative.ddpm import DDPMConfig
>>> from lucid.models.generative.ddpm._model import DDPMUNet
>>> cfg = DDPMConfig(sample_size=32, base_channels=32,
...                  channel_mult=(1, 2), num_res_blocks=1,
...                  attention_resolutions=(16,), resnet_groups=16)
>>> unet = DDPMUNet(cfg).eval()
>>> x = lucid.randn((1, 3, 32, 32))
>>> t = lucid.tensor([42]).long()
>>> out = unet(x, t)
>>> out.shape   # (1, 3, 32, 32)
(1, 3, 32, 32)

Used by 2

Constructors

dunder

init

→None

__init__(config: DDPMConfig)

source edit

Instance methods

forward

→Tensor

forward(sample: Tensor, timestep: Tensor)

source edit

Predict noise (or [noise, σ²]) at the given timestep.

Parameters

sampleTensor

(B, in_channels, H, W) noisy image x_t.

timestepTensor

(B,) or () int timestep tensor. Scalar tensors are broadcast.

Returns

Tensor

(B, out_channels_effective, H, W) raw network output.

>>> import lucid >>> from lucid.models.generative.ddpm import DDPMConfig >>> from lucid.models.generative.ddpm._model import DDPMUNet >>> cfg = DDPMConfig(sample_size=32, base_channels=32, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=(16,), resnet_groups=16) >>> unet = DDPMUNet(cfg).eval() >>> x = lucid.randn((1, 3, 32, 32)) >>> t = lucid.tensor([42]).long() >>> out = unet(x, t) >>> out.shape # (1, 3, 32, 32) (1, 3, 32, 32)