ChangyuanWang/IBFlow-Qwen-2NFE

IB-Flow

IB-Flow: Information Bottleneck-Guided CFG Distillation for Few-Step Text-to-Image Generation

Yiting Wang¹, Jingyi Zhang², Wenhu Zhang³, Ke Chao⁴, Yves Liang¹, Kun Cheng², Kang Zhao²

¹Tsinghua University; ²Wan Team, Alibaba Group; ³HKUST; ⁴Beijing Normal University

Overview

IB-Flow enables high-quality text-to-image generation with only 2 NFEs.

Few-step distillation greatly reduces the inference cost of diffusion models, but fixed guidance does not account for how generation changes from global structure formation to detail refinement. IB-Flow formulates CFG distillation from an information-bottleneck perspective, combining instance-aware supervision targets with an entropy-aware guidance schedule. This repository provides the Qwen-Image inference, prompt preprocessing, and distributed training code.

News

• [2026-07-14] The pretrained IBFlow-Qwen-2NFE adapter is released.
• [2026-07-13] The paper is available on arXiv.
• [2026-07-12] 🔥 The Qwen-Image training and inference code for IB-Flow is released.

Quickstart

Environment Setup

conda create -y -n ibflow python=3.10 ninja
conda activate ibflow

pip install torch==2.6.0 torchvision==0.21.0

pip install -r requirements.txt --no-build-isolation

Inference

IBFlow-Qwen uses Qwen-Image as its base model. The released adapter is designed for 2-NFE inference with timestep_ratio=1.0; both are the defaults used by the command below. IBFLOW_MODEL may be a Hugging Face model ID or a local adapter directory.

export QWEN_IMAGE_MODEL=Qwen/Qwen-Image
export IBFLOW_MODEL=ChangyuanWang/IBFlow-Qwen-2NFE

python inference_qwen.py \
  --base-model "${QWEN_IMAGE_MODEL}" \
  --adapter "${IBFLOW_MODEL}" \
  --prompt "A cinematic close-up portrait of an elderly Tibetan artisan in a sunlit workshop, weathered skin, fine wrinkles, warm brown eyes, both hands carefully carving a small wooden bird, realistic anatomy, shallow depth of field, 85mm photography." \
  --output outputs/ibflow_qwen.png

Data Preparation

Training uses prompts from Lakonik/t2i-prompts-3m. Since loading the text encoder together with the teacher and student requires additional GPU memory, we precompute the prompt embeddings on one node with 8 GPUs:

torchrun --nnodes=1 --nproc_per_node=8 \
  cache_image_prompt_data.py configs/qwen/ibflow_qwen_2nfe_k16.py \
  --text-encoder configs/qwen/_text_encoder.py \
  --max-size 2304128 \
  --launcher pytorch \
  --diff_seed

The cache is saved to data/t2i_prompts_3m/preproc_qwen/ by default and requires approximately 380GB of storage. Change data_root in configs/qwen/_data_trainval.py to use another location.

Training

Run the IB-Flow configuration on one node with 8 GPUs. The default configuration uses four samples per GPU and requires approximately 70GB of memory per GPU, so 8 A100-80GB GPUs are recommended.

torchrun --nnodes=1 --nproc_per_node=8 \
  train.py configs/qwen/ibflow_qwen_2nfe_k16.py \
  --launcher pytorch \
  --diff_seed

Checkpoints are saved to checkpoints/<experiment_name>/iter_*.pth, while training logs and TensorBoard events are written to work_dirs/<experiment_name>/. The experiment name is defined by name in the training config.

Visualize the training losses with TensorBoard:

tensorboard --logdir work_dirs/

Exporting a Trained Model

Training checkpoints are .pth files. Export the EMA adapter to the safetensors format used by the inference pipeline:

export IBFLOW_CHECKPOINT=/path/to/iter_5000.pth

python export_ibflow_to_diffusers.py \
  configs/qwen/ibflow_qwen_2nfe_k16.py \
  --ckpt "${IBFLOW_CHECKPOINT}" \
  --out-dir outputs/ibflow_qwen_adapter

Use the exported adapter for inference:

python inference_qwen.py \
  --base-model Qwen/Qwen-Image \
  --adapter outputs/ibflow_qwen_adapter \
  --prompt "A cinematic close-up portrait of an elderly Tibetan artisan in a sunlit workshop, weathered skin, fine wrinkles, warm brown eyes, both hands carefully carving a small wooden bird, realistic anatomy, shallow depth of field, 85mm photography." \
  --output outputs/ibflow_qwen_custom.png

Acknowledgments

This project benefits from the open-source work of Qwen-Image, ArcFlow, pi-Flow, Diffusers, and MMGeneration. The training prompts are provided by Lakonik/t2i-prompts-3m. We thank their authors and maintainers for making these resources publicly available.

Citation

If you find IB-Flow useful in your research, please cite:

@misc{wang2026ibflow,
  title={IB-Flow: Information Bottleneck-Guided CFG Distillation for Few-Step Text-to-Image Generation},
  author={Wang, Yiting and Zhang, Jingyi and Zhang, Wenhu and Chao, Ke and Liang, Yves and Cheng, Kun and Zhao, Kang},
  year={2026},
  eprint={2607.09133},
  archivePrefix={arXiv},
  primaryClass={cs.CV}
}