DDPM

class dmme.diffusion_models.DDPM(model: Module, timesteps: int = 1000, start: float = 0.0001, end: float = 0.02)

Training and Sampling for DDPM

Parameters:

• model – model predicting noise from data, \(\epsilon_\theta(x_t, t)\)

• timesteps – total timesteps \(T\)

• start – linear variance schedule start value

• end – linear variance schedule end value

training_step(x_0: Tensor) → Tensor

Training step except for optimization

Parameters:

x_0 – image from dataset

Returns:

loss, \(L_\text{simple}\)

sampling_step(x_t: Tensor, t: Tensor) → Tensor

Denoise image by sampling from \(p_\theta(x_{t-1}|x_t)\)

Parameters:

• x_t – image of shape \((N, C, H, W)\)

• t – starting \(t\) to sample from, a tensor of shape \((N,)\)

Returns:

denoised image of shape \((N, C, H, W)\)

generate(img_size: Tuple[int, int, int, int]) → Tensor

Generate image of shape \((N, C, H, W)\) by running the full denoising steps

Parameters:

img_size – image size to generate as a tuple \((N, C, H, W)\)

Returns:

generated image of shape \((N, C, H, W)\) as a tensor

forward(x: Tensor, t: Tensor) → Tensor

Applies forward to internal model

Parameters:

• x – input image passed to internal model

• t – timestep passed to internal model