LakonLab

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

Official PyTorch implementation of the papers:

• Asymmetric Flow Models [README] In NeurIPS 2026 Hansheng Chen, Jan Ackermann, Minseo Kim, Gordon Wetzstein, Leonidas Guibas Stanford University Project Page | arXiv | ComfyUI | AsymFLUX.2 klein Demo🤗

• pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation [README] In ICLR 2026 Hansheng Chen1, Kai Zhang2, Hao Tan2, Leonidas Guibas1, Gordon Wetzstein1, Sai Bi2 1Stanford University, 2Adobe Research arXiv | ComfyUI | pi-Qwen Demo🤗 | pi-FLUX Demo🤗 | pi-FLUX.2 Demo🤗

• Gaussian Mixture Flow Matching Models [README] In ICML 2025 Hansheng Chen1, Kai Zhang2, Hao Tan2, Zexiang Xu3, Fujun Luan2, Leonidas Guibas1, Gordon Wetzstein1, Sai Bi2 1Stanford University, 2Adobe Research, 3Hillbot arXiv

🔥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:

  • FlowSDEScheduler: Generic flow SDE solver with an adjustable diffusion coefficient. 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. Powers the GM-SDE solver in GMFlow.
  • FlowMapSDEScheduler: 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://<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:

  • ImageNet 256x256 models with ADM evaluation:

    Model	FID (reproduced)	FID (official)
    SiT-XL/2	2.05	2.06
    JiT-H/16	1.90	-
    DiT-XL RAE (unguided)	1.50	1.51
    REPA-XL/2	1.38	1.42
    REPA-E-XL VAVAE	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 and lakonlab/models/architectures/diffusers.

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