DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior

Abstract

我们提出了DiffBIR，它利用预训练的文本到图像扩散模型来解决盲图像恢复问题。我们的框架采用两阶段pipeline。在第一阶段，我们在多种退化中预训练恢复模块，以提高现实场景中的泛化能力。第二阶段利用潜在扩散模型的生成能力，实现真实的图像恢复。具体来说，我们引入了一个injective modulation 子网络——LAControlNet进行微调，而预训练的Stable Diffusion是为了保持其生成能力。最后，我们引入了一个可控模块，通过在推理过程的去噪过程中引入潜在图像引导，使用户能够平衡质量和保真度。大量的实验已经证明了它在合成和现实世界数据集上的盲图像超分辨率和盲人脸恢复任务方面优于最先进的方法。

1. Introduction

图像恢复的目的是从低质量的观测中重建出高质量的图像。典型的图像恢复问题，如图像去噪、去模糊和超分辨率，通常是在受限的环境下定义的，其中退化过程是简单和已知的(例如高斯噪声和双三次降采样)。他们成功地推广了大量优秀的恢复算法，但天生具有有限的泛化能力。为了处理现实世界中退化的图像，盲图像恢复 (BIR) 成为一个很有前途的方向。BIR 的最终目标是在具有一般退化的一般图像上实现真实的图像重建。BIR不仅扩展了经典图像修复任务的边界，而且具有广泛的实际应用领域(如旧照片/胶片修复)。

BIR的研究还处于初级阶段，需要对其现状进行更多的解释。根据问题设置的不同，现有的BIR方法大致可以分为三个研究方向，

• 盲图像超分辨率(BSR)
• zero-shot图像恢复(ZIR)
• 盲人脸恢复(BFR)

• 盲图像超分辨率（Blind Image Super-Resolution，简称BSR）是指在图像超分辨率重建任务中，不需要事先了解输入低分辨率图像的真实内容或放大倍数的情况下进行超分辨率处理。算法试图自动估计图像的真实内容和放大倍数，然后通过学习从低分辨率到高分辨率的映射关系来进行图像超分辨率处理。
• 零样本图像恢复（Zero-shot Image Restoration）是指在没有特定任务的训练样本的情况下，通过模型处理具体任务的图像恢复问题。在图像恢复中，任务通常包括超分辨率、去噪、去模糊等。零样本图像恢复的目标是在没有特定任务的先验知识的情况下，使模型能够有效地处理新任务的图像恢复需求。
• 盲人脸恢复（Blind Face Restoration）旨在从低质量的人脸图像中恢复高质量的人脸图像，广泛应用于监控图像修复、老照片修复和人脸图像超分辨率等各种场景。

它们都取得了显著的进步，但也有明显的局限性。BSR最初是为了解决现实世界的超分辨率问题而提出的，其中低分辨率图像包含未知的退化。根据最近的BSR综述调查[37]，最流行的解决方案可能是BSRGAN[64]和Real-ESRGAN[55]。他们将BSR表述为一个有监督的大规模退化过拟合问题。

有监督的大规模退化过拟合问题是指当在大规模有标签的数据集上进行训练时，模型会过度依赖于训练集中的噪声或异常样本，导致在未见过的数据上表现较差的现象。
在有监督学习中，模型通过使用包含特征和相应标签的训练数据来学习从输入到输出的映射关系。然而，当训练集非常庞大时，可能存在一些噪声或错误标签的样本。如果模型过度拟合这些异常样本，它可能无法在真实世界中泛化，而且对未知数据的预测性能会下降。
退化过拟合问题还可能与训练数据的分布不一致有关。如果训练集与真实世界中的数据分布有显著差异，模型可能会在训练数据上表现良好，但在未知数据上表现较差。

