flow_matching is a PyTorch library for Flow Matching algorithms, featuring continuous and discrete implementations. It includes examples for both text and image modalities. This repository is part of Flow Matching Guide and Codebase.
This repository requires Python 3.9 and Pytorch 2.1 or greater. To install the latest version run:
pip install flow_matching
The core and example folders are structured in the following way:
.
├── flow_matching # Core library
│ ├── loss # Loss functions
│ │ └── ...
│ ├── path # Path and schedulers
│ │ ├── ...
│ │ └── scheduler # Schedulers and transformations
│ │ └── ...
│ ├── solver # Solvers for continuous and discrete flows
│ │ └── ...
│ └── utils
│ └── ...
└── examples # Synthetic, image, and text examples
├── ...
├── image
│ └── ...
└── text
└── ...To create a conda environment with all required dependencies, run:
conda env create -f environment.yml
conda activate flow_matching
Install pre-commit hook. This will ensure that all linting is done on each commit
pre-commit install
Install the flow_matching package in an editable mode:
pip install -e .
We provide training examples. Under this folder, you can find synthetic data for continuous, discrete, and Riemannian Flow Matching. We also provide full training examples (continuous and discrete) on CIFAR10 and face-blurred ImageNet, and a scalable discrete Flow Matching example for text modeling.
In this version, we don't release pre-trained models. All models under examples can be trained from scratch by a single running command.
Please follow the contribution guide.
The code in this repository is CC BY-NC licensed. See the LICENSE for details.
If you found this repository useful, please cite the following.
@misc{lipman2024flowmatchingguidecode,
title={Flow Matching Guide and Code},
author={Yaron Lipman and Marton Havasi and Peter Holderrieth and Neta Shaul and Matt Le and Brian Karrer and Ricky T. Q. Chen and David Lopez-Paz and Heli Ben-Hamu and Itai Gat},
year={2024},
eprint={2412.06264},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2412.06264},
}