Metadata-Version: 2.4
Name: rectified-flows
Version: 0.0.13
Summary: Add your description here
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: beartype>=0.21.0
Requires-Dist: cloudpickle>=3.1.1
Requires-Dist: datasets>=3.6.0
Requires-Dist: diffrax>=0.7.0
Requires-Dist: einops>=0.8.1
Requires-Dist: einx>=0.3.0
Requires-Dist: equinox>=0.12.2
Requires-Dist: grain>=0.2.11
Requires-Dist: jax>=0.6.2
Requires-Dist: matplotlib>=3.10.3
Requires-Dist: ml-collections>=1.1.0
Requires-Dist: numpy>=2.3.1
Requires-Dist: optax>=0.2.5
Requires-Dist: pillow>=11.3.0
Requires-Dist: tqdm>=4.67.1
Dynamic: license-file

<h1 align='center'>Rectified Flow Matching</h1>

Cutting-edge and feature-rich implementation of Rectified Flows from [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flows](https://arxiv.org/abs/2209.03003) in `jax` and `equinox`.

#### Features
- Deterministic and stochastic sampling of the associated ODE and SDE respectively,
- Mixed precision optimisation,
- Array-typed to-the-teeth for dependable execution with `jaxtyping` and `beartype`.

To implement:
- [ ] Guidance by score of conditioning
- [x] Mixed precision
- [x] EMA
- [x] AdaLayerNorm
- [x] Stochastic sampling
- [x] ODE Sampling
- [x] Likelihoods
- [x] DiT
- [ ] Hyperparameter/model saving

#### Usage

```
pip install rectified_flows
```

<p align="center">
  <picture>
    <img src="assets/flow_matching_baryons.gif" alt="Your image description">
  </picture>
</p>


#### Citations

```bibtex
@misc{liu2022flowstraightfastlearning,
      title={Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow}, 
      author={Xingchao Liu and Chengyue Gong and Qiang Liu},
      year={2022},
      eprint={2209.03003},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2209.03003}, 
}
```

```bibtex
@misc{lipman2023flowmatchinggenerativemodeling,
    title={Flow Matching for Generative Modeling}, 
    author={Yaron Lipman and Ricky T. Q. Chen and Heli Ben-Hamu and Maximilian Nickel and Matt Le},
    year={2023},
    eprint={2210.02747},
    archivePrefix={arXiv},
    primaryClass={cs.LG},
    url={https://arxiv.org/abs/2210.02747}, 
}
```

```bibtex
@misc{singh2024stochasticsamplingdeterministicflow,
      title={Stochastic Sampling from Deterministic Flow Models}, 
      author={Saurabh Singh and Ian Fischer},
      year={2024},
      eprint={2410.02217},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2410.02217}, 
}
```

```bibtex
@misc{yang2024consistencyflowmatchingdefining,
      title={Consistency Flow Matching: Defining Straight Flows with Velocity Consistency}, 
      author={Ling Yang and Zixiang Zhang and Zhilong Zhang and Xingchao Liu and Minkai Xu and Wentao Zhang and Chenlin Meng and Stefano Ermon and Bin Cui},
      year={2024},
      eprint={2407.02398},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2407.02398}, 
}
```

```bibtex
@misc{peebles2023scalablediffusionmodelstransformers,
      title={Scalable Diffusion Models with Transformers}, 
      author={William Peebles and Saining Xie},
      year={2023},
      eprint={2212.09748},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2212.09748}, 
}
```
