Diffusion Gym

Library for reward adaptation of any pre-trained flow model on any data modality.

Diffusion Gym Teaser

Installation

In order to install diffusiongym, execute the following command:

pip install diffusiongym

High-level overview

Diffusion and flow models are largely agnostic to their data modality. They only require that the underlying data type supports a small set of operations. Building on this idea, diffusiongym is designed to be fully modular. You only need to provide the following:

  • Data type YourDataType that implements DDProtocol, which defines some functions necessary for interacting with it as a flow model.

  • Base model BaseModel[YourDataType], which defines the scheduler, how to sample \(p_0\), how to compute the forward pass, and how to preprocess and postprocess data.

  • Reward function Reward[YourDataType].

Once these are defined, you can sample from the flow model and apply reward adaptation methods, such as Value Matching.

Table of contents

How To

API Reference