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

AblatedDenoiser

Creates an ablated denoiser.

Functions

load_model

Loads a pre-trained ADM denoiser.

Descriptions

class azula.plugins.adm.AblatedDenoiser(backbone, schedule=None, clip_mean=False, learn_var=False, discrete_schedule='linear', discrete_steps=1000)

Creates an ablated denoiser.

Parameters:

  • backbone (Module) – A time conditional network.

  • schedule (Schedule) – A noise schedule. If None, use azula.noise.VPSchedule instead.

  • 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^2_\phi(x_t)\) is learned or not. For pre-trained models, the learned variance is indicative, but inexact.

forward(x_t, t, label=None, **kwargs)
Parameters:

  • x_t (Tensor) – A noisy tensor \(x_t\), with shape \((B, 3, H, W)\).

  • t (Tensor) – The time \(t\), with shape \(()\) or \((B)\).

  • label (Tensor | None) – The class label \(c\) as an integer, with shape \((B)\).

  • kwargs – Optional keyword arguments.

Returns:

The Gaussian \(\mathcal{N}(X \mid \mu_\phi(x_t \mid c), \sigma^2_\phi(x_t \mid c))\).

Return type:

GaussianPosterior

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

Loads a pre-trained ADM denoiser.

Parameters:

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

  • kwargs – Keyword arguments passed to torch.load.

Returns:

A pre-trained denoiser.

Return type:

Denoiser