LatentINDIGO: An INN-Guided Latent Diffusion Algorithm for Image Restoration

TL;DR

LatentINDIGO introduces wavelet-inspired invertible neural networks to guide latent diffusion-based image restoration without requiring explicit degradation models. It presents two implementations: PixelINN (pixel-domain) and LatentINN (latent-domain), both alternating INN-guided updates with forward-model refinement and a manifold-regularization term to maintain natural-image plausibility. The methods achieve state-of-the-art or competitive results on synthetic and real degraded images, including blind face restoration and natural images, and support arbitrary output sizes via patch-based latent inference. The framework integrates with existing LDM pipelines with no retraining of the pretrained models, offering a practical, scalable solution with strong restoration fidelity and perceptual quality.

Abstract

There is a growing interest in the use of latent diffusion models (LDMs) for image restoration (IR) tasks due to their ability to model effectively the distribution of natural images. While significant progress has been made, there are still key challenges that need to be addressed. First, many approaches depend on a predefined degradation operator, making them ill-suited for complex or unknown degradations that deviate from standard analytical models. Second, many methods struggle to provide a stable guidance in the latent space and finally most methods convert latent representations back to the pixel domain for guidance at every sampling iteration, which significantly increases computational and memory overhead. To overcome these limitations, we introduce a wavelet-inspired invertible neural network (INN) that simulates degradations through a forward transform and reconstructs lost details via the inverse transform. We further integrate this design into a latent diffusion pipeline through two proposed approaches: LatentINDIGO-PixelINN, which operates in the pixel domain, and LatentINDIGO-LatentINN, which stays fully in the latent space to reduce complexity. Both approaches alternate between updating intermediate latent variables under the guidance of our INN and refining the INN forward model to handle unknown degradations. In addition, a regularization step preserves the proximity of latent variables to the natural image manifold. Experiments demonstrate that our algorithm achieves state-of-the-art performance on synthetic and real-world low-quality images, and can be readily adapted to arbitrary output sizes.

Figures (14)

• Figure 1: Comparison of 4× blind super-resolution (SR) using DiffBIR lin2023diffbir (b) and our proposed approach (c).
• Figure 2: Illustration of two-level lifting scheme on a one-dimensional signal $\bm{x} \textlbrackdbl n \textrbrackdbl$. (a) shows the forward transform: at each level, the input is split into odd and even samples, and predict/update operations are applied to generate coarse and detail components. The coarse component is then passed to the next level for further decomposition, forming a multi-resolution representation. (b) shows the inverse transform, which reverses the process by using the same predict/update operations and a merge operator to combine the coarse and detail components at each level, ultimately reconstructing the original signal. In a wavelet-inspired invertible neural network, the predict and update operators are implemented using trainable neural networks.
• Figure 3: Architecture of the proposed PixelINN. (a) Forward pass: The first lifting level splits the input image $\bm{x}$ into coarse $\bm{x}_{c0}$ and detail $\bm{x}_{d0}$ components. The predict–update modules (PM and UM) then refine these components, producing $\bm{x}_{c1}$ and $\bm{x}_{d1}$. The next lifting level processes $\bm{x}_{c1}$ similarly, yielding a new coarse–detail pair $\bm{x}_{c2}, \bm{x}_{d2}$. (b) Inverse pass: The same PM and UM modules, combined with a merge operator, are used to recombine the coarse and detail subbands level by level. Starting from the final-stage subbands, the network reverses each lifting level and, owing to its perfect reconstruction property, can ultimately recover the original image $\bm{x}$. PixelINN is conditioned to the vector $\bm{\gamma}_{\mathrm{deg}}$ which embeds information about the degradation model.
• Figure 4: Architecture of the proposed PM/UM. Each PM/UM starts with a convolutional layer for channel adjustment, followed by our Modulated Residual Swin Transformer Blocks (MRSTBs), and ends with another convolutional layer that refines the feature maps and projects them into the desired output channels. The MRSTB module is composed of Swin Transformer Layers (STLs) with multi-head self-attention (MSA) and multi-layer perceptron (MLP), along with a learnable modulation mechanism conditioned on the degradation embedding $\bm{\gamma}_{\mathrm{deg}}$.
• Figure 5: The training framework of our LatentINN. During the forward process of our LatentINN, the latent code $\bm{z}_{x}$ of a clean image is transformed into the coarse part $\bm{z}_{c}$ and the detail part $\bm{z}_{d}$. Then we use inverse transform $g_\mathbf{\Theta}^{-1}([\bm{z}_{y};\bm{z}_{d}])$ with the latent code of the degraded measurement $\bm{z}_{y}=\mathcal{E}(\bm{y})$ to reconstruct $\bm{z}_{inv}.$