class

DiTModel

extendsPretrainedModel

DiTModel(config: DiTConfig)

source edit

The denoising network: latent patches in, noise prediction out.

Parameters

configDiTConfig

The variant to build.

Attributes

patch_embedlucid.nn.Conv2d

Patchifying projection — stride equals kernel, so patches do not overlap.

pos_embedTensor

Frozen sine-cosine table, registered as a buffer.

blockslucid.nn.ModuleList

The transformer blocks, of whichever kind the conditioning names.

finalDiTFinalLayer

Modulated norm and the projection back to patch space.

Notes

Reference: Peebles and Xie, "Scalable Diffusion Models with Transformers", ICCV, 2023 (arXiv:2212.09748). Backbone shapes are Table 1.

Examples

>>> import lucid
>>> from lucid.models.generative.dit import DiTConfig, DiTModel
>>> config = DiTConfig(sample_size=8, patch_size=2, hidden_size=32,
...                    depth=2, num_heads=4, num_classes=10)
>>> model = DiTModel(config).eval()
>>> latent = lucid.randn((1, 4, 8, 8))
>>> labels = lucid.tensor([3], dtype=lucid.int64)
>>> model(latent, lucid.tensor([10.0]), labels).sample.shape
(1, 8, 8, 8)

Used by 2

Constructors

dunder

init

→None

__init__(config: DiTConfig)

source edit

Initialise the network. See the class docstring for parameters.

Instance methods

forward

→DiffusionModelOutput

forward(latent: Tensor, timesteps: Tensor, labels: Tensor | None = None)

source edit

Predict the noise (and covariance) in a noised latent.

Parameters

latentTensor

(B, in_channels, H, W) — the noised latent

z_t

timestepsTensor

(B,) diffusion timesteps.

labelsTensor or None= None

(B,) class indices. None uses the null token.

Returns

DiffusionModelOutput

sample of shape (B, out_channels, H, W) — noise, and the covariance beside it when learn_sigma. loss stays None: the objective lives in DiTForImageGeneration.

>>> import lucid >>> from lucid.models.generative.dit import DiTConfig, DiTModel >>> config = DiTConfig(sample_size=8, patch_size=2, hidden_size=32, ... depth=2, num_heads=4, num_classes=10) >>> model = DiTModel(config).eval() >>> latent = lucid.randn((1, 4, 8, 8)) >>> labels = lucid.tensor([3], dtype=lucid.int64) >>> model(latent, lucid.tensor([10.0]), labels).sample.shape (1, 8, 8, 8)