rectified-flow-pytorch: lucidrains' PyTorch implementation of rectified flow and its follow-up research

What rectified-flow-pytorch is for

Rectified flow is a generative modelling method: instead of learning a reverse diffusion process, it learns a velocity field that transports a simple distribution to the data distribution along paths that are trained to be straight. The repository implements that core idea and, in the author's words, "some of its followup research / improvements in Pytorch". The audience is narrow and identifiable. It is a research implementation, not an application. The evidence is in the repository layout: the top level holds roughly twenty standalone training scripts, one per method, including train_oxford.py, train_mean_flow.py, train_self_flow.py, train_split_mean_flow.py, train_soflow.py, train_lsd_flow.py, train_recursive_flow.py, train_uaflow.py, train_fpo.py and several reinforcement-learning flavoured ones such as train_value_flow_ppo.py and train_prob_flow_ppo.py. Each script is a thin driver over the package. If you want a library that hides the research, this is the wrong shape. If you want to see how a specific paper's loss is wired up and then change it, the layout is exactly what you want. The package metadata classifies it as Development Status 4 - Beta, and the version in pyproject.toml is 0.11.2 while the most recent GitHub release listed is 0.6.6, so the release tags trail the source tree.

The mechanism: a model, a velocity loss, and a reflow loop

The README's usage example is the clearest description of the data flow. You construct a model, in the example a Unet with dim = 64, and wrap it in RectifiedFlow. Calling the wrapper on a batch of images returns a loss, and loss.backward() is what trains. The wrapper owns the interpolation between noise and data and the velocity target; the model only has to predict the velocity. Sampling is a method on the same object, rectified_flow.sample(), and the README asserts that the sampled tensor keeps the channel, height and width of the input images. That assertion is worth reading as a contract: the wrapper is expected to produce samples shaped like the training data without you reshaping anything.

The second mechanism is reflow, and it is the part people usually get wrong. Reflow does not train the same model again on the same pairs. You first train a RectifiedFlow model on real images, then wrap that trained object in Reflow and take a second loss from it. The README notes you can repeat this by redefining Reflow(reflow.model) and looping, which is the paper's iterative straightening procedure expressed as object construction. There is no separate reflow schedule, no teacher checkpoint file and no config key for the number of reflow rounds. The loop is your code.

A Trainer built on accelerate is the third piece. It takes the rectified flow object, a dataset and a step count, and writes samples to a results folder periodically. That is the whole documented training surface. Anything about checkpoint resumption, distributed settings or metric logging is not described in the README.

Install and run a first training step

The package is on PyPI and the README gives a single install command. Use a virtual environment; the dependency list in pyproject.toml is long (torch, torchvision, accelerate, einops, einx, ema-pytorch, hyper-connections, rotary-embedding-torch, x-transformers, torchdiffeq, scipy, pillow) and you do not want it in your system interpreter.

$ pip install rectified-flow-pytorch

After that, the smallest real use is the README's own example: build a Unet, wrap it, run one forward pass on a random tensor, backpropagate, and sample. This tells you whether the install and the tensor shapes line up before you point anything at a dataset.

import torch
from rectified_flow_pytorch import RectifiedFlow, Unet

model = Unet(dim = 64)
rectified_flow = RectifiedFlow(model)

images = torch.randn(1, 3, 256, 256)
loss = rectified_flow(images)
loss.backward()

sampled = rectified_flow.sample()
assert sampled.shape[1:] == images.shape[1:]

The assertion is the thing to watch. If it fires, the model you passed in does not match the shape contract the wrapper assumes, and no amount of training will fix it.

For an actual run, the README shows the ImageDataset and Trainer path. The dataset takes a folder of images and an image size, and the trainer takes a step count and a results folder where samples are written periodically.

import torch
from rectified_flow_pytorch import RectifiedFlow, ImageDataset, Unet, Trainer

model = Unet(dim = 64)
rectified_flow = RectifiedFlow(model)

img_dataset = ImageDataset(
    folder = './path/to/your/images',
    image_size = 256
)

