azula.plugins.adm

Ablated diffusion model (ADM) plugin.

This plugin depends on the guided_diffusion module in the openai/guided-diffusion repository. To use it, clone the repository to your machine

git clone https://github.com/openai/guided-diffusion

and add it to your Python path before importing the plugin.

import sys; sys.path.append("path/to/guided-diffusion")
...
from azula.plugins import adm

References

Diffusion Models Beat GANs on Image Synthesis (Dhariwal et al., 2021)
https://arxiv.org/abs/2105.05233

Classes

BetaSchedule

Creates a named beta schedule.

ImprovedDenoiser

Creates an improved DDPM denoiser.

Functions

list_models

Returns the list of available pre-trained models.

load_model

Loads a pre-trained ADM denoiser.

Descriptions

class azula.plugins.adm.BetaSchedule(name='linear', steps=1000)

Creates a named beta schedule.

Parameters:

  • name (str) – The schedule name.

  • steps (int) – The number of steps.

class azula.plugins.adm.ImprovedDenoiser(backbone, schedule)

Creates an improved DDPM denoiser.

References

Improved Denoising Diffusion Probabilistic Models (Nichol et al., 2021)
https://arxiv.org/abs/2102.09672
Parameters:

  • backbone (Module) – A discrete time conditional network.

  • schedule (Schedule) – A beta schedule.

azula.plugins.adm.list_models()

Returns the list of available pre-trained models.

azula.plugins.adm.load_model(key, **kwargs)

Loads a pre-trained ADM denoiser.

Parameters:

  • key (str) – The pre-trained model key.

  • kwargs – Keyword arguments passed to torch.load.

Returns:

A pre-trained denoiser.

Return type:

GaussianDenoiser