A unified PyTorch codebase for diffusion and flow models, rectified-flow simplifies training and inference by offering a flexible, easy-to-use platform for rapid prototyping. It targets researchers and engineers working with generative models, providing a streamlined approach to implementing and experimenting with Rectified Flow (RF) and diffusion techniques. The library's core benefit is its automated algorithm derivation and seamless integration capabilities.

How It Works

The project employs a unified Ordinary Differential Equation (ODE) framework for both Rectified Flow and diffusion models, incorporating techniques like flow matching and reflow for accelerated sampling. Its key innovation is a symbolic solver that automates the derivation of algorithms and formulas for various model representations (score functions, velocity fields, noise predictions), eliminating manual derivation. This approach also facilitates easy conversion between different interpolation schemes and integration with state-of-the-art models like Flux.

Quick Start & Requirements

• Installation:
  • pip install rectified-flow
  • From source: Requires Python 3.10 (via conda environment), cloning the repository, installing requirements.txt, and then pip install -e ..
• Prerequisites: Python 3.10, standard ML dependencies listed in requirements.txt.
• Links:
  • Blog: https://rectifiedflow.github.io
  • Lecture Notes: https://www.cs.utexas.edu/~lqiang/PDF/flow_book.pdf

Highlighted Details

• Achieves strong performance on CIFAR10 with UNet (FID 2.496) and DiT (FID 3.678) models.
• Features a symbolic solver for automated derivation of flow/diffusion algorithms and automatic conversion between interpolation schemes.
• Provides Diffusers-style training scripts utilizing Accelerate for efficient multi-GPU training.
• Offers straightforward integration with state-of-the-art models such as Flux.

Maintenance & Community

The repository includes author information via a citation but does not explicitly list community channels (e.g., Discord, Slack) or detailed contributor/sponsorship information.

Licensing & Compatibility

The codebase is distributed under the MIT License, which is permissive for commercial use and integration into closed-source projects.

Limitations & Caveats

The MyStochasticSampler is theoretically sound only when the source distribution ($\pi_0$) is a zero-mean Gaussian and the coupling $(X_0, X_1)$ is independent. Users should proceed with caution if these conditions are not met.