Code for "Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors"

1) Install packages

conda create -n pnpdm python=3.10
conda activate pnpdm
pip install -r requirements.txt

2) Download pretrained checkpoint

Download the corresponding checkpoint from the links below and move it to ./models/.

• FFHQ (color)
• FFHQ (grayscale)
• Blackhole (grayscale)

4) Modify the dataset and model paths in config files

You need to modify the paths in the following files so that the dataset and models can be loaded properly:

• ./configs/data/ffhq.yaml
• ./configs/data/ffhq_grayscale.yaml
• ./configs/model/edm_unet_adm_blackhole.yaml
• ./configs/model/edm_unet_adm_dps_ffhq.yaml
• ./configs/model/edm_unet_adm_gray_ffhq.yaml

4) Run experiments

All the commands for running our experiments are provided in commands.sh. The experiments are configured using hydra. Please see its documentation for detailed usage.

5) Citation

Thank you for your interest in our work! Please consider citing

@misc{wu2024principled,
      title={Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors}, 
      author={Zihui Wu and Yu Sun and Yifan Chen and Bingliang Zhang and Yisong Yue and Katherine L. Bouman},
      year={2024},
      eprint={2405.18782},
      archivePrefix={arXiv},
      primaryClass={eess.IV},
      url={https://arxiv.org/abs/2405.18782}, 
}

Please email [email protected] if you run into any problems.