File size: 3,919 Bytes
ac4d199 b1349e2 ac4d199 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | ---
base_model:
- black-forest-labs/FLUX.2-dev
datasets:
- Lakonik/t2i-prompts-3m
library_name: diffusers
license: other
license_name: flux-dev-non-commercial-license
license_link: LICENSE.md
pipeline_tag: text-to-image
tags:
- flux
- flow-matching
- distillation
---
# pi-Flow: Policy-Based Flow Models
4-step FLUX.2 models distilled from FLUX.2 dev, using the pi-Flow method proposed in the paper:
**pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation**
<br>
[Hansheng Chen](https://lakonik.github.io/)<sup>1</sup>,
[Kai Zhang](https://kai-46.github.io/website/)<sup>2</sup>,
[Hao Tan](https://research.adobe.com/person/hao-tan/)<sup>2</sup>,
[Leonidas Guibas](https://geometry.stanford.edu/?member=guibas)<sup>1</sup>,
[Gordon Wetzstein](http://web.stanford.edu/~gordonwz/)<sup>1</sup>,
[Sai Bi](https://sai-bi.github.io/)<sup>2</sup><br>
<sup>1</sup>Stanford University, <sup>2</sup>Adobe Research
<br>
[[arXiv](https://arxiv.org/abs/2510.14974)] [[Code](https://github.com/Lakonik/piFlow)] [[pi-Qwen Demo🤗](https://huggingface.co/spaces/Lakonik/pi-Qwen)] [[pi-FLUX Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.1)] [[pi-FLUX.2 Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.2)]
## Usage
Please first install the [official code repository](https://github.com/Lakonik/piFlow).
We provide diffusers pipelines for easy inference. The following code demonstrates how to sample images from the distilled FLUX.2 models.
### 4-NFE GM-FLUX.2 (GMFlow Policy)
Note: GM-FLUX.2 supports elastic inference. Feel free to set `num_inference_steps` to any value above 4.
```python
import torch
from lakonlab.models.diffusions.schedulers import FlowMapSDEScheduler
from lakonlab.pipelines.pipeline_piflux2 import PiFlux2Pipeline
from diffusers.utils import load_image
pipe = PiFlux2Pipeline.from_pretrained(
'diffusers/FLUX.2-dev-bnb-4bit',
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.2',
subfolder='gmflux2_k8_piid_4step',
target_module_name='transformer')
pipe.scheduler = FlowMapSDEScheduler.from_config( # use fixed shift=3.2
pipe.scheduler.config, shift=3.2, use_dynamic_shifting=False, final_step_size_scale=0.5)
pipe = pipe.to('cuda')
# Text-to-image generation example
prompt = "Realistic macro photograph of a hermit crab using a soda can as its shell, partially emerging from the can, captured with sharp detail and natural colors, on a sunlit beach with soft shadows and a shallow depth of field, with blurred ocean waves in the background. The can has the text `BFL Diffusers` on it and it has a color gradient that start with #FF5733 at the top and transitions to #33FF57 at the bottom."
out = pipe(
prompt=prompt,
width=1360,
height=768,
num_inference_steps=4,
generator=torch.Generator().manual_seed(42),
).images[0]
out.save('gmflux2_4nfe.png')
# Image editing example
prompt = "Add a hat on top of the cat."
cat_image = load_image("https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic/resolve/main/cat.png")
out = pipe(
prompt=prompt,
image=[cat_image], # optional multi-image input
width=1360,
height=768,
num_inference_steps=4,
generator=torch.Generator().manual_seed(42),
).images[0]
out.save('gmflux2_edit_4nfe.png')
```
## 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},
}
``` |