Flow Matching

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.

Installation

This repository requires Python 3.9 and Pytorch 2.1 or greater. To install the latest version run:

pip install flow_matching

Repository structure

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 
            └── ...

Development

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 .

FAQ

I want to train a Flow Matching model, where can I find the training code?

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.

Do you release pre-trained models?

In this version, we don't release pre-trained models. All models under examples can be trained from scratch by a single running command.

How to contribute to this codebase?

Please follow the contribution guide.

License

The code in this repository is CC BY-NC licensed. See the LICENSE for details.

Citation

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}, 
}