piFlow/README.md

pi-Flow: Policy-Based Flow Models

Official PyTorch implementation of the paper:

pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation
Hansheng Chen¹, Kai Zhang², Hao Tan², Leonidas Guibas¹, Gordon Wetzstein¹, Sai Bi^2
1Stanford University, ²Adobe Research
arXiv | ComfyUI | pi-Qwen Demo🤗 | pi-FLUX Demo🤗

🔥News

• [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.

Highlights

• Novel Framework: pi-Flow stands for policy-based flow models. The network does not output a denoised state; instead, it outputs a fast policy that rolls out multiple ODE substeps to reach the denoised state.

  

• Simple Distillation: pi-Flow adopts policy-based imitation distillation (pi-ID). No JVPs, no auxiliary networks, no GANs—just a single L2 loss between the policy and the teacher.

  

• Diversity and Teacher Alignment: pi-Flow mitigates the quality–diversity trade-off, generating highly diverse samples while maintaining high quality. It also remains highly faithful to the teacher’s style. The example below shows that pi-Flow samples generally align with the teacher’s outputs and exhibit significantly higher diversity than those from DMD students (e.g., SenseFlow, Qwen-Image Lightning).

  

• Texture Details: pi-Flow excels in generating fine-grained texture details. When using additional photorealistic style LoRAs, this advantage becomes very prominent, as shown in the comparison below (zoom in for best view).

  

• Scalability: pi-Flow scales from ImageNet DiT to 20-billion-parameter text-to-image models (Qwen-Image). This codebase is highly optimized for large-scale experiments. See the Codebase section for details.

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 . from the repository root to install the LakonLab codebase and its dependencies.

An example of installation commands is shown below:

# Create conda environment
conda create -y -n piflow python=3.10 ninja
conda activate piflow

# Install Pytorch. Goto https://pytorch.org/get-started/previous-versions/ to select the appropriate version
pip install torch==2.6.0 torchvision==0.21.0

# Move to this repository (the folder with setup.py) after cloning
cd <PATH_TO_YOUR_LOCAL_REPO>
# Install LakonLab in editable mode
pip install -e .

Additional notes:

• To access FLUX models, please accept the conditions here, and then run huggingface-cli login to login with your HuggingFace account.
• Optionally, if you would like to use AWS S3 for dataset and checkpoint storage, please also install the AWS CLI.
• This codebase may work on Windows systems, but it has not been tested extensively.

Inference: Diffusers Pipelines

We provide diffusers pipelines for easy inference. The following code demonstrates how to sample images from the distilled Qwen-Image and FLUX models.

4-NFE GM-Qwen (GMFlow Policy)

Note: GM-Qwen supports elastic inference. Feel free to set num_inference_steps to any value above 4.

import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from lakonlab.pipelines.piqwen_pipeline import PiQwenImagePipeline

pipe = PiQwenImagePipeline.from_pretrained(
    'Qwen/Qwen-Image',
    torch_dtype=torch.bfloat16)
adapter_name = pipe.load_piflow_adapter(  # you may later call `pipe.set_adapters([adapter_name, ...])` to combine other adapters (e.g., style LoRAs)
    'Lakonik/pi-Qwen-Image',
    subfolder='gmqwen_k8_piid_4step',
    target_module_name='transformer')
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(  # use fixed shift=3.2
    pipe.scheduler.config, shift=3.2, shift_terminal=None, use_dynamic_shifting=False)
pipe = pipe.to('cuda')

out = pipe(
    prompt='Photo of a coffee shop entrance featuring a chalkboard sign reading "π-Qwen Coffee 😊 $2 per cup," with a neon '
           'light beside it displaying "π-通义千问". Next to it hangs a poster showing a beautiful Chinese woman, '
           'and beneath the poster is written "e≈2.71828-18284-59045-23536-02874-71352".',
    width=1920,
    height=1080,
    num_inference_steps=4,
    generator=torch.Generator().manual_seed(42),
).images[0]
out.save('gmqwen_4nfe.png')

4-NFE GM-FLUX (GMFlow Policy)

Note: For the 8-NFE version, replace gmflux_k8_piid_4step with gmflux_k8_piid_8step and set num_inference_steps=8.

import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from lakonlab.pipelines.piflux_pipeline import PiFluxPipeline

pipe = PiFluxPipeline.from_pretrained(
    'black-forest-labs/FLUX.1-dev',
    torch_dtype=torch.bfloat16)
adapter_name = pipe.load_piflow_adapter(  # you may later call `pipe.set_adapters([adapter_name, ...])` to combine other adapters (e.g., style LoRAs)
    'Lakonik/pi-FLUX.1',
    subfolder='gmflux_k8_piid_4step',
    target_module_name='transformer')
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(  # use fixed shift=3.2
    pipe.scheduler.config, shift=3.2, use_dynamic_shifting=False)
pipe = pipe.to('cuda')

out = pipe(
    prompt='A portrait photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the Sydney Opera House holding a sign on the chest that says "Welcome Friends"',
    width=1360,
    height=768,
    num_inference_steps=4,
    generator=torch.Generator().manual_seed(42),
).images[0]
out.save('gmflux_4nfe.png')

4-NFE DX-Qwen and DX-FLUX (DX Policy)

See example_dxqwen_pipeline.py and example_dxflux_pipeline.py for examples of using the DX policy.

Inference: Gradio Apps

We provide Gradio apps for interactive inference with the distilled GM-Qwen and GM-FLUX models. Official apps are available on HuggingFace Spaces: pi-Qwen Demo🤗 and pi-FLUX Demo🤗.

Run the following commands to launch the apps locally:

python demo/gradio_gmqwen.py --share  # GM-Qwen elastic inference

python demo/gradio_gmflux.py --share  # GM-FLUX 4-NFE and 8-NFE inference

Toy Models

To aid understanding, we provide minimal toy model training scripts that overfit the teacher behavior on a fixed initial noise using a static GMFlow policy (without student network).

Run the following command to distill a toy model from a ImageNet DiT (REPA):

python demo/train_piflow_dit_imagenet_toymodel.py

Run the following command to distill a toy model from Qwen-Image (requires 40GB VRAM):

python demo/train_piflow_qwen_toymodel.py

The results of these toy models demonstrate the expressiveness of the GMFlow policy—a GMFlow policy with 32 components can fit the entire ODE trajectory from $t=1$ to $t=0$, making it theoretically possible for 1-NFE generation. In practice, the bottleneck is often the student network, not the policy itself, thus more NFEs are still needed.

Training and Evaluation

Follow the instructions in the following links to reproduce the main results in the paper:

• Distilling ImageNet DiT
• Distilling Qwen-Image
• Distilling FLUX

By default, checkpoints will be saved into checkpoints/, logs will be saved into work_dirs/, and sampled images will be saved into viz/. These directories can be changed by modifying the config file (AWS S3 URLs are supported). If existing checkpoints are found, training will automatically resume from the latest checkpoint. The training logs can be plotted using Tensorboard. Run the following command to start Tensorboard:

tensorboard --logdir work_dirs/

To use Wandb logging, please export your authentication key to the WANDB_API_KEY environment variable, and then enable Wandb logging by appending the following code to the hooks list in the log_config part of the config file:

dict(
    type='WandbLoggerHook',
    init_kwargs=dict(project='PiFlow'),  # init_kwargs are passed to wandb.init()
)

Essential Code

• Training
  • train_piflow_dit_imagenet_toymodel.py and train_piflow_qwen_toymodel.py: Toy model distillation scripts with self-contained training loops.
  • piflow.py: The forward_train method contains the full training loop.
• Inference
  • piqwen_pipeline.py and piflux_pipeline.py: Full sampling code in the style of Diffusers.
  • piflow.py: The forward_test method contains the same full sampling loop.
• Policies
  • gmflow.py: GMFlow policy.
  • dx.py: DX policy.
• Networks
  • gmflow and dxflow: Student networks with modified output layers to predict the flow policy.

Codebase

pi-Flow and GMFlow are powered by LakonLab, 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.
  • FlowAdapterScheduler: Adapts diffusion solvers for flow matching models. Supports UniPCMultistep, DPMSolverMultistep, DPMSolverSinglestep, DEISMultistep, SASolver, etc.
• 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, HPSv3, and more.
• Diffusers-friendly: Diffusers models can be wrapped and integrated into LakonLab. See examples in configs/misc and lakonlab/models/architecture/diffusers.

LakonLab uses the configuration system and code structure from MMCV.

Citation

@misc{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},
      year={2025},
      eprint={2510.14974},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2510.14974}, 
}

@inproceedings{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},
  booktitle={ICML},
  year={2025},
}