Restora-Flow: Mask-Guided Image Restoration with Flow Matching

Abstract: Flow matching has emerged as a promising generative approach that addresses the lengthy sampling times associated with state-of-the-art diffusion models and enables a more flexible trajectory design, while maintaining high-quality image generation. This capability makes it suitable as a generative prior for image restoration tasks. Although current methods leveraging flow models have shown promising results in restoration, some still suffer from long processing times or produce over-smoothed results. To address these challenges, we introduce Restora-Flow, a training-free method that guides flow matching sampling by a degradation mask and incorporates a trajectory correction mechanism to enforce consistency with degraded inputs. We evaluate our approach on both natural and medical datasets across several image restoration tasks involving a mask-based degradation, i.e., inpainting, super-resolution and denoising. We show superior perceptual quality and processing time compared to diffusion and flow matching-based reference methods.

• The paper introduces flow matching guided by degradation masks to efficiently restore images in tasks like denoising, super-resolution, and inpainting.
• It employs a trajectory correction mechanism that refines the sampling process, achieving impressive LPIPS, SSIM, and PSNR metrics.
• Empirical evaluations on diverse datasets demonstrate that Restora-Flow outperforms state-of-the-art methods with lower computational costs and enhanced perceptual quality.

Restora-Flow: Mask-Guided Image Restoration with Flow Matching

Introduction

The paper "Restora-Flow: Mask-Guided Image Restoration with Flow Matching" (2511.20152) introduces a novel approach to address image restoration tasks, utilizing the principles of flow matching as a generative prior. This method particularly targets tasks like denoising, super-resolution, and inpainting, all framed as inverse problems involving mask-based degradation. The cornerstone of this approach is leveraging flow matching to achieve high-quality image restoration with reduced processing time compared to traditional diffusion models.

Flow Matching and Its Role in Restoration

Flow matching, a methodology for generative modeling, defines optimal data transformation paths using velocity fields, offering more efficient sampling compared to diffusion models that suffer from prolonged sampling due to complex trajectories. The methodology involves Ordinary Differential Equations (ODEs) that guide the flow from a simple base distribution to the target. Restora-Flow extends this by introducing a mask-guided sampling procedure and a trajectory correction mechanism. These enhancements ensure that reconstructed images not only maintain high fidelity to the data but also align closely with the user-provided degradation masks.

In the sampling process, Restora-Flow iteratively integrates flow matching updates that are guided by degradation masks, ensuring consistency and relevance to the input data. One crucial component introduced is trajectory correction, which aligns the flow trajectory with the true data distribution, refining the sample quality.

Figure 1: Restora-Flow samples with and without correction steps. Empirically, one correction step ($C=1$) offers the best trade-off between high reconstruction quality and fast processing.

Experimental Results

The experiments conducted demonstrate Restora-Flow's superior performance across multiple datasets, including CelebA, AFHQ-Cat, COCO, and X-ray Hand. The evaluation encompasses significant image restoration metrics like LPIPS, SSIM, PSNR, and processing time. Remarkably, Restora-Flow consistently outperforms state-of-the-art methods, delivering superior perceptual quality while maintaining lower computational demands.

The results on CelebA showcase the paradigm's proficiency across various tasks, achieving the best LPIPS scores for most scenarios and demonstrating close to optimal SSIM and PSNR scores. Such outcomes assert Restora-Flow's capability in producing high-quality restorations promptly.

Figure 2: Qualitative results on CelebA. Shown are the degraded image (col 1), original image (col 2), restored images using related work as indicated (cols 3-8) and restored image of Restora-Flow (col 9). Rows refer to denoising (row 1), box inpainting (row 2), super-resolution (row 3) and random inpainting (row 4). RePaint is not applicable to denoising (N/A). Differences can be best seen in the pdf version.

Across diverse datasets like AFHQ-Cat and COCO, Restora-Flow maintains its leading performance, confirming its adaptability and effectiveness in different contexts and image types. The research emphasizes the empirical trade-offs in correction steps (C), highlighting the efficacy of minimal interventions for optimal results (Figure 1).

Implications and Future Work

Restora-Flow presents critical advancements in image restoration by significantly reducing the computational cost of generating high-resolution, artifact-free images. This positions it as a valuable tool in applications where rapid and accurate image processing is essential, such as medical imaging and visual content restoration.

Future work can explore the extension of Restora-Flow to scenarios involving non-mask-based degradation operators, broadening its applicability and enhancing its robustness in handling more complex image restoration challenges. Such exploration could substantially impact AI-driven imaging technologies by further reducing processing times and improving reconstruction fidelity.

Conclusion

Restora-Flow introduces a compelling methodology for image restoration leveraging flow matching, demonstrating superior quality and efficiency across multiple image datasets. By optimizing reconstruction fidelity through mask-guided sampling and trajectory correction, it sets new benchmarks for image restoration tasks, paving the way for advancing generative modeling applications in AI. The versatility and robustness evidenced through empirical evaluation invite exploration into further enhancements and broader applications in AI imaging domains.