Image Restoration via Primal Dual Hybrid Gradient and Flow Generative Model

TL;DR

This paper introduces a flexible, efficient Plug-and-Play image restoration method that incorporates flow generative models and supports various fidelity terms, improving robustness to different noise types.

Contribution

It develops a primal-dual hybrid gradient based PnP algorithm that handles diverse data fidelity terms and integrates flow generative priors for enhanced image restoration.

Findings

Supports both $$ and $$ fidelity terms

Outperforms traditional methods under non-Gaussian noise

Effective in denoising, super-resolution, deblurring, inpainting

Abstract

Regularized optimization has been a classical approach to solving imaging inverse problems, where the regularization term enforces desirable properties of the unknown image. Recently, the integration of flow matching generative models into image restoration has garnered significant attention, owing to their powerful prior modeling capabilities. In this work, we incorporate such generative priors into a Plug-and-Play (PnP) framework based on proximal splitting, where the proximal operator associated with the regularizer is replaced by a time-dependent denoiser derived from the generative model. While existing PnP methods have achieved notable success in inverse problems with smooth squared $\ell_2$ data fidelity--typically associated with Gaussian noise--their applicability to more general data fidelity terms remains underexplored. To address this, we propose a general and efficient PnP algorithm inspired by the primal-dual hybrid gradient (PDHG) method. Our approach is computationally efficient, memory-friendly, and accommodates a wide range of fidelity terms. In particular, it supports both $\ell_1$ and $\ell_2$ norm-based losses, enabling robustness to non-Gaussian noise types such as Poisson and impulse noise. We validate our method on several image restoration tasks, including denoising, super-resolution, deblurring, and inpainting, and demonstrate that $\ell_1$ and $\ell_2$ fidelity terms outperform the conventional squared $\ell_2$ loss in the presence of non-Gaussian noise.