D²-DPM: Dual Denoising for Quantized Diffusion Probabilistic Models

Overview

This is the official implementation of paper "D $^2$-DPM: Dual Denoising for Quantized Diffusion Probabilistic Models" [arXiv], which presents a dynamic quantization error correction strategy for diffusion models based on Gaussian modeling.

Getting Started

Requirements

Set up a virtual environment and install the required dependencies as specified in LDM’s instructions.

Sampling with the FP32 Model

python quant_scripts/sample_fp32_imagenet.py

Pre-trained models

Utilize the pre-trained models provided by LDM.

• For example, download the pre-trained LDM-4 model used in class-conditional generation experiments.

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 

Pipeline

1. Collect an ImageNet calibration dataset with evenly spaced time interval sampling.

python quant_scripts/collect_imagenet_calibration_set.py

2. Perform weight quantization using the BRECQ algorithm.

python quant_scripts/quantize_weight.py

• (Optional) Optionally adjust the distribution of timesteps to generate calibration samples.

  This requires:

  • Uncomment line 101 in quant_scripts/quantize_weight.py.
  • Comment out lines 96-98.

3. Perform activation quantization using the BRECQ algorithm.

python quant_scripts/quantize_weight_activation.py

• (Optional) Optionally adjust the distribution of timesteps to generate calibration samples, ensuring consistency with the selection in 2.

  This requires:

  • Uncomment line 101 in quant_scripts/quantize_weight.py.
  • Comment out lines 95-98.

4. Collect statistics for timestep-aware Gaussian modeling in preparation for the next step.

python quant_scripts/collect_timestep_aware_statistics.py

5. Perform timestep-aware Gaussian modeling, optionally choosing between tensor-wise or channel-wise modeling.

python quant_scripts/gaussian_modeling.py

6. Sample using the quantized model and corrected sampler, selecting from the following corrected samplers based on the provided descriptions:

• DDIMSampler_gaussian_quantCorrection_imagenet:

  Utilizes tensor-wise modeling and applies the S-D² denoising strategy.

• DDIMSampler_implicit_gaussian_quantCorrection_imagenet:

  Utilizes tensor-wise modeling and applies the D-D² denoising strategy.

• DDIMSampler_channel_wise_explicit_gaussian_quantCorrection_imagenet:

  Utilizes channel-wise modeling and applies the S-D² denoising strategy.

• DDIMSampler_channel_wise_implicit_gaussian_quantCorrection_imagenet:

  Utilizes channel-wise modeling and applies the D-D² denoising strategy.

python quant_scripts/sample_quantized_ldm_imagenet.py

Evaluation

1. Generate at least 50,000 samples. You may also choose whether to perform calibration dataset resampling.

python quant_scripts/generate_evaluation_samples_dist.py

2. Evaluate using OpenAI’s evaluator and the open-source reference batch available at guided-diffusion.

python evaluator.py --ref_batch VIRTUAL_imagenet256_labeled.npz --sample_batch your_synthesized_samples.npz

Acknowledgement

This repository is built upon LDM and PTQD. We thank the authors for their open-sourced code.