class

DDIMScheduler

DDIMScheduler(config: StableDiffusionConfig)

source edit

Deterministic (or partly stochastic) sampling over a step subset.

Parameters

configStableDiffusionConfig

Read for num_train_timesteps, the betas and the schedule.

Attributes

alphas_cumprodTensor

\bar\alpha_t

, (num_train_timesteps,).

Notes

Reference: Song, Meng and Ermon, "Denoising Diffusion Implicit Models", ICLR, 2021 (arXiv:2010.02502); the schedule constants are the released Stable Diffusion configuration.

Kept as a plain object rather than a lucid.nn.Module: it holds no parameters, and making it a Module would put its buffers in the checkpoint of every model that uses one.

Examples

>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
...     DDIMScheduler, StableDiffusionConfig)
>>> scheduler = DDIMScheduler(StableDiffusionConfig())
>>> len(scheduler.timesteps(50))
50
>>> scheduler.timesteps(50)[0] > scheduler.timesteps(50)[-1]
True
>>> scheduler.timesteps(10)[:4]
[901, 801, 701, 601]

Used by 2

Constructors

dunder

init

→None

__init__(config: StableDiffusionConfig)

source edit

Build the schedule. See the class docstring for parameters.

Instance methods

add_noise

→Tensor

add_noise(latent: Tensor, noise: Tensor, timestep: int)

source edit

The forward process — $z_t$ from $z_0$ .

Parameters

latentTensor

(B, C, h, w) clean latent.

noiseTensor

The same shape.

timestepint

Index into the training schedule.

Returns

Tensor

$\sqrt{\bar\alpha_t} z_0 + \sqrt{1-\bar\alpha_t}\epsilon$ .

Examples

>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
...     DDIMScheduler, StableDiffusionConfig)
>>> scheduler = DDIMScheduler(StableDiffusionConfig())
>>> ones, zeros = lucid.ones((1, 4, 8, 8)), lucid.zeros((1, 4, 8, 8))
>>> scheduler.add_noise(ones, zeros, 500).shape
(1, 4, 8, 8)
A clean latent of ones with no noise reads off the signal
coefficient \sqrt{\bar\alpha_t}, and swapping the two reads
off the noise one.  At the first step the latent is almost
untouched; by the last it is almost all noise.
>>> def coefficients(t):
...     signal = scheduler.add_noise(ones, zeros, t)[0, 0, 0, 0].item()
...     noise = scheduler.add_noise(zeros, ones, t)[0, 0, 0, 0].item()
...     return signal, noise
>>> [round(c, 4) for c in coefficients(0) + coefficients(999)]
[0.9996, 0.0292, 0.0683, 0.9977]
The squares sum to one at every step, which is what keeps a
unit-variance latent at unit variance as the noise replaces it.
>>> signal, noise = coefficients(500)
>>> round(signal**2 + noise**2, 6)
1.0

step

→Tensor

step(model_output: Tensor, timestep: int, previous_timestep: int, latent: Tensor, eta: float = 0.0)

source edit

One reverse step.

Parameters

model_outputTensor

\epsilon_\theta(z_t, t, \tau_\theta(y))

timestepint

Current

t

previous_timestepint

The next index in the subsequence, or -1 for the last step, where

\bar\alpha

is taken as 1.

latentTensor

z_t

etafloat= 0.0

Interpolates DDIM (0) to DDPM (1). At 0 the trajectory is deterministic, which is what makes a seed reproducible.

Returns

Tensor

$z_{t-1}$ .

Raises

ValueError

If eta is outside [0, 1].

Examples

>>> import lucid
>>> from lucid.models.generative.stable_diffusion import (
...     DDIMScheduler, StableDiffusionConfig)
>>> scheduler = DDIMScheduler(StableDiffusionConfig())
>>> lucid.manual_seed(0)
>>> clean = lucid.randn((1, 4, 8, 8))
>>> noise = lucid.randn((1, 4, 8, 8))
>>> noised = scheduler.add_noise(clean, noise, 501)
Handed the noise that was actually added — a perfect
\epsilon_\theta — the step recovers \hat z_0 = z_0,
so going from 501 to 251 lands on the forward process at 251.
>>> stepped = scheduler.step(noise, 501, 251, noised)
>>> target = scheduler.add_noise(clean, noise, 251)
>>> bool(lucid.allclose(stepped, target, atol=1e-5))
True
At eta=0 that is the whole update, so it repeats exactly; any
positive eta adds fresh noise on top.
>>> bool((scheduler.step(noise, 501, 251, noised) == stepped).all())
True
>>> bool((scheduler.step(noise, 501, 251, noised, eta=1.0) == stepped).all())
False

