rectified-flow-pytorch: a reference implementation for flow straightening, reflow and its follow-ups

The problem: turning a curved probability path into something a solver can integrate in few steps

Rectified flow, as described in the Liu, Gong and Liu paper cited in the README, learns a velocity field that transports a noise distribution to a data distribution along paths that are as straight as possible. Straight paths are the point. A generative model that integrates a curved ODE needs many solver steps; a model whose paths are nearly straight can be sampled with far fewer. The package exists so that you can train that velocity field, then run the reflow procedure that straightens it further, without reimplementing the objective from the paper each time. The audience is narrow: researchers and engineers who already know what a flow matching objective looks like and want a working PyTorch reference to modify. The README does not claim to serve application developers who want a text-to-image endpoint. It shows a Unet, a noise tensor, a loss, and a sample.

RectifiedFlow, Reflow and Trainer: the three objects the README actually exposes

The usage example builds a Unet with dim = 64, wraps it in RectifiedFlow(model), and calls the wrapper on a batch of images shaped (1, 3, 256, 256). That call returns a loss. The wrapper is therefore a training module, not a sampler alone: calling it produces the objective, and loss.backward() follows. Sampling is a separate call, rectified_flow.sample(), and the README asserts the sampled tensor matches the image shape on all dimensions except the batch. Reflow is layered on top. The README constructs Reflow(rectified_flow) from an already-trained wrapper, then computes reflow_loss = reflow() with no arguments, and notes that you can reflow repeatedly by redefining Reflow(reflow.model) and looping. That sentence is the most informative line in the document about the intended workflow: reflow consumes the previous stage's model, and the straightening is iterative by design. Trainer sits above both. It takes the flow object, a dataset, a step count and a results folder, and is invoked by calling the trainer instance. The README states that samples are saved periodically into results_folder. Nothing in the supplied material describes the sampling schedule, the optimizer, the batch size, or how the trainer decides when to write a sample.

Install and the three commands that get an example running

The install line is pip install rectified-flow-pytorch. For the bundled example, the README gives pip install .[examples], which implies an examples extra defined in the packaging metadata, followed by python train_oxford.py. The README captions an Oxford Flowers sample image with the note 32 batch size, 11k steps oxford flowers, so the example script is the path by which the maintainer produced that figure. The Trainer constructor in the README takes four named arguments: the flow object positionally, then dataset = img_dataset, num_train_steps = 70_000, and results_folder = './results'. The dataset is built with ImageDataset(folder = './path/to/your/images', image_size = 256). Those are the only configuration keys visible in the supplied text. There is no mention of a config file, environment variables, or CLI flags beyond running the script directly. If you need to control precision, gradient accumulation, or checkpoint frequency, the README does not tell you where those knobs live.

Where the documentation runs out: no checkpointing, no distributed story, no evaluation

The README shows training a Unet at dim = 64 on 256x256 images and a 70,000-step Trainer run, but it never mentions saving or loading model weights. For a package whose whole purpose is a multi-stage procedure (train, reflow, reflow again), the absence of a stated checkpoint format is a practical gap, because reflow needs the previous stage's model and the README's own example assumes you still have it in memory. There is also no stated support for distributed training, even though the Trainer is described as based on accelerate, a library whose main purpose is exactly that. The two facts sit awkwardly together: accelerate is named, but nothing in the README says whether the Trainer uses it for multi-GPU, mixed precision, or only for device placement. Evaluation is absent as well. The only quality signal offered is a sample image, which is not a metric. Treat the package as a reference implementation whose correctness you verify against the papers, not as a framework with a tested training contract.

The follow-up methods are named but not mapped to code

The appreciation section credits contributors for LSD flow, SoFlow, Split Mean Flow, Unconstrained Alignment (UA) Flow, and Laplacian multiscale flow matching. The citation list is longer still, covering immiscible diffusion, consistency flow matching, hyper-connections, F5R-TTS, mean flows, noise conditioning, flow Q-learning, flow matching policy gradients, and a denoising-generative-models paper. That is a lot of surface area for one package. What the README does not do is tell you which import name corresponds to which paper. The usage examples only ever import RectifiedFlow, Reflow, Unet, ImageDataset and Trainer. If you want the UA Flow or Split Mean Flow variant, you will be reading the package's module tree rather than its README. This is the single biggest documentation weakness in the supplied material, and it matters because the description explicitly frames the package as covering follow-up research.

The alternative: a full training framework instead of a research reference

The obvious comparison is a complete diffusion or flow training framework rather than a single-author reference package. The difference is in what is assumed for you. A full framework typically ships a configuration system, checkpoint management, dataset registries, evaluation hooks and a serving path, and asks you to adapt your model to its interfaces. rectified-flow-pytorch does the opposite: it hands you a small set of classes and expects you to write the surrounding training loop, or to accept the minimal Trainer as-is. The trade is legibility against operational completeness. If you want to read the reflow objective in about a screen of code and change it, the reference package wins. If you want a run that survives a node failure and resumes from a checkpoint, the reference package does not describe how to do that, and a full framework does. That is the decision, and it is not a close one for anyone shipping a model.

Maintenance, versioning and the MIT licence

The repository is not archived and the last push recorded is 2026-08-02. Releases listed are 0.6.6 on 2026-01-28, 0.6.2 on 2026-01-23, and 0.6.0 on 2026-01-23, so three releases landed within five days of each other in January and the next recorded release came roughly a week later. That cadence is consistent with active development, and it also means the API can move between minor versions. Pin the version you build against. The licence is MIT, which permits commercial use and modification provided the copyright notice and permission notice are included; this is a statement of what the licence text allows, not legal advice, and you should read the LICENSE file in the repository and consult your own counsel if the distinction matters to you. Note that the package's own MIT licence does not automatically cover the papers it implements or any pretrained weights you might pair with it, and the README ships no weights.

Editorial conclusion

Adopt it if you are prototyping rectified flow, reflow, or one of the follow-up variants in PyTorch and you want a small, MIT-licensed codebase to read and modify. Do not adopt it as a production training or inference pipeline: the README shows no checkpointing, no distributed or mixed-precision configuration, and no serving path, and the Trainer is described only through a single constructor call. Before committing, read the module list and the example training script in the repository, and confirm which of the cited papers each exported class actually implements, because the README names the papers but never maps them to modules.

Official sources

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