VETA-DiT: Variance-Equalized and Temporally Adaptive Quantization for Efficient 4-bit Diffusion Transformers.
This repo contains the official code of our paper: VETA-DiT: Variance-Equalized and Temporally Adaptive Quantization for Efficient 4-bit Diffusion Transformers.
Fitst, download and set up the repo.
Then create the environment and install required packages:
conda create -n veta-dit python=3.10
conda activate veta-dit
pip install -r requirements.txtUse the following command to generate the calibration data for VETA-DiT:
cd dit
bash get_calib_data.sh # You can change the save path in the script- Example for quantizing DiT-XL/2 with 100 timesteps into W4A8 on ImageNet 256x256 generation.
python ptq_inference.py \
--image-size 256\
--seed 42 \
--num-sampling-steps 100\
--ptq-config "./configs/w4a8.yaml"\ # You can change the quantization config file here
--log "./logs/w4a8" \ # You can change the log path here
--cfg-scale 1.5 \
--argument_method "inco" \
--smooth_quant \
--gptq- Example for quantizing DiT-XL/2 with 100 timesteps into W4A4 on ImageNet 256x256 generation.
python ptq_inference.py \
--image-size 256\
--seed 42 \
--num-sampling-steps 100\
--ptq-config "./configs/w4a4.yaml"\ # You can change the quantization config file here
--log "./logs/w4a4" \ # You can change the log path here
--cfg-scale 1.5 \
--argument_method "inco" \
--smooth_quant \
--gptqUse the following command to generate the calibration data for VETA-DiT:
cd pixart
bash get_calib_data_pixart.sh # You can change the save path in the script- Example for quantizing Pixart into W4A8 on COCO generation.
python ptq_inference.py --quant-config "./configs/w4a8.yaml" --log "./logs/w4a8" --argument_method "inco" --smooth_quant --gptq --prompt 'assets/coco_1024.txt'- Example for quantizing Pixart into W4A4 on COCO generation.
python ptq_inference.py --quant-config "./configs/w4a4.yaml" --log "./logs/w4a4" --argument_method "inco" --smooth_quant --gptq --prompt 'assets/coco_1024.txt'setup the environment following opensora1.2/OpenSORA/README.md
Use the following command to generate the calibration data for VETA-DiT:
cd dit
python get_calib_data.py configs/software_simulation_w4a8.py --prompt-path ./assets/t2v_samples_1.txt --seed 42 # You can change the save path in the script- Example for quantizing Open-Sora into W4A8.
python ptq_inference.py configs/software_simulation_w4a8.py --save-dir "./logs/w4a8" --prompt-path ./assets/t2v_samples.txt- Example for quantizing Open-Sora into W4A4.
python ptq_inference.py configs/software_simulation_w4a4.py --save-dir "./logs/w4a4" --prompt-path ./assets/t2v_samples.txtWe use the ADM’s evaluation suite to calculate FID, sFID, IS, and Precision.
Our code was developed based on opensora v1.2, PixArt-sigama and ViDiT-Q