timesteps

→list of int

timesteps(num_inference_steps: int)

source edit

The decreasing subsequence to sample over.

Parameters

num_inference_stepsint

How many network evaluations to spend.

Returns

list of int

Descending indices into the training schedule.

Raises

ValueError

If more steps are requested than the schedule has.

Examples

>>> from lucid.models.generative.stable_diffusion import (
...     DDIMScheduler, StableDiffusionConfig)
>>> scheduler = DDIMScheduler(StableDiffusionConfig())
>>> scheduler.timesteps(4)
[751, 501, 251, 1]
Evenly strided by 1000 // 4 and shifted by steps_offset, so
the last step is 1 rather than 0.  Asking for more steps than the
schedule was trained with is an error, not a clamp.
>>> scheduler.timesteps(1001)
Traceback (most recent call last):
    ...
ValueError: num_inference_steps must be in [1, 1000], got 1001

>>> import lucid >>> from lucid.models.generative.stable_diffusion import ( ... DDIMScheduler, StableDiffusionConfig) >>> scheduler = DDIMScheduler(StableDiffusionConfig()) >>> len(scheduler.timesteps(50)) 50 >>> scheduler.timesteps(50)[0] > scheduler.timesteps(50)[-1] True >>> scheduler.timesteps(10)[:4] [901, 801, 701, 601]

>>> import lucid >>> from lucid.models.generative.stable_diffusion import ( ... DDIMScheduler, StableDiffusionConfig) >>> scheduler = DDIMScheduler(StableDiffusionConfig()) >>> ones, zeros = lucid.ones((1, 4, 8, 8)), lucid.zeros((1, 4, 8, 8)) >>> scheduler.add_noise(ones, zeros, 500).shape (1, 4, 8, 8) A clean latent of ones with no noise reads off the signal coefficient \sqrt{\bar\alpha_t}, and swapping the two reads off the noise one. At the first step the latent is almost untouched; by the last it is almost all noise. >>> def coefficients(t): ... signal = scheduler.add_noise(ones, zeros, t)[0, 0, 0, 0].item() ... noise = scheduler.add_noise(zeros, ones, t)[0, 0, 0, 0].item() ... return signal, noise >>> [round(c, 4) for c in coefficients(0) + coefficients(999)] [0.9996, 0.0292, 0.0683, 0.9977] The squares sum to one at every step, which is what keeps a unit-variance latent at unit variance as the noise replaces it. >>> signal, noise = coefficients(500) >>> round(signal**2 + noise**2, 6) 1.0

>>> import lucid >>> from lucid.models.generative.stable_diffusion import ( ... DDIMScheduler, StableDiffusionConfig) >>> scheduler = DDIMScheduler(StableDiffusionConfig()) >>> lucid.manual_seed(0) >>> clean = lucid.randn((1, 4, 8, 8)) >>> noise = lucid.randn((1, 4, 8, 8)) >>> noised = scheduler.add_noise(clean, noise, 501) Handed the noise that was actually added — a perfect \epsilon_\theta — the step recovers \hat z_0 = z_0, so going from 501 to 251 lands on the forward process at 251. >>> stepped = scheduler.step(noise, 501, 251, noised) >>> target = scheduler.add_noise(clean, noise, 251) >>> bool(lucid.allclose(stepped, target, atol=1e-5)) True At eta=0 that is the whole update, so it repeats exactly; any positive eta adds fresh noise on top. >>> bool((scheduler.step(noise, 501, 251, noised) == stepped).all()) True >>> bool((scheduler.step(noise, 501, 251, noised, eta=1.0) == stepped).all()) False

>>> from lucid.models.generative.stable_diffusion import ( ... DDIMScheduler, StableDiffusionConfig) >>> scheduler = DDIMScheduler(StableDiffusionConfig()) >>> scheduler.timesteps(4) [751, 501, 251, 1] Evenly strided by 1000 // 4 and shifted by steps_offset, so the last step is 1 rather than 0. Asking for more steps than the schedule was trained with is an error, not a clamp. >>> scheduler.timesteps(1001) Traceback (most recent call last): ... ValueError: num_inference_steps must be in [1, 1000], got 1001