azula.plugins.edm

Elucidated diffusion model (EDM) plugin.

This plugin depends on the torch_utils and training modules in the NVlabs/edm repository. To use it, clone the repository to your machine

git clone https://github.com/NVlabs/edm

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

import sys; sys.path.append("path/to/edm")
...
from azula.plugins import edm

References

Elucidating the Design Space of Diffusion-Based Generative Models (Karras et al., 2022)
https://arxiv.org/abs/2206.00364

Classes

ElucidatedSchedule

Creates an elucidated noise schedule.

ElucidatedDenoiser

Creates an elucidated denoiser.

Functions

model_cards

Returns a key-card mapping of available pre-trained models.

load_model

Loads a pre-trained EDM denoiser.

Descriptions

class azula.plugins.edm.ElucidatedSchedule(sigma_min=0.002, sigma_max=80.0, rho=7.0)

Creates an elucidated noise schedule.

\[\begin{split}\alpha_t & = 1 \\ \sigma_t & = \left( (1 - t) \, {\sigma_\min}^\frac{1}{\rho} + t \, {\sigma_\max}^\frac{1}{\rho} \right)^\rho\end{split}\]
Parameters:

  • sigma_min (float) – The initial noise scale \(\sigma_\min \in \mathbb{R}_+\).

  • sigma_max (float) – The final noise scale \(\sigma_\max \in \mathbb{R}_+\).

  • rho (float) – A hyper-parameter \(\rho \in \mathbb{R}_+\).

class azula.plugins.edm.ElucidatedDenoiser(backbone, schedule=None)

Creates an elucidated denoiser.

Parameters:
azula.plugins.edm.model_cards()

Returns a key-card mapping of available pre-trained models.

azula.plugins.edm.load_model(key)

Loads a pre-trained EDM denoiser.

Parameters:

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

Returns:

A pre-trained denoiser.

Return type:

GaussianDenoiser