AutoDIR: Automatic All-in-One Image Restoration with Latent Diffusion

Yitong Jiang, Zhaoyang Zhang, Tianfan Xue, Jinwei Gu

The official PyTorch implementation of the paper [AutoDIR: Automatic All-in-One Image Restoration with Latent Diffusion]

News: Accepted by ECCV2024

| Github | Page | Arxiv |

We propose an all-in-one image restoration system with latent diffusion, named AutoDIR, which can automatically detect and restore images with multiple unknown degradations. Our main hypothesis is that many image restoration tasks, such as super-resolution, motion deblur, denoising, low-light enhancement, dehazing, and deraining can often be decomposed into some common basis operators which improve image quality in different directions. AutoDIR aims to learn one unified image restoration model capable of performing these basis operators by joint training with multiple image restoration tasks. Specifically, AutoDIR consists of a Blind Image Quality Assessment (BIQA) module based on CLIP which automatically detects unknown image degradations for input images, an All-in-One Image Restoration (AIR) module based on latent diffusion which handles multiple types of image degradation, and a Structural-Correction Module (SCM) which further recovers the image structures. Extensive experimental evaluation demonstrates that AutoDIR outperforms state-of-the-art approaches for a wider range of image restoration tasks. The design of AutoDIR also enables flexible user control (via text prompt) and generalization to new tasks as a foundation model of image restoration.

Demo Video

Set up a conda environment

conda env create -f environment.yaml
conda activate autodir
cd NAFNet
python setup.py develop --no_cuda_ext
pip install git+https://github.com/openai/consistencydecoder.git

Updates

• 2023.01.19: Add Colab demo of AutoDIR.
• 2024.01.19: Inference code and pre-trained models are released.

TODO

• With Latent-Consistency-model
• Code for training from scratch
• Code for finetuning to new image restoration tasks
• CodeLab release
• Pre-trained model release
• Inference code release

Sections.

1. Inference

In this section, we provide the script to evaluate automatically handling images with unknown degradations and use-customization image restoration.

easy start

Handling images with unknown degradations iteratively

python eval_w_SCM.py
--input
test_images\unknown.png
--steps
200
--output
\path\to\our\output_folder
--ckpt
\path\to\the\pretrained\checkpoints
--cfg-text
1
--config
"configs/generate.yaml"

Example of open-vocabulary user control

Example of multiple unknown artifacts

Example of low-resolution artifact

Example of blur artifact

Example of noise artifact

Example of raindrop artifact

Example of rain artifact

Example of haze artifact

Example of underexposure artifact

User-customized image restoration

python eval_w_SCM.py
--customize
"A photo needs underexposure artifact reduction"
--input
test_images\noise.png
--steps
200
--output
\path\to\our\output_folder
--ckpt
\path\to\the\pretrained\checkpoints
--cfg-text
1
--config
"configs/generate.yaml"

Example of user-customized image restoration

Result structure:

results/
    └── imageName_input.png   #input
    └── imageName_result.png   #result of All-in-one image restoration(Air) Module
    └── imageName_result_w_SCM.png   #result with Structure-correction-module(SCM)
    └── imageName_result_colorcorrect.png   #Air Module result after color correction (if appliable)
    └── imageName_result_w_SCM_colorcorrect.png   #Result after SCM after color correction (if appliable)

Comments

This implementation is based on StableDiffusion, NAFNet, InstructPix2Pix, stablesr

Citation

If you find our paper useful for your research, please consider citing our work 😊 :

@article{jiang2023autodir,
  title={Autodir: Automatic all-in-one image restoration with latent diffusion},
  author={Jiang, Yitong and Zhang, Zhaoyang and Xue, Tianfan and Gu, Jinwei},
  journal={arXiv preprint arXiv:2310.10123},
  year={2023}
}