LakonLab: Official Codebase for AsymFlow, pi-Flow, and GMFlow

June 1, 2026 · View on GitHub

Official PyTorch implementation of the papers:

🔥News

  • [May 20, 2026] AsymFlow is now supported in our ComfyUI extension.

  • [May 14, 2026] AsymFlow 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🤗. Please re-install the latest version of LakonLab (this repository) to use pi-FLUX.2.

  • [Nov 7, 2025] 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.

Installation

The code has been tested in the following environment:

  • Linux (tested on Ubuntu 20 and above)
  • PyTorch 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:

# Move to this repository (the folder with setup.py) after cloning
cd <PATH_TO_YOUR_LOCAL_REPO>

# 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 and FLUX.1 dev conditions, 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:

  • 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://<HF_REPO_NAME>/<PATH_TO_MODEL>) and HTTP/HTTPS URLs directly.

  • Streamlined training and evaluation: Supports online evaluation using common metrics, including FID, KID, IS, Precision, Recall, CLIP similarity, VQAScore, HPSv2, and HPSv3. Supports exporting results to offline evaluators, including HPSv3 Benchmark, DPG-Bench and GenEval.

  • 3rd-party model inference reproduction:

LakonLab uses the configuration system and code structure from 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},
}