Zero-Shot Image Restoration Using Few-Step Guidance of Consistency Models (and Beyond)

TL;DR

This paper introduces a zero-shot image restoration method using Consistency Models that achieves high-quality results with as few as 4 neural function evaluations, significantly reducing computational cost.

Contribution

It proposes a novel noise injection mechanism combined with better initialization and back-projection guidance for efficient zero-shot restoration with CMs, requiring fewer NFEs.

Findings

Achieves effective image restoration with only 4 NFEs.

Noise injection improves performance of guided diffusion models.

Method outperforms existing zero-shot approaches in super-resolution, deblurring, and inpainting.

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.