Zero-Shot Image Restoration via Few-Step Guidance of Consistency Models

Abstract

Recently, it has become popular to tackle image restoration tasks with a single pretrained (unconditional) denoising diffusion model (DDM) and data-fidelity guidance, instead of training a dedicated deep neural network per task. However, such "zero-shot" restoration schemes require many Neural Function Evaluations (NFEs). This follows from the need of iterative schemes with many NFEs already in the original generative functionality of the DDMs. Very recently, faster variants of DDMs have been explored for image generation. A prominent alternative are 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. Clearly, the latter is not a zero-shot strategy and, as such, leads to performance drop if the assumptions during the fine-tuning (e.g., the noise level) 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 and inpainting.