摘要
扩散模型在生成高度复杂图像分布方面表现出强大的图像生成能力，可作为图像恢复的有力先验。现有方法常利用 ControNet 等技术，以低质量图像为条件从这些先验中采样高质量图像。然而，ControNet 通常需要复制原网络的大部分，导致随着先验规模增大，参数数量显著增加。本文提出一种相对轻量级的 Adapter，利用预训练先验的强大生成能力，实现照片级真实的图像恢复。该 Adapter 可适配于降噪 U-Net 和 DiT，表现卓越。
引言
图像恢复旨在从退化观测中恢复高质量图像。早期工作多使用 CNN 和 Transformer 等架构建立低质量到高质量图像的映射关系。最近，基于生成式先验的图像恢复理念兴起。该方法利用预训练先验的强大生成能力，以低质量图像为条件生成高质量图像。此前工作多使用 GAN 等作为先验，如 StyleGAN。这些基于 StyleGAN 的方法通常在 StyleGAN 的潜在空间中运行，实现可控生成。尽管 GAN 倒置方法效果有限，但这些方法在人脸图像恢复中取得较好效果。然而，由于 GAN 的局限性，这些方法对未对齐人脸数据集效果不佳。扩散模型的崛起为大规模和多样化数据集的分布拟合开启了新的可能。结合潜在扩散模型（LDM）和文本条件，模型生成高质量图像的能力得到进一步提升。目前，扩散模型已成为图像恢复中最常用的预训练先验。
使用预训练扩散模型作为生成式先验进行图像恢复，主要涉及微调模块，使先验能以低质量图像为条件进行条件生成。ControNet 可有效用于边缘检测等条件控制生成任务。然而，ControNet 参数量大，训练和推理负担重。与深度图和边缘条件不同，低质量图像通常包含更有意义的特征。因此，认为生成高质量图像的过程更像从分布中找到与低质量图像最接近的样本。基于此，提出用较小模块代替 ControNet，以更高效地实现基于先验的图像恢复。
为此，本文引入 Restoration Adapter，比 ControNet 更轻量，并插入原模型。Adapter 将低质量图像作为条件嵌入原模型特征，实现条件采样。此外，预训练先验通常在超大规模数据集上训练。在微调先验进行图像恢复任务时，通常收集少量高质量数据用于训练，且这些数据可能未出现在原数据集中。因此，决定微调先验。同时，认为微调先验可减轻 Restoration Adapter 的负担，提升结果。为高效微调，将 Adapter 引入预训练网络参数。由于 DiTs 逐渐受到青睐，本文方法可适配于两种架构，提出灵活框架。此外，为平衡保真度和多样性，引入可轻松适配各类采样的采样策略。如图 1 所示，本方法在各类图像中实现高质量恢复，并基于具体提示实现可控恢复。综上，本文贡献如下：（1）提出扩散恢复适配器框架，将 Adapter 整合入扩散降噪网络；（2）为 UNet 和 DiT 设计特定适配器，方便适配各类扩散先验；（3）引入简单采样策略，确保各种采样器的恢复保真度。
预备知识
扩散模型
扩散模型是一种通过一系列小的、可逆的步骤，将简单分布逐步转换为复杂分布的生成模型。核心思想是定义将数据逐渐叠加高斯噪声的前向扩散过程，以及用于从噪声中恢复数据的反向过程。前向扩散过程可通过随机微分方程（SDE）描述，详细阐述噪声如何随时间添加。反向过程旨在通过反向噪声添加来恢复原始数据，并可使用随机微分方程或常微分方程来建模。选择 SDE 或 ODE 取决于具体方法和采样器（解算器）的有效性。
恢复先验
恢复先验是经过预训练的生成式模型。在图像恢复任务中，低质量图像作为条件输入，通过在先验中进行条件采样，生成相应的高质量图像。LDM 是图像恢复任务中常用的先验，因其能生成高质量图像而基于文本。StableDiffusion 系列模型是基于大规模数据集训练的 LDM。本文选 StableDiffusion XL 和 StableDiffusion 3 作为先验。StableDiffusion XL 的降噪网络是 U-Net，StableDiffusion 3 基于 MM-DiT，类似 DiT。输出层使用零初始化，以确保训练初期稳定。
扩散恢复适配器
所提扩散恢复适配器由两部分组成：插入原降噪网络中的 Restoration Adapters，以及用于微调降噪网络特定参数的 Diffusion Adapters。选择预训练的潜在扩散模型，如基于 U-Net 的 StableDiffusion XL 和基于 DiT 的 StableDiffusion 3，作为扩散先验，因其在文本-图像生成任务中性能出色。对于 StableDiffusion XL，生成分为两个阶段，仅使用第一阶段的 base 模型。
恢复适配器
为使先验从低质量条件输入生成高质量图像，需额外的低质量条件处理模块。ControNet 通常的做法是复制原网络的部分，并与主网络并行处理输入，最后将结果相加。我们不如 ControNet 那种处理额外条件的独立网络，而是希望其更与主网络集成。为解决此问题，为基于 U-Net 的扩散和 DiT 引入 Restoration Adapter。图 2 所示的架构。提出的方法是在原降噪网络中插入少量适配器。将适配器置于其要插入的块之后。注意，适配器可选择性地分配给某些块，无需为所有块配备适配器。在 StableDiffusion XL 中，有三个下采样块、一个中间块和三个上采样块。图 2 中的“块 n”（0≤n≤N）对应这些块。例如，“块 0”表示第一个下采样块，“块 1”表示第二个下采样块，依此类推。对于 StableDiffusion 3，降噪网络由多个变换器块组成，具体是 MM-DiT 块。因此，“块 n”表示 n-th MM-DiT 块。插入适配器后，各块的输出传递给相应的适配器，而非下一个块。如图 3 所示，适配器有两个变体，分别针对 StableDiffusion XL 的 U-Net 和 StableDiffusion 3 的 DiT。适配器接收低质量（LQ）特征、t 时刻的降噪特征 Xt 和时间步嵌入 emb 作为输入。以 U-Net 架构为例，LQ 特征源自前一个适配器的输出，第一个适配器的输入是编码的 LQ 潜在特征。LQ 特征通过一个适配器块，同时时间嵌入经历线性变换和 SiLU 激活，确保其与 Xt 形状匹配。这些特征相加后，经过带有残差连接的两个适配器块，最后通过零初始化的线性层，并与 Xt 相加。适配器的输出传递给下一个降噪网络块。DiT 的适配器架构类似，但图 3 右侧显示了一些变化。
扩散适配器
恢复先验通常拟合于超大规模数据集。使用小型训练数据子集和一些未见过的数据对先验微调，同时希望微调原降噪网络以减轻 Restoration Adapter 的压力。此外，在微调先验时，要想整体效率，因此需设计高效微调方法。为充分利用先验的强大能力来生成高质量图像，从而增强恢复任务，在降噪网络的自注意力模块中添加适配器，并与我们的 Restoration Adapter 一起训练。为保持简单高效，选择常见的 LoRA 作为适配器。
训练数据和目标
训练 Diffusion Restoration Adapter 需高质量图像及其对应的退化图像。退化图像由遵循 RealESRGAN 设计的退化管道生成。此外，由于先验是文本-图像生成模型，利用多模态语言模型提取训练图像的短标题。为提高对所有退化程度的稳健性，以及在文本输入为空时的性能，可设置文本为空描述。此外，LQ 图像可作为模型的训练目标，而输入提示设置为退化描述，如“低质量”或“低分辨率”以及具体的退化应用，例如“压缩”。为了改进在 StableDiffusion XL 上使用有限训练资源的性能，为 StableDiffusion XL 应用稳健的降解编码器。训练过程使用 Diffusion Restoration Adapter 作为 LQ 图像条件处理器，将条件转换为预训练网络的中间特征。因此，从条件分布微调预训练先验，其中 LQ 图像作为条件。冻结预训练网络的参数，并仅对插入的适配器进行训练。训练目标严格遵循先验的原始训练过程。对于 StableDiffusion 3，由于其使用条件流匹配损失训练，故遵循相同的损失函数。对于 StableDiffusion XL，训练目标遵循 DDPM 使用的损失。
恢复采样策略
许多方法旨在提升扩散模型的生成质量，一般提升多样性和整体质量，但这些方法并非为恢复任务设计。因此，使用这些采样调度器可能导致高度不忠实的样本。提出一种可轻松适配任何采样技术、无需改动核心采样器代码的简单方法。此方法通过在生成采样循环中调整降噪方向。类似于分类器引导和 SUPIR，使用 LQ 图像作为引导。通过计算 t 时刻去噪潜在 zt 与 LQ 潜在之间的非归一化方向，然后将去噪潜在 z’ t 朝该方向调整一个因素。该因素随时间变化，并由映射函数 g 创建。 zt = z’ t + w·g(t, T)·(clq - z’ t) (1) 其中，clq 是 LQ 潜在，z’ t 是 t 时刻原始采样方法的输出，T 是采样步数总数，w 是加权超参数，g 是映射函数。g 的选择应遵循规则，赋予较大时间步更多的权重，反之亦然。这样，在降噪循环的早期阶段，降噪步骤更关注保真度，而在后期则考虑高质量和多样性。为简化起见，使用分段线性映射函数。该函数将 t 映射到[0, 1)，阈值超参数 a 确定分段行为。 g(t, T; a) = {(t/(T - a)) / (1 - a) if t/T > a else 0} (2) 其中，0 < t < T。 通过将 z’ t 朝 clq 移动，当 z’ t 接近 clq 时，可实现保真度更高的结果。
实验
定量比较
如表 2 所示，对于居中裁剪部分，基于 SD3 的方法在所有质量评估指标上获得最佳结果，而基于 SDXL 的方法在 ClipIQA（0.6517）和 MUSIQ（71.92）指标上排名第二。对于 8× 下采样的图像，基于 SD3 的方法在 ClipIQA（0.6659）和 MUSIQ（72.2）指标上获得最佳结果，而基于 SDXL 的方法在 MUSIQ（71.92）指标上排名第二。对于随机裁剪部分，基于 SD3 的方法在 ClipIQA（0.6868）和 MUSIQ（70.56）指标上获得最佳结果。基于 SDXL 的方法表现出竞争力，尤其是在 ManIQA（0.5358）和 LPIPS（0.334）指标上，并在 MUSIQ（70.49）指标上排名第二。尽管在像 PSNR 和 SSIM 这样的指标上表现不佳，但视觉质量结果表明，本方法在保持良好的视觉保真度方面表现出色。我们将在消融研究中解决这个问题。还对模型大小进行了比较。表格显示了 157M（SDXL）和 80M（SD3）的可训练参数，而 ControNet 的参数量大致是原网络的一半。使用 ControNet 作为 LQ 条件输入的方法，如 DiffBIR、SeeSR 和 SUPIR，当先验参数增加时，尤其对于 2B 的 StableDiffusion 3 来说，存在问题。相比之下，本方法仅使用 80M 参数就可实现良好的恢复。
定性比较
对于视觉质量比较，呈现了 RealPhoto60 和 DIV2K 验证集的 1024×1024 结果。从 DIV2K 验证集中生成 LQ 图像的方式为混合退化。如图 4 所示，本方法在保真度和质量之间表现出良好的平衡。与 SUPIR 和 SeeSR 等 SOTA 方法相比，本方法大多数情况下可产生良好的恢复，但有时它们会生成模糊的结果或保真度较低的结果。例如，SUPIR 为第 3 行的石头生成了错误的纹理，对猫的毛发进行了模糊处理。SeeSR 对第 1 行的面部生成细节较少，对第 2 行和第 3 行的图像进行了模糊处理。最后一行表明，本方法能对帽子和衣服生成精细的纹理，而其他方法难以做到。本方法的结果还具有良好的视觉保真度，尽管在 PSNR 和 SSIM 指标上表现不佳，这表明它们作为视觉保真度的鲁棒性指标不够。最后两列显示了本方法的结果，足够详细和锐利的边缘，没有明显的退化，同时保持视觉保真度。我们还在 DIV2K 验证集的随机裁剪图像上进行了实验。如图 5 所示，一些方法会产生模糊的结果或错误的纹理，而本方法始终能提供高质量的图像，并具有丰富的细节。
5 结论
设计了一个名为扩散恢复适配器的框架，由恢复适配器和扩散适配器组成，用于利用预训练扩散模型作为真实世界图像恢复的先验。该框架可与使用降噪 U-Net 或 DiT 的扩散模型一起工作。具体为恢复适配器设计了适用于两种扩散先验的架构。此外，还提出了恢复采样策略，以平衡不同采样器之间的质量和保真度。定量和定性实验已证明本方法的有效性。
END
来源：AIWalker
推荐阅读
- 英伟达提出首个 Mamba-Transformer 视觉骨干网络！打破精度/吞吐瓶颈
- 自回归视觉生成破局：Next-X 预测，开辟视觉新路径
- YOLOv12 来袭，更高性能、更快速度~
- 从文字到视觉：EvalMuse-40K 如何评价 T2I 模型的进化
- FoundIR: 释放百万规模高质量训练数据，助推图像复原基础模型
本文章著作权归作者所有，任何形式的转载都请注明出处。更多动态滤波，图像质量，超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
