IDDPM

class dmme.diffusion_models.IDDPM(model: Module, timesteps: int = 1000, loss_type='hybrid', gamma=0.001, schedule: str = 'cosine', offset=0.008, start: float = 0.0001, end: float = 0.02)

Improved DDPM with cosine variance schedule and learned variance

Parameters:

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

• timesteps – total timesteps \(T\)

• loss_type – loss type to use either “hybrid” or “simple”

• gamma – \(\gamma\) in hybrid loss

• shcedule – variance schedule to use either “linear” or “cosine”

• offset – default offset to use if cosine schedule is used

• start – default linear variance schedule start value

• end – default linear variance schedule end value

training_step(x_0)

Computes hybrid loss for improved DDPM

Parameters:

x_0 – sample image to add noise and denoise for training

Returns:

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

sampling_step(x_t, t)

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)\)

forward_model(x_t, t, beta_t, alpha_bar_t, alpha_bar_t_minus_one)

Applies forward to internal model

Parameters:

• x – input image passed to internal model

• t – timestep passed to internal model