• 为了模拟真实的退化，分别提出了一种退化混洗策略和高阶退化建模。然后是对抗性损失[31;17;56;41;49]以端到端方式学习重建过程。它们确实消除了一般图像上的大部分退化，但无法生成逼真的细节。此外，它们的退化设置仅限于×4/ × 8超分辨率，这对于BIR问题来说是不完整的。
• 第二组ZIR是一个新出现的方向。代表作品有DDRM[26]、DDNM[57]、GDP[16]。它们将强大的扩散模型作为附加先验，因此比基于GAN的方法具有更大的生成能力。通过适当的退化假设，它们可以在经典IR任务中实现令人印象深刻的zero-shot恢复。但是，ZIR的问题设置与BIR不一致。他们的方法只能处理明确定义的退化(线性或非线性)，但不能很好地推广到未知的退化。换句话说，它们可以在一般图像上实现逼真的重建，但不能在一般退化图像上实现。
• 第三类是BFR，主要研究人脸修复。最先进的方法可以参考CodeFormer[68]和VQFR[18]。它们具有与BSR方法相似的pipeline，但在退化模型和生成网络上有所不同。由于图像空间较小，这些方法可以利用VQGAN和Transformer在真实世界的人脸图像上取得令人惊讶的好结果。然而，BFR只是BIR的一个子域。它通常假设固定的输入大小和有限的图像空间，因此不能应用于一般图像。

由以上分析可知，我们可以看到，现有的BIR方法不能实现(1)真实的图像重建，(2)一般图像，(3)一般退化。因此，我们希望一种新的BIR方法来克服这些限制。

在这项工作中，我们提出了DiffBIR，将以往工作的优点整合到一个统一的框架中。具体来说，DiffBIR

• 采用了一种扩展的退化模型，可以推广到现实世界的退化；
• 利用训练良好的Stable Diffusion作为先验来提高生成能力；
• 引入了一个两阶段的pipeline来保证真实和保真。

我们也做了专门的设计来实现这些策略。首先，为了提高泛化能力，我们将BSR的多种退化类型和BFR的广泛退化范围结合起来，建立了一个更实用的退化模型。这有助于DiffBIR处理各种极端退化情况。其次，为了利用Stable Diffusion，我们引入了一个injective modulation 子网络——LAControlNet，可以针对我们的特定任务进行优化。与ZIR类似，预训练的Stable Diffusion在微调期间是固定的，以保持其生成能力。第三，为了实现真实和逼真的图像重建，我们首先应用恢复模块(即SwinIR)来减少大多数退化，然后微调生成模块(即LAControlNet)来生成新的纹理。如果没有这个pipeline，模型可能会产生过度平滑的结果(删除生成模块)或生成错误的细节(删除恢复模块)。此外，为了满足用户多样化的需求，我们进一步提出了一个可控模块，可以实现第一阶段的恢复结果和第二阶段的生成结果之间的连续过渡效果。这是通过在去噪过程中引入潜在图像引导而无需重新训练来实现的。适用于潜在图像距离的梯度尺度可以调整以权衡真实感和保真度。

