add dps

diffusion-posterior-sampling/Dockerfile

@@ -0,0 +1,28 @@
+FROM nvidia/cuda:11.3.1-devel-ubuntu20.04
+
+ENV TZ=Asiz/Seoul
+ENV TERM=xterm-256color
+
+RUN ln -fs /usr/share/zoneinfo/Asia/Seoul /etc/localtime
+
+#### 0. Install python and pip
+RUN apt-get -y update && apt-get install -y git wget curl
+RUN apt-get update
+RUN apt-get upgrade python3 -y
+RUN apt-get install python3-pip -y
+RUN alias python='python3'
+
+#### 1. Install Pytorch
+RUN pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113
+
+#### 2. Install other dependencies
+WORKDIR /usr/app
+COPY . ./
+RUN pip install -r ./requirements.txt
+
+#### 3. Clone external codes
+RUN git clone https://github.com/VinAIResearch/blur-kernel-space-exploring bkse
+RUN git clone https://github.com/LeviBorodenko/motionblur motionblur
+
+#### 4. change user
+RUN useradd docker_user -u 1000 -m

diffusion-posterior-sampling/README.md

@@ -0,0 +1,158 @@
+# Diffusion Posterior Sampling for General Noisy Inverse Problems (ICLR 2023 spotlight)
+
+![result-gif1](./figures/motion_blur.gif)
+![result-git2](./figures/super_resolution.gif)
+<!-- See more results in the [project-page](https://jeongsol-kim.github.io/dps-project-page) -->
+
+## Abstract
+In this work, we extend diffusion solvers to efficiently handle general noisy (non)linear inverse problems via the approximation of the posterior sampling. Interestingly, the resulting posterior sampling scheme is a blended version of the diffusion sampling with the manifold constrained gradient without strict measurement consistency projection step, yielding more desirable generative path in noisy settings compared to the previous studies.
+
+![cover-img](./figures/cover.jpg)
+
+
+## Prerequisites
+- python 3.8
+
+- pytorch 1.11.0
+
+- CUDA 11.3.1
+
+- nvidia-docker (if you use GPU in docker container)
+
+It is okay to use lower version of CUDA with proper pytorch version.
+
+Ex) CUDA 10.2 with pytorch 1.7.0
+
+<br />
+
+## Getting started 
+
+### 1) Clone the repository
+
+```
+git clone https://github.com/DPS2022/diffusion-posterior-sampling
+
+cd diffusion-posterior-sampling
+```
+
+<br />
+
+### 2) Download pretrained checkpoint
+From the [link](https://drive.google.com/drive/folders/1jElnRoFv7b31fG0v6pTSQkelbSX3xGZh?usp=sharing), download the checkpoint "ffhq_10m.pt" and paste it to ./models/
+```
+mkdir models
+mv {DOWNLOAD_DIR}/ffqh_10m.pt ./models/
+```
+{DOWNLOAD_DIR} is the directory that you downloaded checkpoint to.
+
+:speaker: Checkpoint for imagenet is uploaded.
+
+<br />
+
+
+### 3) Set environment
+### [Option 1] Local environment setting
+
+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
+```
+
+Install dependencies
+
+```
+conda create -n DPS python=3.8
+
+conda activate DPS
+
+pip install -r requirements.txt
+
+pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113
+```
+
+<br />
+
+### [Option 2] Build Docker image
+
+Install docker engine, GPU driver and proper cuda before running the following commands.
+
+Dockerfile already contains command to clone external codes. You don't have to clone them again.
+
+--gpus=all is required to use local GPU device (Docker >= 19.03)
+
+```
+docker build -t dps-docker:latest .
+
+docker run -it --rm --gpus=all dps-docker
+```
+
+<br />
+
+### 4) Inference
+
+```
+python3 sample_condition.py \
+--model_config=configs/model_config.yaml \
+--diffusion_config=configs/diffusion_config.yaml \
+--task_config={TASK-CONFIG};
+```
+
+
+:speaker: For imagenet, use configs/imagenet_model_config.yaml
+
+<br />
+
+## Possible task configurations
+
+```
+# Linear inverse problems
+- configs/super_resolution_config.yaml
+- configs/gaussian_deblur_config.yaml
+- configs/motion_deblur_config.yaml
+- configs/inpainting_config.yaml
+
+# Non-linear inverse problems
+- configs/nonlinear_deblur_config.yaml
+- configs/phase_retrieval_config.yaml
+```
+
+### Structure of task configurations
+You need to write your data directory at data.root. Default is ./data/samples which contains three sample images from FFHQ validation set.
+
+```
+conditioning:
+    method: # check candidates in guided_diffusion/condition_methods.py
+    params:
+        scale: 0.5
+
+data:
+    name: ffhq
+    root: ./data/samples/
+
+measurement:
+    operator:
+        name: # check candidates in guided_diffusion/measurements.py
+
+noise:
+    name:   # gaussian or poisson
+    sigma:  # if you use name: gaussian, set this.
+    (rate:) # if you use name: poisson, set this.
+```
+
+## Citation
+If you find our work interesting, please consider citing
+
+```
+@inproceedings{
+chung2023diffusion,
+title={Diffusion Posterior Sampling for General Noisy Inverse Problems},
+author={Hyungjin Chung and Jeongsol Kim and Michael Thompson Mccann and Marc Louis Klasky and Jong Chul Ye},
+booktitle={The Eleventh International Conference on Learning Representations },
+year={2023},
+url={https://openreview.net/forum?id=OnD9zGAGT0k}
+}
+```
+

diffusion-posterior-sampling/requirements.txt

@@ -0,0 +1,20 @@
+certifi==2022.9.14
+charset-normalizer==2.1.1
+contourpy==1.0.5
+cycler==0.11.0
+fonttools==4.37.2
+idna==3.4
+kiwisolver==1.4.4
+matplotlib==3.6.0
+numpy==1.23.3
+packaging==21.3
+Pillow==9.2.0
+pyparsing==3.0.9
+python-dateutil==2.8.2
+PyYAML==6.0
+requests==2.28.1
+scipy==1.9.1
+six==1.16.0
+tqdm==4.64.1
+typing-extensions==4.3.0
+urllib3==1.26.12

diffusion-posterior-sampling/sample_condition.py

@@ -0,0 +1,121 @@
+from functools import partial
+import os
+import argparse
+import yaml
+
+import torch
+import torchvision.transforms as transforms
+import matplotlib.pyplot as plt
+
+from guided_diffusion.condition_methods import get_conditioning_method
+from guided_diffusion.measurements import get_noise, get_operator
+from guided_diffusion.unet import create_model
+from guided_diffusion.gaussian_diffusion import create_sampler
+from data.dataloader import get_dataset, get_dataloader
+from util.img_utils import clear_color, mask_generator
+from util.logger import get_logger
+
+
+def load_yaml(file_path: str) -> dict:
+    with open(file_path) as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    return config
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_config', type=str)
+    parser.add_argument('--diffusion_config', type=str)
+    parser.add_argument('--task_config', type=str)
+    parser.add_argument('--gpu', type=int, default=0)
+    parser.add_argument('--save_dir', type=str, default='./results')
+    args = parser.parse_args()
+   
+    # logger
+    logger = get_logger()
+    
+    # Device setting
+    device_str = f"cuda:{args.gpu}" if torch.cuda.is_available() else 'cpu'
+    logger.info(f"Device set to {device_str}.")
+    device = torch.device(device_str)  
+    
+    # Load configurations
+    model_config = load_yaml(args.model_config)
+    diffusion_config = load_yaml(args.diffusion_config)
+    task_config = load_yaml(args.task_config)
+   
+    #assert model_config['learn_sigma'] == diffusion_config['learn_sigma'], \
+    #"learn_sigma must be the same for model and diffusion configuartion."
+    
+    # Load model
+    model = create_model(**model_config)
+    model = model.to(device)
+    model.eval()
+
+    # Prepare Operator and noise
+    measure_config = task_config['measurement']
+    operator = get_operator(device=device, **measure_config['operator'])
+    noiser = get_noise(**measure_config['noise'])
+    logger.info(f"Operation: {measure_config['operator']['name']} / Noise: {measure_config['noise']['name']}")
+
+    # Prepare conditioning method
+    cond_config = task_config['conditioning']
+    cond_method = get_conditioning_method(cond_config['method'], operator, noiser, **cond_config['params'])
+    measurement_cond_fn = cond_method.conditioning
+    logger.info(f"Conditioning method : {task_config['conditioning']['method']}")
+   
+    # Load diffusion sampler
+    sampler = create_sampler(**diffusion_config) 
+    sample_fn = partial(sampler.p_sample_loop, model=model, measurement_cond_fn=measurement_cond_fn)
+   
+    # Working directory
+    out_path = os.path.join(args.save_dir, measure_config['operator']['name'])
+    os.makedirs(out_path, exist_ok=True)
+    for img_dir in ['input', 'recon', 'progress', 'label']:
+        os.makedirs(os.path.join(out_path, img_dir), exist_ok=True)
+
+    # Prepare dataloader
+    data_config = task_config['data']
+    transform = transforms.Compose([transforms.ToTensor(),
+                                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+    dataset = get_dataset(**data_config, transforms=transform)
+    loader = get_dataloader(dataset, batch_size=1, num_workers=0, train=False)
+
+    # Exception) In case of inpainting, we need to generate a mask 
+    if measure_config['operator']['name'] == 'inpainting':
+        mask_gen = mask_generator(
+           **measure_config['mask_opt']
+        )
+        
+    # Do Inference
+    for i, ref_img in enumerate(loader):
+        logger.info(f"Inference for image {i}")
+        fname = str(i).zfill(5) + '.png'
+        ref_img = ref_img.to(device)
+
+        # Exception) In case of inpainging,
+        if measure_config['operator'] ['name'] == 'inpainting':
+            mask = mask_gen(ref_img)
+            mask = mask[:, 0, :, :].unsqueeze(dim=0)
+            measurement_cond_fn = partial(cond_method.conditioning, mask=mask)
+            sample_fn = partial(sample_fn, measurement_cond_fn=measurement_cond_fn)
+
+            # Forward measurement model (Ax + n)
+            y = operator.forward(ref_img, mask=mask)
+            y_n = noiser(y)
+
+        else: 
+            # Forward measurement model (Ax + n)
+            y = operator.forward(ref_img)
+            y_n = noiser(y)
+         
+        # Sampling
+        x_start = torch.randn(ref_img.shape, device=device).requires_grad_()
+        sample = sample_fn(x_start=x_start, measurement=y_n, record=True, save_root=out_path)
+
+        plt.imsave(os.path.join(out_path, 'input', fname), clear_color(y_n))
+        plt.imsave(os.path.join(out_path, 'label', fname), clear_color(ref_img))
+        plt.imsave(os.path.join(out_path, 'recon', fname), clear_color(sample))
+
+if __name__ == '__main__':
+    main()