diffusion

Denoising diffusion schedules and samplers.

Minimal, dependency-free implementation of the DDIM sampler from Denoising Diffusion Implicit Models, matching the diffusers-style API (add_noise / get_velocity / set_timesteps / step) used by latent diffusion models.

Classes:

DDIMScheduler –

DDIM noise schedule and sampling step.

DDIMScheduler ¶

DDIMScheduler(
    num_train_timesteps: int = 1000,
    beta_start: float = 0.0001,
    beta_end: float = 0.02,
    prediction_type: str = "epsilon",
    set_alpha_to_one: bool = True,
)

DDIM noise schedule and sampling step.

Implements the deterministic-to-stochastic DDIM update of formula (12) in DDIM over a linear \(\beta\) schedule, with epsilon, sample and v_prediction parameterizations.

Parameters:

num_train_timesteps (int, default: 1000 ) –

Number of diffusion steps \(T\) used at training time.
beta_start (float, default: 0.0001 ) –

First value of the linear \(\beta\) schedule.
beta_end (float, default: 0.02 ) –

Last value of the linear \(\beta\) schedule.
prediction_type (str, default: 'epsilon' ) –

Quantity predicted by the model, one of "epsilon", "sample" or "v_prediction".
set_alpha_to_one (bool, default: True ) –

Use \(\bar\alpha_{-1} = 1\) for the final denoising step instead of \(\bar\alpha_0\).

Example

scheduler = DDIMScheduler(prediction_type="v_prediction")
scheduler.set_timesteps(100)
sample = torch.randn(1024, 8)
for t in scheduler.timesteps:
    sample = scheduler.step(model(sample, t), int(t), sample)

Methods:

set_timesteps –

Select the num_inference_steps evenly spaced timesteps used for sampling.
step –

Run one reverse diffusion step \(x_t \rightarrow x_{t-1}\).
add_noise –

Diffuse original_samples to the given timesteps (forward process).
get_velocity –

Compute the v_prediction target \(v_t = \sqrt{\bar\alpha_t}\,\epsilon - \sqrt{1-\bar\alpha_t}\,x_0\).

set_timesteps ¶

set_timesteps(
    num_inference_steps: int,
    device: Union[str, device, None] = None,
) -> None

Select the num_inference_steps evenly spaced timesteps used for sampling.

Parameters:

num_inference_steps (int) –

Number of denoising steps.
device (Union[str, device, None], default: None ) –

Device for the timestep tensor.

step ¶

step(
    model_output: Tensor,
    timestep: int,
    sample: Tensor,
    eta: float = 1.0,
    generator: Optional[Generator] = None,
) -> Tensor

Run one reverse diffusion step \(x_t \rightarrow x_{t-1}\).

Parameters:

model_output (Tensor) –

Model prediction at timestep (interpreted per prediction_type).
timestep (int) –

Current discrete timestep \(t\).
sample (Tensor) –

Current sample \(x_t\).
eta (float, default: 1.0 ) –

Noise scale \(\eta\) of formula (16); \(\eta = 0\) is deterministic DDIM, \(\eta = 1\) matches DDPM-level stochasticity.
generator (Optional[Generator], default: None ) –

Random generator for the added noise.

Returns:

Tensor –

The previous sample \(x_{t-1}\).

Shape

model_output: \((N, C)\) or any shape.
sample: same shape as model_output.
Output: same shape as sample.

add_noise ¶

add_noise(
    original_samples: Tensor,
    noise: Tensor,
    timesteps: Tensor,
) -> Tensor

Diffuse original_samples to the given timesteps (forward process).

Parameters:

original_samples (Tensor) –

Clean samples \(x_0\).
noise (Tensor) –

Gaussian noise of the same shape.
timesteps (Tensor) –

Per-row timesteps.

Returns:

Tensor –

The noisy samples \(x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1 - \bar\alpha_t}\,\epsilon\).

Shape

original_samples, noise: \((N, C)\) or any shape.
timesteps: \((N,)\) or broadcastable to the leading dimension.
Output: same shape as original_samples.

get_velocity ¶

get_velocity(
    sample: Tensor, noise: Tensor, timesteps: Tensor
) -> Tensor

Compute the v_prediction target \(v_t = \sqrt{\bar\alpha_t}\,\epsilon - \sqrt{1-\bar\alpha_t}\,x_0\).

Parameters:

sample (Tensor) –

Clean samples \(x_0\).
noise (Tensor) –

Gaussian noise of the same shape.
timesteps (Tensor) –

Per-row timesteps.

Returns:

Tensor –

The velocity target.

Shape

sample, noise: \((N, C)\) or any shape.
timesteps: \((N,)\) or broadcastable to the leading dimension.
Output: same shape as sample.