https://github.com/cloneofsimo/minrf
Minimal implementation of scalable rectified flow transformers, based on SD3's approach
https://github.com/cloneofsimo/minrf
Last synced: about 1 year ago
JSON representation
Minimal implementation of scalable rectified flow transformers, based on SD3's approach
- Host: GitHub
- URL: https://github.com/cloneofsimo/minrf
- Owner: cloneofsimo
- License: apache-2.0
- Created: 2024-05-04T05:36:50.000Z (over 2 years ago)
- Default Branch: main
- Last Pushed: 2024-07-01T20:26:13.000Z (about 2 years ago)
- Last Synced: 2025-05-07T16:12:12.563Z (over 1 year ago)
- Language: Jupyter Notebook
- Homepage:
- Size: 114 MB
- Stars: 519
- Watchers: 8
- Forks: 47
- Open Issues: 11
-
Metadata Files:
- Readme: README.md
- License: LICENSE
- Citation: CITATION.cff
Awesome Lists containing this project
README
# Minimal Implementation of Scalable Rectified Flow Transformers
> Left is the naive RF, right is the logit-normal time-sampling RF. Both are trained on MNIST.
This repository contains a minimal implementation of the rectified flow models. I've taken [SD3](https://arxiv.org/abs/2403.03206) approach of training along with [LLaMA-DiT](https://github.com/Alpha-VLLM/LLaMA2-Accessory) architecture. [Unlike my previous repo](https://github.com/cloneofsimo/minDiffusion) this time I've decided to split the file into 2: The model implementation and actual code, but you don't have to look at the model code.
Everything is still self-contained, minimal, and hopefully easy to hack. There is nothing complicated goin on if you understood the math.
# 1. *Simple* Rectified Flow, for beginners
Install torch, pil, torchvision
```
pip install torch torchvision pillow
```
Run
```bash
python rf.py
```
to train the model on MNIST from scratch.
If you are cool and want to train CIFAR instead, you can do that.
```bash
python rf.py --cifar
```
On 63'th epoch, your output should be something like:
# 2. *Massive* Rectified Flow, muP Support
This is for gigachads who wants to train Imagenet instead. Don't worry! IMO Imagenet is the new MNIST, and we will use my [imagenet.int8](https://huggingface.co/datasets/cloneofsimo/imagenet.int8) dataset for this.
First go to advanced dir, download the dataset.
```bash
cd advanced
pip install hf_transfer # just do install this.
bash download.sh
```
This shouldn't take more than 5 min if your network is decent.
Run
```bash
bash run.sh
```
to train the model. This will train Imagenet from scratch, do a muP grid search to find the aligned basin for the loss function, you unlock the zero-shot LR transfer for Rectified Flow models!
This uses multiple techniques and codebases I have developed over the year. Its a natural mixture of [min-max-IN-dit](https://github.com/cloneofsimo/min-max-in-dit), [min-max-gpt](https://github.com/cloneofsimo/min-max-gpt), [ez-muP](https://github.com/cloneofsimo/ezmup)
# Citations
If you use this material, please cite this repository with the following:
```bibtex
@misc{ryu2024minrf,
author = {Simo Ryu},
title = {minRF: Minimal Implementation of Scalable Rectified Flow Transformers},
year = 2024,
publisher = {Github},
url = {https://github.com/cloneofsimo/minRF},
}
```