Mean Flow Distillation

Official implementation of "Mean Flow Distillation: Robust and Stable Distillation for Flow Matching Models (ICML 2026)".

This repository provides the training code for multi-GPU Mean Flow Distillation (MFD) on SANA 1.6B. The current release focuses on prompt-only text-to-image distillation with LoRA adapters built on top of the Hugging Face diffusers SANA pipeline.

📌 Paper link: https://arxiv.org/abs/2606.11155

• Mean Flow Distillation training for flow matching text-to-image models.
• SANA 1.6B 1024px BF16 backbone support.
• LoRA-based student and auxiliary adapter training.
• Multi-GPU training via accelerate.

• Add paper link.
• Add pretrained MFD LoRA checkpoints.

.
├── README.md
├── requirements.txt
└── t2i_SANA
    ├── inference_sana.py          # Example inference script with trained LoRA weights
    ├── train_sana_mfd.sh          # Default multi-GPU launch script
    └── scripts
        └── train_sana_mfd.py      # Main MFD training script

We recommend using a fresh Conda environment with Python 3.10.

conda create -n mfd python=3.10
conda activate mfd

Install the PyTorch version that matches your CUDA and GPU setup. Please refer to the official PyTorch installation page:

# Example only. Choose the command that matches your CUDA version.
pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu118

Install the latest diffusers from source:

git clone https://github.com/huggingface/diffusers.git
cd diffusers
pip install -e .
cd ..

Install the remaining dependencies:

pip install -r requirements.txt

The training script checks for diffusers>=0.36.0.dev0, so using the source installation above is recommended.

We use LAION-Aesthetic-6.5+ as the prompt source for distillation. Only the prompt text is used during training.

Download the cleaned prompt file from here and place it at:

t2i_SANA/datasets/labels_cleaned.tsv

The TSV file is expected to contain at least four tab-separated fields per line. The script uses the second field as the prompt.

Minimum:

• 1 GPU with at least 24 GB VRAM

Recommended:

• 8 GPUs with at least 40 GB VRAM each

cd t2i_SANA
bash train_sana_mfd.sh

The default script uses:

• Base model: Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers
• Resolution: 1024
• LoRA rank / alpha: 64 / 64
• Batch size per GPU: 1
• Mixed precision: bf16
• Number of GPUs: 8
• Output directory: t2i_SANA/outputs

Edit t2i_SANA/train_sana_mfd.sh to change GPU IDs, learning rates, output directory, validation prompt, training length, etc.

Common arguments:

Training writes files under --output_dir, including:

• checkpoint-* directories with saved LoRA weights and accelerator state.
• Intermediate decoded samples under pics/.

After training, update the LoRA checkpoint path in t2i_SANA/inference_sana.py:

pipe.load_lora_weights("outputs/checkpoint-1000/pytorch_lora_weights.safetensors")

Then run:

cd t2i_SANA
python inference_sana.py

The script loads the SANA 1.6B BF16 pipeline, applies the trained LoRA weights, generates one 1024x1024 image with 4 inference steps, and saves it to:

t2i_SANA/sana.png

If you find this repository useful, please consider citing our work.

@misc{zhao2026MeanFlowDistillation,
      title={Mean Flow Distillation: Robust and Stable Distillation for Flow Matching Models}, 
      author={An Zhao and Shengyuan Zhang and Zhongjian Sun and Yixiang Zhou and Zejian Li and Ling Yang and Tianrun Chen and Lingyun Sun},
      year={2026},
      eprint={2606.11155},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2606.11155}, 
}

This repository builds on and modifies code from the excellent open-source projects below:

• Hugging Face diffusers

We thank the authors and contributors for their valuable work.