## Codebase for paper 8702: Unleashing the Denoising Capability of Diffusion Prior for Solving Inverse Problems.

The experimental codes are arranged into four folders: `inverse_celeba`, `inverse_ffhq`, `inverse_imagenet`, and  `msdm`.

The `msdm` is for the source separation and partial generation task, which is constructed based on [Multi-Source-Diffusion-Models](https://gladia-research-group.github.io/multi-source-diffusion-models/). `inverse_celeba`,  `inverse_ffhq`,  `inverse_imagenet` are for the image restoration tasks, which are constructed based on [DDRM](https://github.com/bahjat-kawar/ddrm).

### Source separation and partial generation

#### **Setup**

The environment is written in file `env.yaml` .

#### **Run** 

To reproduce the source separation and partial generation tasks,  please run scripts shown below:

```shell
cd ./msdm
# separation tasks - MSDM
python3 eval_separation_msdm.py
# separation tasks - ISDM
python3 eval_separation_isdm.py
# separation tasks - ProjDiff (Ours)
python3 projdiff_partial_generation.py
```

**Note**: 

+ The datasets and pre-trained models are from [Multi-Source-Diffusion-Models](https://gladia-research-group.github.io/multi-source-diffusion-models/). 

### Image restoration

#### Setup

The environment is written in file `environment.yaml` .

#### Run

The codes are arranged into three folders coresponding to the dataset, named `inverse_celeba`, `inverse_imagenet`, `inverse_ffhq`. Each experiment are separated into one folder, named by the settings. 

To reproduce the image restoration tasks, please run scripts shown below:

```shell
# move to the directory of a specific dataset (celeba, imagenet, ffhq)
cd ./ProjDiff_code_submit/inverse_{dataset_name} 
# run python scripts
python3 projdiff_{task-name}.py \
--ni \
--config {CONFIG_FILE} \
--doc {DATASET} \
--timesteps {STEPS} \
--deg {DEGRADATION} \
--sigma_0 {SIGMA_0} \
-i {IMAGE_FOLDER}
```

**Note**: 

+ The datasets and pre-trained models for imagenet are from [DDRM](https://github.com/bahjat-kawar/ddrm). 
+ The datasets and pre-trained models for CelebA are from [DDNM](https://github.com/wyhuai/DDNM)
+ The datasets and pre-trained models for FFHQ are from [Diffusion Posterior Sampling for General Noisy Inverse Problems (ICLR 2023 spotlight)](https://github.com/DPS2022/diffusion-posterior-sampling) 