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 you machine

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

and add it to your Python path.

import sys; sys.path.append("path/to/guided-diffusion")

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 model.

make_model

Builds an ADM model.

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 model.

Parameters:

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

  • kwargs – Keyword arguments passed to torch.hub.load.

Returns:

A pre-trained denoiser.

Return type:

ImprovedDenoiser

azula.plugins.adm.make_model(learned_var=True, schedule_name='linear', timesteps=1000, attention_resolutions={8, 16, 32}, channel_mult=(1, 2, 3, 4), dropout=0.0, image_size=64, num_channels=128, num_classes=None, num_heads=1, num_head_channels=64, num_res_blocks=3, **kwargs)

Builds an ADM model.

Parameters:

  • learned_var (bool) – Whether the variance term is learned or not.

  • schedule_name (str) – The beta schedule name.

  • timesteps (int) – The number of schedule time steps.

The remaining arguments are for the guided_diffusion.unet.UNetModel backbone.

Returns:

A denoiser.

Return type:

ImprovedDenoiser