This is the official implementation of "DMPlug: A Plug-in Method for Solving Inverse Problems with Diffusion Models". You can find our paper via arXiv.
In this paper, we advocate viewing the reverse process in DMs as a function and propose a novel plug-in method for solving IPs using pretrained DMs, dubbed DMPlug. DMPlug addresses the issues of manifold feasibility and measurement feasibility in a principled manner, and also shows great potential for being robust to unknown types and levels of noise. Through extensive experiments across various IP tasks, including two linear and three nonlinear IPs, we demonstrate that DMPlug consistently outperforms state-of-the-art methods, often by large margins especially for nonlinear IPs.
git clone https://github.com/sun-umn/DMPlug.git
cd DMPlug
From the link, download the checkpoint "celebahq_p2.pt" and paste it to ./models/;
From the link, download the checkpoint "ffhq_10m.pt" and paste it to ./models/;
From the link, download the checkpoint "lsun_bedroom.pt" and paste it to ./models/.
mkdir models
mv {DOWNLOAD_DIR}/celebahq_p2.pt ./models/
mv {DOWNLOAD_DIR}/ffqh_10m.pt ./models/
mv {DOWNLOAD_DIR}/lsun_bedroom.pt ./models/
{DOWNLOAD_DIR} is the directory that you downloaded checkpoint to.
We use the external codes for motion-blurring and non-linear deblurring.
git clone https://github.com/VinAIResearch/blur-kernel-space-exploring bkse
git clone https://github.com/LeviBorodenko/motionblur motionblur
From the link, download the checkpoint "GOPRO_wVAE.pth" and paste it to ./experiments/pretrained/.
mv {DOWNLOAD_DIR}/GOPRO_wVAE.pt ./experiments/pretrained/
{DOWNLOAD_DIR} is the directory that you downloaded checkpoint to.
Install dependencies
conda env create -f environment.yml
conda activate dmplug
# Superresolution, inpainting, nonlinear deblurring
python sr_inp_nonlinear.py --task 'super_resolution'
python sr_inp_nonlinear.py --task 'inpainting'
python sr_inp_nonlinear.py --task 'nonlinear_deblur'
# BID
python bid.py --kernel 'motion'
python bid.py --kernel 'gaussian'
# BID with turbulence
python turbulence.py
This repo is developed based on DPS and BlindDPS, especially for forward operations. Please also consider citing them if you use this repo.
- Hengkang Wang, wang9881@umn.edu, Homepage
- Xu Zhang, spongezhang@gmail.com, Homepage
- Taihui Li, lixx5027@umn.edu, Homepage
- Tiancong Chen, chen6271@umn.edu, Homepage
- Yuxiang Wan, wan01530@umn.edu, Homepage
- Ju Sun, jusun@umn.edu, Homepage