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

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)[source]

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)[source]

Creates an elucidated denoiser.

Parameters:
forward(x_t, t, label=None, **kwargs)[source]
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 a one-hot vector, with shape \((*, 1000)\).

  • kwargs – Optional keyword arguments.

Returns:

The Dirac delta \(\delta(X - \mu_\phi(x_t \mid c))\).

Return type:

DiracPosterior

azula.plugins.edm.load_model(name)[source]

Loads a pre-trained EDM denoiser.

Parameters:

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

Returns:

A pre-trained denoiser.

Return type:

Denoiser