class
DiTForImageGeneration
extends
ImageGenerationModelDiTForImageGeneration(config: DiTConfig)DiT posed as a diffusion model: the objective and a DDPM sampler.
Parameters
configDiTConfigThe variant to build.
Attributes
ditDiTModelThe denoising network.
Notes
Reference: Peebles and Xie, arXiv:2212.09748, 2023. The diffusion hyperparameters are ADM's, which the paper retains wholesale: a linear variance schedule over 1000 steps from 1e-4 to 2e-2.
The training objective here is the simple one — mean squared error on
the noise. The covariance the decoder emits is returned rather than
trained; ADM's full objective adds a variational term for it, which
the paper inherits but does not modify and which a caller can build
from DiTOutput.variance_pred.
Examples
>>> import lucid
>>> from lucid.models.generative.dit import (
... DiTConfig, DiTForImageGeneration)
>>> config = DiTConfig(sample_size=8, patch_size=2, hidden_size=32,
... depth=2, num_heads=4, num_classes=10)
>>> model = DiTForImageGeneration(config).eval()
>>> model.generate(2, steps=3).samples.shape
(2, 4, 8, 8)Used by 2
Constructors
1Instance methods
2generate(n_samples: int = 1, labels: Tensor | None = None, steps: int | None = None, eta: float = 0.0, noise: Tensor | None = None, device: str | None = None)Sample by running the reverse diffusion process.
Parameters
n_samplesint= 1How many to draw.
(n_samples,) class indices. None samples the
unconditional field.stepsint or None= NoneDenoising steps. Defaults to the full training schedule.
etafloat= 0.0Interpolates the reverse step between DDIM at
0 and DDPM
at 1. The paper's numbers are at the DDPM end — it
follows ADM and reports FID over 250 DDPM steps — so
steps=250, eta=1.0 is the protocol to compare against.
The default is deterministic instead, because a sampler that
draws on every call cannot be reproduced without a seed and
this family takes no seed.Starting latent. Drawn when absent.
devicestr or None= NoneWhere to draw. Defaults to the model's own device.
Returns
GenerationOutputsamples of shape (n_samples, in_channels, H, W) — a
latent, which a VAE decoder turns into pixels.
Raises
ValueErrorIf
steps is not positive, or eta is outside [0, 1].Examples
>>> import lucid
>>> from lucid.models import dit_small_2_gen
>>> model = dit_small_2_gen(
... sample_size=8, hidden_size=32, depth=1, num_heads=4)
>>> model.generate(1, steps=2).samples.shape
(1, 4, 8, 8)
The paper's protocol is the stochastic end of eta:
>>> model.generate(1, steps=2, eta=1.0).samples.shape
(1, 4, 8, 8)