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

model_cards

Returns a key-card mapping 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, clip_mean=False, learn_var=False)

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.

  • clip_mean (bool) – Whether the mean \(\mu_\phi(x_t)\) is clipped to \([-1, 1]\) or not during evaluation.

  • learn_var (bool) – Whether the variance \(\Sigma_\phi(x_t)\) is learned or not. For pre-trained models, the learned variance is indicative, but inexact.

azula.plugins.adm.model_cards()

Returns a key-card mapping 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