pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation
Paper • 2510.14974 • Published • 10
import torch
from diffusers import DiffusionPipeline
# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("Lakonik/pi-Flow-ImageNet", dtype=torch.bfloat16, device_map="cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]Distilled 1-step and 2-step ImageNet DiTs proposed in the paper:
pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation
Hansheng Chen1,
Kai Zhang2,
Hao Tan2,
Leonidas Guibas1,
Gordon Wetzstein1,
Sai Bi2
1Stanford University, 2Adobe Research
[arXiv] [Code] [pi-Qwen Demo🤗] [pi-FLUX Demo🤗]
@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},
}