# pi-Flow: Policy-Based Flow Models Official PyTorch implementation of the paper: **pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation**
[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) ## šŸ”„News - [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). ## 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](https://github.com/XingtongGe/SenseFlow), [Qwen-Image Lightning](https://github.com/ModelTC/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](#codebase) section for details. ## 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 .` from the repository root to install the LakonLab codebase and its dependencies. An example of installation commands is shown below: ```bash # 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 # Install LakonLab in editable mode pip install -e . ``` Additional notes: - To access FLUX models, please accept the conditions [here](https://huggingface.co/black-forest-labs/FLUX.1-dev), 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](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html). - 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)](demo/example_gmqwen_pipeline.py) Note: GM-Qwen supports elastic inference. Feel free to set `num_inference_steps` to any value above 4. ```python 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)](demo/example_gmflux_pipeline.py) Note: For the 8-NFE version, replace `gmflux_k8_piid_4step` with `gmflux_k8_piid_8step` and set `num_inference_steps=8`. ```python 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](demo/example_dxqwen_pipeline.py) and [example_dxflux_pipeline.py](demo/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šŸ¤—](https://huggingface.co/spaces/Lakonik/pi-Qwen) and [pi-FLUX DemošŸ¤—](https://huggingface.co/spaces/Lakonik/pi-FLUX.1). Run the following commands to launch the apps locally: ```bash python demo/gradio_gmqwen.py --share # GM-Qwen elastic inference ``` ```bash 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): ```bash python demo/train_piflow_dit_imagenet_toymodel.py ``` Run the following command to distill a toy model from Qwen-Image (requires 40GB VRAM): ```bash 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](configs/piflow_imagenet/README.md) - [Distilling Qwen-Image](configs/piqwen/README.md) - [Distilling FLUX](configs/piflux/README.md) By default, checkpoints will be saved into [checkpoints/](checkpoints/), logs will be saved into [work_dirs/](work_dirs/), and sampled images will be saved into [viz/](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: ```bash 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: ```python dict( type='WandbLoggerHook', init_kwargs=dict(project='PiFlow'), # init_kwargs are passed to wandb.init() ) ``` ## Essential Code - Training - [train_piflow_dit_imagenet_toymodel.py](demo/train_piflow_dit_imagenet_toymodel.py) and [train_piflow_qwen_toymodel.py](demo/train_piflow_qwen_toymodel.py): Toy model distillation scripts with self-contained training loops. - [piflow.py](lakonlab/models/diffusions/piflow.py): The `forward_train` method contains the full training loop. - Inference - [piqwen_pipeline.py](lakonlab/pipelines/piqwen_pipeline.py) and [piflux_pipeline.py](lakonlab/pipelines/piflux_pipeline.py): Full sampling code in the style of Diffusers. - [piflow.py](lakonlab/models/diffusions/piflow.py): The `forward_test` method contains the same full sampling loop. - Policies - [gmflow.py](lakonlab/models/diffusions/piflow_policies/gmflow.py): GMFlow policy. - [dx.py](lakonlab/models/diffusions/piflow_policies/dx.py): DX policy. - Networks - [gmflow](lakonlab/models/architecture/gmflow) and [dxflow](lakonlab/models/architecture/dxflow): Student networks with modified output layers to predict the flow policy. ## Codebase pi-Flow and [GMFlow](configs/gmflow/README.md) 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](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. - [FlowAdapterScheduler](lakonlab/models/diffusions/schedulers/flow_adapter.py): 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:///`) 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), [HPSv3](https://github.com/MizzenAI/HPSv3), and more. - **Diffusers-friendly**: Diffusers models can be wrapped and integrated into LakonLab. See examples in [configs/misc](configs/misc) and [lakonlab/models/architecture/diffusers](lakonlab/models/architecture/diffusers). LakonLab uses the configuration system and code structure from [MMCV](https://github.com/open-mmlab/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}, } ```