Lakonik/pi-FLUX.2

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
Hansheng Chen¹, Kai Zhang², Hao Tan², Leonidas Guibas¹, Gordon Wetzstein¹, Sai Bi^2
1Stanford University, ²Adobe Research
[arXiv] [Code] [pi-Qwen Demo🤗] [pi-FLUX Demo🤗] [pi-FLUX.2 Demo🤗]

Usage

Please first install the official code repository.

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.

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