设置了上述组件，所提出的DiffBIR在合成和现实数据集上的BSR和BFR任务中都表现出优异的性能。值得注意的是，DiffBIR在一般图像恢复方面实现了很大的性能飞跃，优于现有的BSR和BFR方法(如BSRGAN[64]、Real-ESRGAN[55]、CodeFormer[68]等)。我们可以观察到这些方法在某些方面的差异。对于复杂的纹理，BSR方法往往会产生不真实的细节，而DiffBIR方法可以产生视觉上令人愉悦的结果，参见图1(第一行)。对于语义区域，BSR方法倾向于实现过度平滑的效果，而DiffBIR可以重建语义细节，参见图1(第二行)。对于微小的条纹，BSR方法倾向于删除这些细节，而DiffBIR方法仍然可以增强它们的结构，参见图1(第三行)。此外，DiffBIR能够处理极端的降级并重新生成逼真而生动的语义内容，参见图1(最后一行)。这些都表明DiffBIR成功地打破了现有BSR方法的瓶颈。对于盲人面部恢复，DiffBIR在处理一些困难的情况下表现出优势，例如在被其他物体遮挡的面部区域保持良好的保真度(见图1 (b)第一行)，在面部区域之外成功恢复(见图1 (b)第一行）。综上所述，我们的DiffBIR首次能够在统一的框架内获得具有竞争力的BSR和BFR任务性能。广泛而深入的实验证明了我们提出的DiffBIR优于现有的最先进的BSR和BFR方法。

2. Relate work

盲图像超分辨率(BSR)。BSR的最新进展[37]探索了更复杂的退化模型来近似真实世界的退化。BSRGAN[64]旨在基于随机混洗策略合成更实用的退化，Real-ESRGAN[55]利用“高阶”退化建模。它们都利用GAN [17;41;49个;31;56]学习复杂退化下的图像重建过程。SwinIR-GAN[36]使用新的主流骨干Swin Transformer[38]来实现更好的图像恢复性能。FeMaSR[6]将SR表述为基于预训练的VQ-GAN的特征匹配问题[15]。尽管BSR方法可以有效地消除现实世界中的退化，但它们并不擅长生成真实的细节。此外，他们通常假设低质量图像输入被一些特定的尺度(例如×4/× 8)下采样，这对于BIR问题是有限的。

Zero-shot图像恢复(ZIR)。ZIR旨在通过以无监督的方式利用预训练的先验网络来实现图像恢复。早期工作[2;10;40;44]主要集中于在预训练的GAN的潜在空间中搜索潜在编码。该领域的最新进展包括使用去噪扩散概率模型[21;51;52;46;45;48]。DDRM[26]引入了一种基于SVD的方法来有效地处理线性图像恢复任务。同时，DDNM[57]从理论上分析了向量的范围-零空间分解，并基于零空间设计了采样方案。受分类器引导[12]的启发，GDP[16]引入了一种更方便有效的引导方法，可以在推理过程中对退化模型进行估计。虽然这些工作促进了zero-shot图像恢复技术的进步，但对于现实世界的低质量图像，ZIR方法仍然不能达到令人满意的恢复效果。

盲人脸图像修复(BFR)。人脸图像作为一般图像的特定子域，通常承载着更多的结构和语义信息。早期的尝试利用几何先验(例如面部解析图[5]、面部地标[9;27]，面部成分热图[62])或参考先验[34;33;32;13]作为辅助信息来指导面部恢复过程。随着生成网络的快速发展，许多BFR方法都采用了强大的生成先验算法来重建高真实感的人脸。基于GAN先验的代表性方法[54;61;19;4]已经证明了他们在实现高质量和高保真的面部重建方面的能力。最先进的工作[68;18;59]引入HQ码本，通过利用矢量量化(VQ)字典学习来生成令人惊讶的逼真面部细节[53;15)。

3. Methodology

在这项工作中，我们的目标是利用强大的生成 prior-Stable Diffusion来解决普通图像和人脸图像的盲恢复问题。我们提出的框架采用有效、稳健和灵活的两阶段pipeline。首先，我们使用恢复模块来去除损坏，例如使用回归损失的噪声或失真伪影。由于丢失的局部纹理和粗/细细节仍然不存在，我们然后利用 Stable Diffusion来弥补信息丢失。整个框架如图2所示。具体来说，我们首先在大规模数据集上预训练一个SwinIR[36]，以实现跨多种退化的初步退化去除(第3.1节)。然后，利用生成先验来产生逼真的恢复结果(第3.2节)。此外，还引入了一种基于潜在图像引导的可控模块，用于在真实感和保真度之间进行权衡(第3.3节)。

3.1 退化去除预训练

退化模型。BIR旨在从具有未知和复杂退化的低质量(LQ)图像中恢复干净图像。通常，模糊，噪声，压缩伪影和低分辨率经常涉及。为了更好地覆盖LQ图像的退化空间，我们采用了考虑多样化退化和高阶退化的综合退化模型。在所有退化中，模糊、调整大小和噪声是现实场景中的三个关键因素[64]。我们的多样化退化涉及

• 模糊：各向同性高斯核和各向异性高斯核；
• 调整大小：区域调整大小，双线性插值和双立方调整大小；
• 噪声：加性高斯噪声、泊松噪声和JPEG压缩噪声。

关于高阶退化，我们按照[55]使用二阶退化，它重复经典退化模型：模糊-调整大小-噪声过程两次。请注意，我们的退化模型是为图像恢复而设计的，因此所有退化的图像将被重新调整到原始大小。

所谓的高阶退化模型通俗的说就是将经典退化算法排列组合

恢复模块。为了构建一个鲁棒的生成图像恢复pipeline，我们采用了一种保守但可行的解决方案，首先去除LQ图像中的大部分退化(特别是噪声和压缩伪像)，然后使用随后的生成模块来重现丢失的信息。这种设计将促进潜在扩散模型更多地关注纹理/细节的生成，而不受噪声损坏的干扰，并且在没有错误细节的情况下获得更真实/清晰的结果(参见4.3节)。我们修改了SwinIR 作为我们的恢复模块。具体来说，我们利用像素重排[50]操作以8的比例因子对原始低质量输入 $I_{LQ}$进行下采样。然后，采用3 × 3卷积层进行浅层特征提取。所有后续的Transformer操作都在低分辨率空间中进行，类似于潜在扩散模型。深度特征提取采用多个残差Transformer块(Residual Swin Transformer block, RSTB)，每个RSTB有多个SwinTransformer层(STL)。将增加浅层和深层特征，以同时维护低频和高频信息。为了将深度特征上采样回原始图像空间，我们执行了三次最接近插值，每次插值之后是一个卷积层和一个Leaky ReLU激活层。我们通过最小化$L_2$像素损失来优化恢复模块的参数。公式如下：
$$I_{reg}=SwinIR(I_{LQ}),{\mathcal{L}}_{reg}=\Vert I_{reg}-I_{HQ}{\Vert }_2^2\text{\hspace{1em}(1)}$$其中$I_{HQ}$和$I_{LQ}$分别表示高质量图像和低质量对应图像。$I_{reg}$是通过回归学习得到的，将用于对潜在扩散模型的微调。

3.2 利用生成先验进行图像重建

初步: Stable Diffusion。在本文中，我们基于大规模文本到图像的潜在扩散模型——Stable Diffusion来实现我们的方法。扩散模型通过估计数据分布分数的去噪序列来学习生成数据样本。为了获得更好的效率和稳定的训练，Stable Diffusion预训练了一个自动编码器[29]，该编码器使用编码器 $\mathcal{E}$将图像 $x$ 转换为潜在 $z$，并使用解码器$\mathcal{D}$进行重建。这种潜在表示是通过使用VAE[30]、Patch-GAN[23]和LPIPS[67]的混合目标来学习的。扩散和去噪过程在潜在空间中进行。在扩散过程中，将时刻 $t$ 方差${\beta }_t\in (0,1)$的高斯噪声加入到编码潜在函数$z=\mathcal{E}(x)$中，产生噪声潜在特征：
$$z_t=\sqrt{{\bar{\alpha }}_t}z+\sqrt{1-{\bar{\alpha }}_t}ϵ,\text{\hspace{1em}(2)}$$其中 $ϵ\sim N(0,I),{\alpha }_t=1-{\beta }_t$和${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$。当 $t$ 足够大时，潜在的$z_t$接近于标准高斯分布。

通过在随机选择的时间步 $t$上预测以$c$ (即文本提示)为条件的噪声 $\lambda$ 来学习网络 ${ϵ}_{\theta }$。潜在扩散模型的优化定义如下：
$${\mathcal{L}}_{ldm}={\mathbb{E}}_{z,c,t,ϵ}[||ϵ-{ϵ}_{\theta }(z_t=\sqrt{{\bar{\alpha }}_t}z+\sqrt{1-{\bar{\alpha }}_t}ϵ,c,t)|{|}_2^2],\text{\hspace{1em}(3)}$$

其中$x,c$ 表示从数据集中采样，$z=\mathcal{E}(x)$， t 均匀采样， $ϵ$ 从标准高斯分布中采样。

LAControlNet。虽然第一阶段可以消除大多数退化，但获得的图像通常过于平滑，距离高质量自然图像的分布还很远。然后，我们利用预训练的Stable Diffusion与我们获得的$I_{reg}-I_{HQ}$对进行图像重建。首先，我们利用Stable Diffusion预训练的VAE编码器将$I_{reg}$ 映射到潜在空间中，得到条件潜在$\mathcal{E}(I_{reg})$。UNet[47]去噪器执行潜在扩散，它包含一个编码器、一个中间块和一个解码器。特别是，解码器从编码器接收特征并在不同尺度上进行融合。在这里，我们创建一个并行模块 (在图2中用橙色表示)，它包含与UNet去噪器中相同的编码器和中间块。然后，我们将条件潜在$\mathcal{E}(I_{reg})$与随机采样的噪声 $z_t$ 合并起来作为并行模块的输入。由于此串联操作将增加并行模块中第一个卷积层的通道数，因此我们将新添加的参数初始化为零，其中所有其他权重都从预训练的UNet去噪检查点初始化。并行模块的输出被添加到原始UNet解码器。此外，在每个尺度的加法运算之前，应用一个1 × 1的卷积层。在微调过程中，并行模块和这些1×1卷积层同时进行优化，其中提示条件设置为空。我们的目标是尽量极小化以下潜在扩散目标：

$${\mathcal{L}}_{ldm}={\mathbb{E}}_{z_t,c,t,\mathcal{E}(I_{reg})}[||ϵ-{ϵ}_{\theta }(z_t,c,t,\mathcal{E}(I_{reg})|{|}_2^2],\text{\hspace{1em}(4)}$$

这一阶段得到的结果记为$I_{dif}$ 。总而言之，只有UNet去噪器中的跃迁连接特征才会针对我们的特定任务进行调整。该策略缓解了小训练数据集的过拟合问题，并继承了Stable Diffusion的高质量生成。更重要的是，与ControlNet[66]相比，我们的条件调节机制对于图像重建任务更直接有效，ControlNet使用从头开始训练的附加条件网络来编码条件信息。在我们的LAControlNet中，训练有素的VAE编码器能够将条件图像投影到与潜在变量相同的表示空间中。该策略显著减轻了潜在扩散模型中内部知识与外部条件信息对齐的负担。在实验中，直接使用ControlNet进行图像重建会导致严重的颜色偏移。

3.3 保真度-真实性权衡的潜在图像引导

虽然上述两阶段的方法已经可以获得很好的恢复效果，但由于用户的不同偏好，仍然需要在真实和保真之间进行权衡。因此，我们提出了一个可控的模块，可以引导降噪过程在第一阶段获得的$I_{reg}$，从而在真实和平滑结果之间进行调整。分类器引导是由Dhariwal和Nichol[12]提出的，它利用在噪声图像上训练的分类器来引导生成到目标类。而在大多数情况下，作为指导的预训练模型通常是在干净的图像上训练的。为处理这种情况，工作在[1;16]转向引导中间变量${\tilde{x}}_0$来控制扩散模型的生成过程。具体来说，在采样过程中，他们通过估计 $x_t$中的噪声，从噪声图像 $x_t$中估计出一个干净的图像 $x_0$。在这项工作中，扩散和去噪过程是基于潜在空间的。因此，我们的目标是得到一个干净的潜在$z_0$，通过以下公式：
$${\tilde{z}}_0=\frac{z_t}{\sqrt{{\bar{\alpha }}_t}}-\frac{\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }\left(z_t,c,t,\mathcal{E}\left(I_{reg}\right)\right)}{\sqrt{{\bar{\alpha }}_t}}.\text{\hspace{1em}(4)}$$

然后，将基于潜在的损失${\mathcal{D}}_{latent}$定义为潜在图像引导$\mathcal{E}(I_{reg})$与估计的干净潜在变量${\tilde{z}}_0$之间的 ${\mathcal{L}}_2$ 距离：
$${\mathcal{D}}_{\text{latent }}\left(x,I_{\text{reg }}\right)=\mathcal{L}\left({\tilde{z}}_0,\mathcal{E}\left(I_{\text{reg }}\right)\right)=\sum _j\frac{1}{C_j{H}_j{W}_j}\Vert {\tilde{z}}_0-\mathcal{E}\left(I_{\text{reg }}\right)){\Vert }_2^2\text{\hspace{1em}(5)}$$

上述引导可以迭代地强制潜在特征之间的空间对齐和颜色一致性，并引导生成的潜在保留参考潜在的内容。因此，可以控制从参考图像$I_{reg}$中保留多少信息(如结构、布局和颜色)，从而实现从生成的输出到更平滑的结果的过渡。我们的潜在图像引导的整个算法在算法1中进行了说明。

4. Experiments

4.1 数据集、实现、度量

数据集。我们在ImageNet[11]数据集上以512 × 512分辨率训练DiffBIR。对于BFR，我们使用FFHQ[25]数据集，并将其大小调整为512 × 512。为了合成LQ图像，我们在训练过程中利用提出的退化pipeline对HQ图像进行处理(详见附录A)。对于BSR，我们使用RealSRSet[3]数据集在现实环境中进行比较。为了在现实场景中进行更彻底的比较，我们从互联网上收集了47张图像，记为Real47。它包含了各种场景的一般图像，如自然户外景观，老照片，建筑，人物从肖像到密集的人群，植物，动物等。对于BFR任务，我们在合成数据集CelebA-Test[39]和三个真实数据集：LFW-Test[54]、CelebChild-Test[54]和WIDER-Test[68]上评估了我们的方法。特别是，CelebA-Test包含从CelebA-HQ数据集中选择的3,000张图像，其中LQ图像是在与我们的训练设置相同的退化范围下合成的。

实现。恢复模块采用8个残差Swin Transformer block (RSTB)，每个RSTB包含6个Swin Transformer Layers (STL)。head数设置为6，窗口大小设置为8。我们训练恢复模块，批量大小为96，迭代150k。我们使用Stable Diffusion 2.1-base 3作为生成先验，并对扩散模型进行了微调，迭代次数为25k，批大小为192。我们使用Adam[28]优化器，将学习率设置为10−4。训练过程在512 × 512分辨率下使用8个NVIDIA A100 GPU进行。对于推理，我们采用间隔DDPM采样[43]，时间步长为50。我们的DiffBIR能够处理大于512 × 512的任意大小的图像。对于边长< 512的图像，我们首先将它们的短边放大到512，然后将它们调整回来。

度量。对于基于真实值的评价，我们采用传统的指标：PSNR、SSIM和LPIPS[67]。为了更好地评估BIR任务的真实性，我们还包括了几个无参考图像质量评估(IQA)指标：MANIQA4[60]和NIQE。对于BFR，我们评估了identity preservation-IDS[68]，并采用了广泛使用的感知度量FID[20]。我们还部署了一项用户研究，以进行更彻底的比较。

4.2 与最先进方法的比较

对于BSR，我们将DiffBIR与最先进的BSR方法进行了比较：Real-ESRGAN+[55]、BSRGAN[64]、SwinIR-GAN[36]和FeMaSR[6]。最近最先进的ZIR方法(DDNM[57]和GDP[16])也包括在内。对于BFR任务，我们比较了最新的最先进的方法：DMDNet[35]、GFP-GAN[54]、GPEN[61]、GCFSR[19]、VQFR[18]、CodeFormer[68]、RestoreFormer[59]。

真实数据集上的BSR

我们在表1中提供了对真实数据集的定量比较。可以观察到，我们的DiffBIR在广泛使用的RealSRSet[24]和我们收集的Real47上都在MANIQA中获得了最好的分数。而BSRGAN和Real-ESRGAN+在两个数据集上都可以在MANIQA中获得前3名的结果。

目视比较结果如图3所示。可以看出，DiffBIR能够更自然地还原文本信息，而其他方法往往会扭曲字符或产生模糊的输出。另一方面，我们的DiffBIR也可以为自然图像生成逼真的纹理细节，而其他方法会产生过于光滑的结果。

为了进一步将DiffBIR与其他最先进的方法进行比较，我们对收集的Real47数据集进行了用户研究。本用户研究比较了DiffBIR、SwinIR-GAN、BSRGAN和RealESRGAN+。对于每张图像，用户被要求对四种方法的结果进行排序，并按升序为不同的方法分配1-4点。更准确地说，结果越好，得分越高。在详细的指导下，招募了31名用户进行这项用户研究。各方法得到的分数分布如图4所示。可以看出，DiffBIR的中位数得分最高，其上四分位数超过3。这表明用户倾向于将DiffBIR的结果排在首位。用户研究结果再次证明DiffBIR的视觉结果优于其他方法，这与它在MANIQA上的最高分一致。