This repository provides the implementation for our paper "EFFICIENTDM: Efficient Quantization-Aware Fine-Tuning of Low-Bit Diffusion Models". Our approach introduces a novel method for fine-tuning low-bit diffusion models, focusing on efficiency and performance.

arXiv | BibTeX

Follow the step-by-step tutorial to set up EFFICIENTDM.

Create a virtual environment and install dependencies as specified by LDM.

Download the pretrained models provided by LDM.

mkdir -p models/ldm/cin256-v2/
wget -O models/ldm/cin256-v2/model.ckpt https://ommer-lab.com/files/latent-diffusion/nitro/cin/model.ckpt

Gather input data required for model calibration. Remember to modified the ldm/models/diffusion/ddpm.py as indicated in the quant_scripts/collect_input_4_calib.py.

python3 quant_scripts/collect_input_4_calib.py

We just apply a naive quantization method for model calibration because we will fine-tune it afterwards.

python3 quant_scripts/quantize_ldm_naive.py

Convert the fake quantized model to a packed integer model. Notice the significant reduction in model size.

python3 quant_scripts/save_naive_2_intmodel.py

python3 quant_scripts/train_efficientdm.py

Optionally, downsample the Temporal Activation LSQ parameters for sampling with fewer steps.

python3 quant_scripts/downsample_talsq_ckpt.py

python3 quant_scripts/sample_lora_intmodel.py

If you find this work useful for your research, please consider citing:

@inproceedings{he2024efficientdm,
  title={EfficientDM: Efficient Quantization-Aware Fine-Tuning of Low-Bit Diffusion Models},
  author={He, Yefei and Liu, Jing and Wu, Weijia and Zhou, Hong and Zhuang, Bohan},
  booktitle={International Conference on Learning Representations},
  year={2024}
}