trainer = Trainer(
    rectified_flow,
    dataset = img_dataset,
    num_train_steps = 70_000,
    results_folder = './results'
)

trainer()

If you would rather start from a working script, the repository ships examples. The README instructs installing the example extras and then running the Oxford Flowers script.

$ pip install .[examples]
$ python train_oxford.py

Note that the extras are split: pyproject.toml defines examples, examples_ql, examples_fpo and test separately, so the reinforcement-learning and FPO scripts need their own extras rather than the plain examples group.

Where the documentation stops

The README is a usage sketch, not a reference. It documents RectifiedFlow, Reflow, Unet, ImageDataset and Trainer through short snippets and nothing else. There is no API listing, no description of the loss terms, no explanation of the sampler's step count or solver, and no statement about how many sampling steps you should use. torchdiffeq is a dependency, which suggests an ODE solver is involved, but the README never says which solver is used or how to change it. If solver choice matters for your work, you will be reading the source in rectified_flow_pytorch/ rather than the documentation.

The version situation is a second practical hazard. The most recent release listed is 0.6.6 from 2026-01-28, while pyproject.toml in the repository declares 0.11.2. A pip install therefore does not obviously give you the same code as the main branch. If you are trying to reproduce a result from one of the cited follow-up papers, check which methods are exported by the version you installed before assuming the repository's training scripts will run against it. The tests directory exists, but the README does not describe what it covers.

The last push was on 2026-08-02, so the repository is not archived and has been touched recently, but the release cadence and the source version do not move together and that is the number that affects you as a user.

When to reach for flow matching instead

The honest alternative is not a different library, it is a different formulation. Flow matching, as popularised in the image-generation literature and cited here through Esser et al.'s scaling paper, trains a conditional velocity field against an explicitly chosen probability path between noise and data, with a closed-form target for that path. Rectified flow instead trains on the coupling it currently has and then straightens it by reflowing. The practical difference shows up in the training loop: flow matching gives you a target you can write down per batch, while rectified flow adds a second stage where you generate pairs from your own model and retrain on them. If your goal is a single-stage training run with a fixed objective, flow matching is the simpler object and libraries built around it, such as the diffusers training scripts, are better documented. If your goal is to study the straightening effect, or to cut the number of sampling steps by making the learned paths straighter, reflow is the point and this repository is one of the few places where the loop is short enough to read in an afternoon. The paper list in the README makes the trade-off explicit: it cites both the original rectified flow paper and the scaling rectified flow transformers work, so the author is not claiming one replaces the other.

Licence and the cost of keeping up

The project is MIT licensed, and pyproject.toml declares license = { file = "LICENSE" }. MIT is permissive: you can use, modify and redistribute the code, including in closed products, provided the copyright notice and permission notice are kept. That is the licence text, not legal advice; if you are shipping a model trained with this code, the licence covers the code, not the weights or the data you trained on, and those are separate questions.

The upgrade cost is the more interesting one. The dependency list pins minimum versions rather than upper bounds (torch>=2.0, einops>=0.8.0, x-transformers>=2.3.21, hyper-connections>=0.1.8, ema-pytorch>=0.5.2), and several of those are the same author's packages. That means an upgrade of the package can pull a newer x-transformers or hyper-connections, and the Unet and the various flow variants in this repository are built on them. Because the published release lags the source tree, upgrading to a version that contains a method you want also means accepting whatever dependency movement happened in between. Pinning the whole environment is the cheap insurance here.

Editorial conclusion

Adopt rectified-flow-pytorch if you are a researcher or engineer who wants to read or modify flow-matching training code and already have a PyTorch and accelerate environment. Do not adopt it if you need a documented, stable API or a training service, because the README documents RectifiedFlow, Reflow, ImageDataset and Trainer and little else. Verify first that the method you need is actually exported by the installed version, that the model you pass in satisfies the Unet-shaped interface the examples assume, and that the optional example dependencies install on your Python version.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. lucidrains/rectified-flow-pytorch on GitHub
4. README
5. Releases