IR-CM: The Fast and General-Purpose Image Restoration Method Based on Consistency Model
Abstract
This paper proposes a fast and general-purpose image restoration method. The key idea is to achieve few-step or even one-step inference by conducting consistency distilling or training on a specific mean-reverting stochastic differential equations. Furthermore, based on this, we propose a novel linear-nonlinear decoupling training strategy, significantly enhancing training effectiveness and surpassing consistency distillation on inference performance. This allows our method to be independent of any pre-trained checkpoint, enabling it to serve as an effective standalone image-to-image transformation model. Finally, to avoid trivial solutions and stabilize model training, we introduce a simple origin-guided loss. To validate the effectiveness of our proposed method, we conducted experiments on tasks including image deraining, denoising, deblurring, and low-light image enhancement. The experiments show that our method achieves highly competitive results with only one-step inference. And with just two-step inference, it can achieve state-of-the-art performance in low-light image enhancement. Furthermore, a number of ablation experiments demonstrate the effectiveness of the proposed training strategy. our code is available at https://github.com/XiaoxuanGong/IR-CM.
Cite
Text
Gong and Ma. "IR-CM: The Fast and General-Purpose Image Restoration Method Based on Consistency Model." Neural Information Processing Systems, 2024. doi:10.52202/079017-2406Markdown
[Gong and Ma. "IR-CM: The Fast and General-Purpose Image Restoration Method Based on Consistency Model." Neural Information Processing Systems, 2024.](https://mlanthology.org/neurips/2024/gong2024neurips-ircm/) doi:10.52202/079017-2406BibTeX
@inproceedings{gong2024neurips-ircm,
title = {{IR-CM: The Fast and General-Purpose Image Restoration Method Based on Consistency Model}},
author = {Gong, Xiaoxuan and Ma, Jie},
booktitle = {Neural Information Processing Systems},
year = {2024},
doi = {10.52202/079017-2406},
url = {https://mlanthology.org/neurips/2024/gong2024neurips-ircm/}
}