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

AngularSchedule

Creates an angular schedule.

VelocityDenoiser

Creates a velocity denoiser.

Functions

list_models

Returns the list of available pre-trained models.

load_model

Loads a pre-trained VDM denoiser.

Descriptions

class azula.plugins.vdm.AngularSchedule

Creates an angular schedule.

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

Creates a velocity denoiser.

Parameters:

  • backbone (Module) – A time conditional network.

  • schedule (Schedule) – A noise schedule.

azula.plugins.vdm.list_models()

Returns the list of available pre-trained models.

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

Loads a pre-trained VDM denoiser.

Parameters:

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

  • kwargs – Keyword arguments passed to torch.load.

Returns:

A pre-trained denoiser.

Return type:

GaussianDenoiser