azula.plugins.vdm

Velocity diffusion model (VDM) plugin.

This plugin depends on the diffusion module in the crowsonkb/v-diffusion-pytorch repository. To use it, clone the repository to your machine

git clone https://github.com/crowsonkb/v-diffusion-pytorch

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

import sys; sys.path.append("path/to/v-diffusion-pytorch")
...
from azula.plugins import vdm

Classes

VelocityDenoiser

Creates a velocity denoiser.

Functions

load_model

Loads a pre-trained VDM denoiser.

Descriptions

class azula.plugins.vdm.VelocityDenoiser(backbone, schedule=None)

Creates a velocity denoiser.

Parameters:

• backbone (Module) – A time conditional network.

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

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

Loads a pre-trained VDM denoiser.

Parameters:

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

• kwargs – Keyword arguments passed to torch.load.

Returns:

A pre-trained denoiser.

Return type:

Denoiser