We propose a new method called Decoupled Annealing Posterior Sampling (DAPS) that relies on a novel noise annealing process to solve posterior sampling with diffusion prior. Specifically, we decouple consecutive steps in a diffusion sampling trajectory, allowing them to vary considerably from one another while ensuring their time-marginals anneal to the true posterior as we reduce noise levels.
This approach enables the exploration of a larger solution space, improving the success rate for accurate reconstructions. We demonstrate that DAPS significantly improves sample quality and stability across multiple image restoration tasks, particularly in complicated nonlinear inverse problems.
-
2025-03: update code structure and usability to support different MCMC algorithms and various types of diffusion models. The previous code structure has been moved to the
legacybranch. Major updates are summarized as below:-
update diffusion schulers in
cores/scheduler.py. -
update MCMC sampler to support different algorithms and approximations in
cores/mcmc.py. -
enhance LatentDAPS with
$\texttt{HMC}$ which sustantially improve the performance.
-
| Link | Description |
|---|---|
 |
Try DAPS on demo datasets with different diffusion models. |
|
Customizing DAPS for New Inverse Problems |
- python 3.8
- PyTorch 2.3
- CUDA 12.1
Lower version of PyTorch with proper CUDA should work but not be fully tested.
# in DAPS folder
conda create -n DAPS python=3.8
conda activate DAPS
pip install -r requirements.txt
# (optional) install PyTorch with proper CUDA
conda install pytorch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 pytorch-cuda=12.1 -c pytorch -c nvidia
We use bkse for nonlinear blurring and motionblur for motion blur. No further action required then.
Download the public available FFHQ and ImageNet checkpoint (ffhq_10m.pt, imagenet256.pt) here.
# in DAPS folder
mkdir checkpoints
mv {DOWNLOAD_DIR}/ffqh_10m.pt checkpoints/ffhq256.pt
mv {DOWNLOAD_DIR}/imagenet256.pt.pt checkpoints/imagenet256.pt
Download the public available LDM checkpoint for FFHQ and ImageNet with following commands:
# in DAPS folder
wget https://ommer-lab.com/files/latent-diffusion/ffhq.zip -P ./checkpoints
unzip checkpoints/ffhq.zip -d ./checkpoints
mv checkpoints/model.ckpt checkpoints/ldm_ffhq256.pt
rm checkpoints/ffhq.zip
wget https://ommer-lab.com/files/latent-diffusion/nitro/cin/model.ckpt -P ./checkpoints/
mv checkpoints/model.ckpt checkpoints/ldm_imagenet256.pt
Checkpoints will be automatically downloaded.
(Optional) For nonlinear deblur task, we need the pretrained model from bkse at here:
# in DAPS folder
mv {DOWNLOAD_DIR}/GOPRO_wVAE.pth forward_operator/bkse/experiments/pretrained
You can download the selected test dataset used here, unzip and move to dataset folder. Otherwise, you can test on our provided 10 demo images at dataset\demo-ffhq and dataset\demo-imagenet.
Now you are ready to run. For phase retrieval with DAPS-1k and ffhq256ddpm model in 4 runs for 10 demo FFHQ images in dataset/demo-ffhq:
python posterior_sample.py \
+data=demo-ffhq \
+model=ffhq256ddpm \
+task=phase_retrieval \
+sampler=edm_daps \
task_group=pixel \
save_dir=results \
num_runs=4 \
sampler.diffusion_scheduler_config.num_steps=5 \
sampler.annealing_scheduler_config.num_steps=200 \
batch_size=10 \
data.start_id=0 data.end_id=10 \
name=phase_retrieval_demo \
gpu=0
It takes about 8 minutes (2 for each run) and 6G GPU memory on a single NVIDIA A100-SXM4-80GB GPU. The results are saved at foloder \results.
Full comands used to reproduce the results in paper are provided in commands folder:
- pixel space diffusion:
commands/pixel.sh - latent diffusion:
commands/ldm.sh - stable diffusion:
commands/sd.sh
| Model | Dataset | Model Config Name | Sampler | Task Group |
|---|---|---|---|---|
| ffhq-256 | FFHQ | ffhq256ddpm | edm_daps | pixel |
| imagenet-256 | ImageNet | imagenet256ddpm | edm_daps | pixel |
| ldm-ffhq-256 | FFHQ | ffhq256ldm | latent_edm_daps | ldm |
| ldm-imagenet-256 | ImageNet | imagenet256ldm | latent_edm_daps | ldm |
| sd-v1.5 | Any | stable-diffusion-v1.5 | sd_edm_daps | sd |
| sd-v2.1 | Any | stable-diffusion-v2.1 | sd_edm_daps | sd |
python posterior_sample.py \
+data={DATASET_CONFIG_NAME} \
+model={MODEL_CONFIG_NAME} \
+task={TASK_CONFIG_NAME} \
+sampler={SAMPLER_CONFIG_NAME} \
task_group={pixel, ldm, sd} # choose the used task parameters group \
save_dir=results \
num_runs={NUMBER_OF_RUNS} \
sampler.diffusion_scheduler_config.num_steps={DIFFUSION_ODE_STEPS} \
sampler.annealing_scheduler_config.num_steps={ANNEALING_STEPS} \
batch_size=100 \
name={SUB_FOLDER_NAME} \
gpu=0
Currently supported tasks are:
-
phase_retrieval: phase retrival of oversample ratio of 2.0 -
down_sampling: super resolution ($\times$4) -
inpainting: 128x128 box inpainting -
inpainting_rand: 70% random inpainting -
gaussian_blur: gaussian deblur of kernel size 61 and intensity 3 -
motion_blur: gaussian deblur of kernel size 61 and intensity 0.5 -
nonlinear_blur: nonlinear deblur of default setting in bkse repo -
hdr: high dynamic range reconstruction of factor 2
If you find our work interesting, please consider citing
@misc{zhang2024improvingdiffusioninverseproblem,
title={Improving Diffusion Inverse Problem Solving with Decoupled Noise Annealing},
author={Bingliang Zhang and Wenda Chu and Julius Berner and Chenlin Meng and Anima Anandkumar and Yang Song},
year={2024},
eprint={2407.01521},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2407.01521},
}