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

2025-11-25

AI Summary

Existing flow-based image inpainting methods suffer from slow sampling and over-smoothed results. This paper proposes a training-free, mask-guided flow matching framework for image restoration: it explicitly steers the flow matching sampling trajectory via a degradation mask and introduces a trajectory correction mechanism to enforce structural consistency between the generation process and the degraded input. To our knowledge, this is the first work to integrate mask guidance with flow matching modeling and to introduce a training-free trajectory correction technique. The method is broadly applicable to denoising, super-resolution, and general image inpainting tasks. It achieves state-of-the-art performance on both natural images (e.g., Places2) and medical images (e.g., BraTS), significantly outperforming diffusion models and existing flow matching approaches. It delivers superior visual quality—higher PSNR, SSIM, and lower LPIPS—faster inference (2–5× speedup), and enhanced detail fidelity.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Addresses lengthy sampling times in image restoration models
• Combats over-smoothed results in flow-based restoration methods
• Solves mask-based degradation tasks including inpainting and super-resolution

Innovation

• Flow matching for faster image generation
• Mask-guided sampling for image restoration
• Trajectory correction for input consistency