https://github.com/Lakonik/LakonLab
Official implementation of AsymFlow, pi-Flow, GMFlow
https://github.com/Lakonik/LakonLab
diffusion-models distillation flow-matching generative-ai generative-model image-generation pixel pytorch text-to-image
Last synced: 5 days ago
JSON representation
Official implementation of AsymFlow, pi-Flow, GMFlow
- Host: GitHub
- URL: https://github.com/Lakonik/LakonLab
- Owner: Lakonik
- License: apache-2.0
- Created: 2025-10-16T12:23:11.000Z (8 months ago)
- Default Branch: main
- Last Pushed: 2026-05-29T21:22:40.000Z (26 days ago)
- Last Synced: 2026-05-29T23:11:15.775Z (26 days ago)
- Topics: diffusion-models, distillation, flow-matching, generative-ai, generative-model, image-generation, pixel, pytorch, text-to-image
- Language: Python
- Homepage: https://hanshengchen.com/research#pub_diffuion
- Size: 19.9 MB
- Stars: 417
- Watchers: 5
- Forks: 25
- Open Issues: 8
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- License: LICENSE
Awesome Lists containing this project
- awesomeopd - piFlow - the-badge&logo=github&logoColor=white&labelColor=181717&color=ffd700" alt="Stars"> | 2025.10 | Multi-org | [arXiv 2510.14974](https://arxiv.org/abs/2510.14974) | π-Flow — image / flow OPD (ICLR 2026) | (🖼️ Multimodal OPD (VLM, Video, Audio, Image) / 🔁 Iterative Self-Bootstrapping)
README
# LakonLab: Official Codebase for AsymFlow, pi-Flow, and GMFlow
Official PyTorch implementation of the papers:
- **Asymmetric Flow Models [[README](docs/AsymFlow.md)]**
arXiv 2026
[Hansheng Chen](https://lakonik.github.io/),
[Jan Ackermann](https://janackermann.info/),
[Minseo Kim](https://soniaminseokim.github.io/),
[Gordon Wetzstein](http://web.stanford.edu/~gordonwz/),
[Leonidas Guibas](https://geometry.stanford.edu/?member=guibas)
Stanford University
[Project Page](https://hanshengchen.com/asymflow) | [arXiv](https://arxiv.org/abs/2605.12964) | [ComfyUI](https://github.com/Lakonik/ComfyUI-piFlow) | [AsymFLUX.2 klein Demo🤗](https://huggingface.co/spaces/Lakonik/AsymFLUX.2-klein)
- **pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation [[README](docs/piFlow.md)]**
In ICLR 2026
[Hansheng Chen](https://lakonik.github.io/)1,
[Kai Zhang](https://kai-46.github.io/website/)2,
[Hao Tan](https://research.adobe.com/person/hao-tan/)2,
[Leonidas Guibas](https://geometry.stanford.edu/?member=guibas)1,
[Gordon Wetzstein](http://web.stanford.edu/~gordonwz/)1,
[Sai Bi](https://sai-bi.github.io/)2
1Stanford University, 2Adobe Research
[arXiv](https://arxiv.org/abs/2510.14974) | [ComfyUI](https://github.com/Lakonik/ComfyUI-piFlow) | [pi-Qwen Demo🤗](https://huggingface.co/spaces/Lakonik/pi-Qwen) | [pi-FLUX Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.1) | [pi-FLUX.2 Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.2)
- **Gaussian Mixture Flow Matching Models [[README](docs/GMFlow.md)]**
In ICML 2025
[Hansheng Chen](https://lakonik.github.io/)1,
[Kai Zhang](https://kai-46.github.io/website/)2,
[Hao Tan](https://research.adobe.com/person/hao-tan/)2,
[Zexiang Xu](https://zexiangxu.github.io/)3,
[Fujun Luan](https://research.adobe.com/person/fujun/)2,
[Leonidas Guibas](https://geometry.stanford.edu/?member=guibas)1,
[Gordon Wetzstein](http://web.stanford.edu/~gordonwz/)1,
[Sai Bi](https://sai-bi.github.io/)2
1Stanford University, 2Adobe Research, 3Hillbot
[arXiv](https://arxiv.org/abs/2504.05304)
## 🔥News
- [May 20, 2026] AsymFlow is now supported in our [ComfyUI extension](https://github.com/Lakonik/ComfyUI-piFlow).
- [May 14, 2026] [AsymFlow](docs/AsymFlow.md) is released.
- [Dec 12, 2025] pi-FLUX.2 is now available for 4-step image generation and editing. Check out the [pi-FLUX.2 Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.2). Please re-install the latest version of LakonLab (this repository) to use pi-FLUX.2.
- [Nov 7, 2025] [ComfyUI-piFlow](https://github.com/Lakonik/ComfyUI-piFlow) is now available. Supports 4-step sampling of Qwen-Image and Flux.1 dev using 8-bit models on a single consumer-grade GPU, powered by [ComfyUI](https://github.com/comfyanonymous/ComfyUI).
## Installation
The code has been tested in the following environment:
- Linux (tested on Ubuntu 20 and above)
- [PyTorch](https://pytorch.org/get-started/previous-versions/) 2.6+
With the above prerequisites, run `pip install -e . --no-build-isolation` from the repository root to install the LakonLab codebase and its dependencies.
An example of installation commands is shown below:
```bash
# Move to this repository (the folder with setup.py) after cloning
cd
# Create uv environment
uv venv --python 3.10
source .venv/bin/activate
# Install Pytorch. Goto https://pytorch.org/get-started/previous-versions/ to select the appropriate version
uv pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu128
# Install LakonLab in editable mode
uv pip install -e . --no-build-isolation
```
Additional notes:
To access FLUX models, please accept the [FLUX.2 klein Base 9B conditions](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B) and [FLUX.1 dev conditions](https://huggingface.co/black-forest-labs/FLUX.1-dev), and then run `hf auth login` to login with your HuggingFace account.
## Codebase
**LakonLab** is a high-performance codebase for experimenting with large diffusion models. Key features of LakonLab include:
- **Performance optimizations**: Seamless switching between DDP, FSDP, and FSDP2, all supporting gradient accumulation and mixed precision.
- **Weight tying**: For LoRA fine-tuning, the base weights of the teacher, student, and EMA models are tied, sharing the same underlying memory. This is compatible with DDP and FSDP.
- **Advanced flow solvers**:
- [FlowSDEScheduler](lakonlab/models/diffusions/schedulers/flow_sde.py): Generic flow SDE solver with an adjustable [diffusion coefficient](https://arxiv.org/pdf/2306.02063). `h=0` corresponds to a flow ODE; `h=1` corresponds to a standard flow SDE; `h='inf'` corresponds to the re-noising sampler in the original [consistency models](https://arxiv.org/pdf/2303.01469). Powers the GM-SDE solver in GMFlow.
- [FlowMapSDEScheduler](lakonlab/models/diffusions/schedulers/flow_map_sde.py): Generic flow SDE solver for few-step flow map models, similar to above.
- **Storage backends**: Most I/O operations (e.g., dataloaders, checkpoint I/O) support both local filesystems and AWS S3. In addition, model checkpoints can be loaded from HuggingFace (link format `huggingface:///`) and HTTP/HTTPS URLs directly.
- **Streamlined training and evaluation**: Supports online evaluation using common [metrics](lakonlab/evaluation/metrics.py), including FID, KID, IS, Precision, Recall, CLIP similarity, [VQAScore](https://github.com/linzhiqiu/t2v_metrics), [HPSv2](https://github.com/tgxs002/HPSv2), and [HPSv3](https://github.com/MizzenAI/HPSv3). Supports exporting results to offline evaluators, including [HPSv3 Benchmark](https://github.com/MizzenAI/HPSv3), [DPG-Bench](https://github.com/TencentQQGYLab/ELLA) and [GenEval](https://github.com/djghosh13/geneval).
- **3rd-party model inference reproduction**:
- ImageNet 256x256 models with [ADM evaluation](https://github.com/openai/guided-diffusion/tree/main/evaluations):
| Model | FID (reproduced) | FID (official) |
|-------------------------------------------------------------------|:----------------:|:--------------:|
| [SiT-XL/2](configs/misc/sit_imagenet_test.py) | 2.05 | 2.06 |
| [JiT-H/16](configs/misc/jit_h_16_imagenet_test.py) | 1.90 | - |
| [DiT-XL RAE (unguided)](configs/misc/rae_dinov2_imagenet_test.py) | 1.50 | 1.51 |
| [REPA-XL/2](configs/misc/repa_imagenet_test.py) | 1.38 | 1.42 |
| [REPA-E-XL VAVAE](configs/misc/repae_imagenet_test.py) | 1.12 | 1.12 |
Note:
We use BF16 inference for all models except RAE.
Original JiT paper uses its own evaluation protocol that differs from ADM evaluation.
- Text-to-image models:
See examples in [configs/misc](configs/misc) and [lakonlab/models/architectures/diffusers](lakonlab/models/architectures/diffusers).
LakonLab uses the configuration system and code structure from [MMCV](https://github.com/open-mmlab/mmcv).
## Citation
```
@article{asymflow,
title={Asymmetric Flow Models},
author={Hansheng Chen and Jan Ackermann and Minseo Kim and Gordon Wetzstein and Leonidas Guibas},
url={https://arxiv.org/abs/2605.12964},
journal={arXiv preprint arXiv:2605.12964},
year={2026},
}
@article{piflow,
title={pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation},
author={Hansheng Chen and Kai Zhang and Hao Tan and Leonidas Guibas and Gordon Wetzstein and Sai Bi},
url={https://arxiv.org/abs/2510.14974},
journal={arXiv preprint arXiv:2510.14974},
year={2025},
}
@article{gmflow,
title={Gaussian Mixture Flow Matching Models},
author={Hansheng Chen and Kai Zhang and Hao Tan and Zexiang Xu and Fujun Luan and Leonidas Guibas and Gordon Wetzstein and Sai Bi},
url={https://arxiv.org/abs/2504.05304},
journal={arXiv preprint arXiv:2504.05304},
year={2025},
}
```