Zero-Shot Image Restoration Using Few-Step Guidance of Consistency Models (and Beyond)
Tomer Garber, Tom Tirer
TL;DR
This work addresses the challenge of zero-shot image restoration with diffusion-based priors by proposing CM4IR, a four-step restoration framework built on pretrained Consistency Models. CM4IR leverages initialization from the observation model, back-projection guidance, and a novel noise-injection mechanism that decouples denoising and injection noise and uses a Polyak-acceleration–style update, enabling high-quality results on super-resolution, deblurring, and inpainting with only 4 NFEs. The authors demonstrate that their approach outperforms other zero-shot methods that require many NFEs and, in several cases, rivals task-specific fine-tuned methods, while also improving the performance of existing guided DM methods when NFEs are drastically reduced. The contribution has practical impact by reducing computational cost and enabling robust restoration under varied degradation models, with code to be released for reproducibility.
Abstract
In recent years, it has become popular to tackle image restoration tasks with a single pretrained diffusion model (DM) and data-fidelity guidance, instead of training a dedicated deep neural network per task. However, such "zero-shot" restoration schemes currently require many Neural Function Evaluations (NFEs) for performing well, which may be attributed to the many NFEs needed in the original generative functionality of the DMs. Recently, faster variants of DMs have been explored for image generation. These include Consistency Models (CMs), which can generate samples via a couple of NFEs. However, existing works that use guided CMs for restoration still require tens of NFEs or fine-tuning of the model per task that leads to performance drop if the assumptions during the fine-tuning are not accurate. In this paper, we propose a zero-shot restoration scheme that uses CMs and operates well with as little as 4 NFEs. It is based on a wise combination of several ingredients: better initialization, back-projection guidance, and above all a novel noise injection mechanism. We demonstrate the advantages of our approach for image super-resolution, deblurring and inpainting. Interestingly, we show that the usefulness of our noise injection technique goes beyond CMs: it can also mitigate the performance degradation of existing guided DM methods when reducing their NFE count.