API Reference

This page provides a summary of the main classes and their interfaces in the denoising-diffusion-pytorch library.

Core 2D Image Models

These are the primary classes for generating 2D images.

denoising_diffusion_pytorch.Unet

The core U-Net architecture for image-based diffusion.

class Unet(
    dim,
    init_dim = None,
    out_dim = None,
    dim_mults = (1, 2, 4, 8),
    channels = 3,
    self_condition = False,
    learned_variance = False,
    flash_attn = False
)

Parameters:

• dim (int): Base channel dimension.
• dim_mults (tuple): Multipliers for dim at each U-Net stage.
• channels (int): Number of input image channels.
• self_condition (bool): Enables self-conditioning on the model's own predictions.
• learned_variance (bool): If True, doubles the output channels to predict variance for LearnedGaussianDiffusion.
• flash_attn (bool): Enables FlashAttention for efficiency.

denoising_diffusion_pytorch.GaussianDiffusion

Manages the DDPM/DDIM process for training and sampling.

class GaussianDiffusion(
    model,
    *,
    image_size,
    timesteps = 1000,
    sampling_timesteps = None,
    objective = 'pred_v',
    beta_schedule = 'sigmoid',
    ddim_sampling_eta = 0.,
    auto_normalize = True
)

Parameters:

• model (nn.Module): The Unet model to use.
• image_size (int|tuple): The size of the input images.
• timesteps (int): Total number of diffusion steps for training.
• sampling_timesteps (int): Number of steps for DDIM sampling. If None, uses DDPM.
• objective (str): The prediction objective ('pred_noise', 'pred_x0', or 'pred_v').
• beta_schedule (str): The noise schedule ('linear', 'cosine', or 'sigmoid').

Methods:

• sample(batch_size=16): The main sampling method. Automatically chooses between DDPM and DDIM.
• p_sample_loop(shape): Performs DDPM sampling.
• ddim_sample(shape): Performs DDIM sampling.
• forward(img): Computes the training loss for a batch of images.

denoising_diffusion_pytorch.Trainer

A helper class to manage the training loop.

class Trainer(
    diffusion_model,
    folder,
    *,
    train_batch_size = 16,
    train_lr = 1e-4,
    train_num_steps = 100000,
    ema_decay = 0.995,
    amp = False,
    calculate_fid = True
)

Methods:

• train(): Starts the training process.
• save(milestone): Saves a model checkpoint.
• load(milestone): Loads a model checkpoint.

Core 1D Sequence Models

These classes are for generating 1D sequences.

denoising_diffusion_pytorch.Unet1D

A U-Net architecture adapted for 1D data.

class Unet1D(
    dim,
    dim_mults=(1, 2, 4, 8),
    channels = 3,
    self_condition = False
)

denoising_diffusion_pytorch.GaussianDiffusion1D

The diffusion process manager for 1D data.

class GaussianDiffusion1D(
    model, 
    *,
    seq_length,
    timesteps = 1000,
    objective = 'pred_noise'
)

denoising_diffusion_pytorch.Trainer1D

A helper class for training 1D diffusion models.

class Trainer1D(
    diffusion_model,
    dataset,
    *,
    train_batch_size = 16,
    train_lr = 1e-4,
    train_num_steps = 100000
)