https://github.com/renyulin-f/MoE-DiffIR
The code source of MoE-DiffIR
https://github.com/renyulin-f/MoE-DiffIR
Last synced: 7 months ago
JSON representation
The code source of MoE-DiffIR
- Host: GitHub
- URL: https://github.com/renyulin-f/MoE-DiffIR
- Owner: renyulin-f
- Created: 2024-07-10T08:34:34.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2024-07-10T08:36:26.000Z (over 1 year ago)
- Last Synced: 2024-08-01T18:30:36.711Z (about 1 year ago)
- Size: 1000 Bytes
- Stars: 2
- Watchers: 2
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
- awesome-diffusion-categorized - [Code
README
[](https://arxiv.org/abs/2407.10833) [](https://renyulin-f.github.io/MoE-DiffIR.github.io/) 
This repository is the official PyTorch implementation of MoE-DiffIR (ECCV 2024).
## :bookmark: News!!!
- [x] 2024-07-15: **Arxiv version has been released.**
- [x] 2024-09-24: **Release partial of code, training and testing will come soon.**
- [x] 2024-12-27: **Release training and test code, Instructions will come soon.**
- [x] 2025-3-21: **Release visual results.**
> We present MoE-DiffIR, an innovative universal compressed image restoration (CIR) method with task-customized diffusion priors. This intends to handle two pivotal challenges in the existing CIR methods: (i) lacking adaptability and universality for different image codecs, e.g., JPEG and WebP; (ii) poor texture generation capability, particularly at low bitrates. Specifically, our MoE-DiffIR develops the powerful mixture-of-experts (MoE) prompt module, where some basic prompts cooperate to excavate the task-customized diffusion priors from Stable Diffusion (SD) for each compression task. Moreover, the degradation-aware routing mechanism is proposed to enable the flexible assignment of basic prompts. To activate and reuse the cross-modality generation prior of SD, we design the visual-to-text adapter for MoE-DiffIR, which aims to adapt the embedding of low-quality images from the visual domain to the textual domain as the textual guidance for SD, enabling more consistent and reasonable texture generation. We also construct one comprehensive benchmark dataset for universal CIR, covering 21 types of degradations from 7 popular traditional and learned codecs. Extensive experiments on universal CIR have demonstrated the excellent robustness and texture restoration capability of our proposed MoE-DiffIR.
## :sparkles: Getting Start
### Prepare Datasets:
We propose a comprehensive benchmark dataset for universal compressed image restoration (CIR), covering 21 types of degradations from 7 popular traditional and learned codecs. Traditional and learned codecs include HIFIC, $$C_{SSIM}$$, $$C_{PSNR}$$. Within each codec, we apply three levels of distortions:
- **JPEG**: QF=10,15,20
- **VVC**: QP=37,42,47
- **HEVC**: QP=37,42,47
- **WEBP**: QF=1,5,10
- **HIFIC**: Mode='low', 'med', 'high'
- $$C_{SSIM}$$: Mode=1,2,3
- $$C_{PSNR}$$: Mode=1,2,3
Our CIR dataset is based on DF2K, with the original DF2K images considered as the ground truth. All compression codecs are applied to these images. Here we release the [datalink](https://drive.google.com/drive/folders/1Kn8SjJWpHITHlg5kuL1Ur7Ml-WNJJ064) of CIR dataset including 'CIR_datasets' and 'CIR_Unseen_Tasks'. The 'CIR_Dataset' branch of this dataset is shown in the following figure. The dataset comprises training and testing sets. In the training set, "DF2K_HR" contains the ground truth (GT) images, while the other folders correspond to 21 types of low-quality (LQ) images resulting from various compression distortions. The five testing sets also include folders corresponding to the 21 types of compressed LQ images, as well as folders for the high-resolution (HR) images.
Here we also release a small test_set including 'LIVE1' and 'DIV2K_Test'. [link](https://rec.ustc.edu.cn/share/3b079e10-fa65-11ef-991a-6512ec9141fb) (password:cqq3)
### Prepare Environment:
- Python 3.9
- PyTorch 1.12.1 + cu113
Install other requiremnets:
```
pip install -r requirements.txt
```
### Inference code:
Here we release three checkpoints: (i) Only using MoE-Prompt (ii) Use Both MoE-Prompt and V2T Adapter (iii) Use MoE-Prompt, V2T Adapter and DA-CLIP degradation Prior.
You can find three weights in this [link](https://rec.ustc.edu.cn/share/84315560-c692-11ef-81e4-a301640e2c6f). (Password: 283i)
Inference code:
```
bash test.sh
```
Here you could download the klvae_ckpt from this [link](https://drive.google.com/drive/folders/1G6Hn9djg7_1-Qijgj619fMd262uXCrNu?usp=sharing).
### Training code:
Training code:
```
bash train.sh
```
### Visual Results:
The link of our Moe-DiffIR Visual Reuslt:[link](https://rec.ustc.edu.cn/share/6a05f850-0644-11f0-bea8-1319946c8ca0)
## Results
## Cite US
Please cite us if this work is helpful to you.
```
@inproceedings{ren2025moe,
title={Moe-diffir: Task-customized diffusion priors for universal compressed image restoration},
author={Ren, Yulin and Li, Xin and Li, Bingchen and Wang, Xingrui and Guo, Mengxi and Zhao, Shijie and Zhang, Li and Chen, Zhibo},
booktitle={European Conference on Computer Vision},
pages={116--134},
year={2025},
organization={Springer}
}
```
## Acknowledgments
The basic code is partially from the below repos.
- [StableSR](https://github.com/IceClear/StableSR)
- [StableDiffusion](https://github.com/Stability-AI/stablediffusion)