noamelata/Blind-Inpainting-Diffusion

VOutliers Detection and Removal Using Diffusion Models Combined with Confidence Intervals

Reliability, Equity, and Reproducibility in Modern Machine Learning - 048100

Noam Elata: GitHub
Shahar Yadin: GitHub

Background

In this project we equip diffusion models with confidence intervals as done in the paper CONffusion: confidence intervals for diffusion models". We demonstrate that using the confidence intervals, we can solve the task of blind image inpainting, removing simple image artifacts.

Preparation

Download the CelebA-HQ256 dataset and the Diffusion Model checkpoint from SDEdit.

Create a conda environment using environment.yml.

Training

Finetune the pretrained model using the training script for the desired quantiles:

python train.py -c configs/celeba_hq_q005.json

python train.py -c configs/celeba_hq_q095.json

Calibrate the models using

Calibrate using the calibration set:

python calibrate.py -c configs/celeba_hq_q005.json \
                    -ml <lower quantile model checkpoint> -mh <higher quantile model checkpoint>

Inference & Visual Results

Run the following script:

python calibrate.py -c configs/celeba_hq_q005.json \
                    -ml <lower quantile model checkpoint> -mh <higher quantile model checkpoint> \
                    --artifact ["red", "rainbow", "butterflys"]