ControlMambaIR: conditional controls with state-space model for image restoration
摘要
Diffusion-based restoration frameworks typically use ControlNet, which is based on a CNN or Transformer, to extract conditional features from low-quality images to guide the restoration sampling process. However, existing diffusion backbones often face the dilemma between global receptive fields and efficient computation, hindering their practical application. Recently, the state-space model Mamba has emerged as a promising solution to the above dilemma, offering the ability to model long-range dependencies with linear complexity. However, the standard Mamba still faces certain challenges in diffusion model such as spatial-temporal feature interactions and multi-scale feature integration. In this paper, we propose ControlMambaIR, a novel framework that adapts the Mamba architecture for diffusion-based image restoration. By introducing dedicated modules for spatial-temporal feature interaction and multi-scale feature integration, ControlMambaIR enhances the original Mamba, enabling it to capture long-range, global image context with high efficiency. This design provides robust feature guidance during the reverse diffusion sampling process. Extensive experiments demonstrate that ControlMambaIR achieves highly competitive results compared with traditional and generative methods.