Repository files navigation

A rectified flow implementation in Flax

This repository implements a rectified flowwhich has been proposed inFlow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flowusing JAX and Flax.

Theexperiments folder contains a use case where samples from the "Wto Moonss" data set are transportedto the "Eight Gaussian" data set. To train a model and make visualizations, call:

cd experiments/eight_gaussians_two_moonspython main.py

Shown below are samples from the two moons data set (black) that have been transportedto the eight Gaussians data set(blue). Each figure shows the transportmap after x training iterations.

To install the latest GitHub , just call the following on thecommand line:

pip install git+https://github.com/dirmeier/rflow@<RELEASE>

Simon Dirmeiersfyrbnd @ pm me