Practicing Model Scaling for Photo-Realistic Image Restoration In the Wild（SUPIR-CVPR2024）-平芜编程栈

Paper：https://arxiv.org/pdf/2401.13627
Code：https://github.com/Fanghua-Yu/SUPIR
Baseline：StableDiffusion-XL (SDXL)

文章目录

前言

一、整体框架

3.1. Model Scaling Up

3.2. Scaling Up Training Data

3.3. Restoration-Guided Sampling

二、实验

数据集

实验设置

三、实验结果

4.2 Comparison with Existing Methods

4.3 Controlling Restoration with Textual Prompts

4.4 Ablation Study

前言

问题：Image Restoration：然而，这些方法（单任务）通常基于特定的退化假设[25,50,56]，因此缺乏对其他退化的泛化能力；

Generative Prior：然而，这些基于扩散的红外方法的性能受到所使用的生成模型的规模的限制，这对进一步提高其有效性提出了挑战。

Model Scaling：然而，扩大规模是一个系统问题，涉及模型设计、数据收集、计算资源和其他限制。许多其他任务尚未能够享受到扩展带来的实质性性能改进。IR就是其中之一。

动机：不断提高生成先验的能力是获得更好的 IR 结果的关键，模型缩放是一个关键和有效的方法。有许多任务已经从缩放中获得了惊人的改进，例如SAM 和大型语言模型(large language models)。这进一步推动了我们对构建能够产生超高质量图像的大规模智能 IR 模型的追求。

然而，由于计算资源、模型架构、训练数据以及生成模型和IR的合作等工程限制，扩大IR模型是一项挑战。

贡献：在这项工作中，我们引入了SUPIR，这是迄今为止最大的 IR 方法，旨在探索在恢复视觉效果和智能方面的更大潜力。

具体来说，SUPIR使用StableDiffusion-XL (SDXL)作为强大的生成先验，它包含26亿个参数。

为了有效地在IR中部署该模型，我们设计并训练了一个large-scale adaptor，该适配器包含一个名为ZeroSFT connector的新组件。

为了最大化模型缩放的好处，我们收集了超过2000万张高质量、高分辨率图像的数据集，每张图像都附有详细的描述性文本。

我们利用130亿个参数的多模态语言模型来提供图像内容提示，大大提高了我们方法的准确性和智能性。

所提出的SUPIR模型在各种IR任务中表现出卓越的性能，实现了最佳的视觉质量，特别是在复杂和具有挑战性的现实世界场景中。此外，该模型通过文本提示提供对恢复过程的灵活控制，极大地扩展了IR的可能性。图1说明了我们模型的影响。

一、整体框架

3.1. Model Scaling Up

（1）Generative Prior

我们选择SDXL的原因如下。

Imagen 和 IF 优先考虑文本到图像的生成，并依赖于分层方法。他们首先生成小分辨率的图像，然后分层次地对它们进行采样。

SDXL 与我们的目标一致，直接生成高分辨率的图像，没有分层设计，有效地使用其参数来提高图像质量，而不是专注于文本解释。

此外，SDXL 采用 Base-Refine 策略。在 Base 模型中，会生成各种不同但质量较低的图像。随后，Refine 模型利用的训练图像质量明显高于 Base 模型，但多样性却低于 Base 模型，从而提高了图像的质量。

考虑到我们使用大量高质量图像数据集进行训练的方法，SDXL 的双阶段设计对于我们的目标来说变得多余。我们选择 Base 模型，它有更多的参数，使其成为理想的生成先验。

（2）Degradation-Robust Encoder

在SDXL中，扩散生成过程在潜在空间中进行。首先通过预训练的编码器将图像映射到潜在空间。为了有效地利用预训练的SDXL，我们的 LQ 图像也应该映射到相同的潜在空间。但是，由于原始编码器没有经过 LQ 图像的训练，使用它进行编码会影响模型对LQ图像内容的判断，从而将伪影误解为图像内容[49-DiffBIR]。

为此，我们对编码器进行微调，使其对退化具有鲁棒的，方法是最小化：

其中是要微调的退化鲁棒编码器，是固定解码器，是 ground truth。

（3）Large-Scale Adaptor Design

这部分就是如何让一个 26 亿参数的 SDXL，在“不被破坏、不失控、不炸显存”的前提下，听懂低质量图像，并进行像素级可控的图像复原。但是：

SDXL 很大（2.6B 参数）
IR 需要 pixel-level control（不是随便生成）
LQ 图像 ≠ 文本 prompt
直接套 ControlNet / LoRA 都不合适

所以必须重新设计一个适配器（Adaptor）。

LoRA[32]、T2I适配器[57] 和 ControlNet[92] 是现有的扩散模型自适应方法，但它们都不满足我们的要求：LoRA限制了生成（只能调风格/语义，不能精确对齐 LQ 图像；且本质是低秩权重扰动）; T2I缺乏LQ图像内容识别能力（倾向于从外部条件引导生成（边缘、深度、姿态））; 而ControlNet的直接复制对于SDXL模型规模来说是一个挑战（太大，直接复制 SDXL encoder 根本训练不动）。

为了解决这个问题，我们设计了一个具有两个关键特性的新适配器，如图 3 图(a)所示。

Trimmed ControlNet：保留 ControlNet 思想，但砍掉一半 encoder（只保留 N/2 个ViT block）。

还能 work 的原因：砍掉一半 block，只要结构仍像原 encoder，初始化来自 SDXL，那么条件信息仍然能对齐到 SDXL latent space。

ZeroSFT Connector：ControlNet 的 zero-conv 不够，需要更强的像素级控制。

ZeroSFT = ZeroConv + SFT + GroupNorm；不仅加一点条件残差，而是用条件特征直接调制主干特征分布。输入三路特征：：主干 decoder feature（fixed）；：encoder shortcut feature；：来自 adaptor（LQ guidance）。

3.2. Scaling Up Training Data

（1）Image Collection

模型的缩放需要对训练数据[38-Scaling Laws for Neural Language Models]进行相应的缩放。但是目前还没有大规模的高质量红外图像数据集。虽然DIV2K[3]和LSDIR[1]提供高图像质量，但它们的数量有限。ImageNet (IN)[17]、LAION-5B[67]、SA-1B[44]等较大的数据集包含的图像较多，但其图像质量不符合我们的高标准。

为此，我们收集了一个大规模的高分辨率图像数据集，其中包括 2000 万张 1024×1024 高质量、纹理丰富的图像。收集到的数据集与现有数据集的尺度对比如图 3 所示。我们还从 FFHQ-raw 数据集[40]中额外加入了 70K 未对齐的高分辨率面部图像，以提高模型的面部恢复性能。在图 5(a) 中，我们显示了与其他知名数据集相比我们数据的相对大小。

（2）Multi-Modality Language Guidance

我们认为文本提示也可以帮助 IR 识别：

(1) 理解图像内容对 IR 至关重要。现有框架往往忽略或隐式处理这种理解[24,29]。通过结合文本提示，我们明确地将对 LQ 图像的理解传达给 IR 模型，促进有针对性地恢复缺失信息。

(2) 在严重退化的情况下，即使是最好的 IR 模型也很难完全恢复丢失的信息。在这种情况下，文本提示可以作为一种控制机制，根据用户偏好有针对性地完成缺失的信息。

(3) 我们还可以通过文本描述期望的图像质量，进一步增强输出的感知质量。一些例子见图 1(b)。

为此，我们做了两个主要的修改：

首先，我们修改了整体框架，将LLaVA多模态大型语言模型[51]纳入我们的 pipeline，如图 2 所示。LLaVA 将经过退化鲁棒处理的 LQ 图像

作为输入，显式地理解图像中的内容，以文本描述的形式输出。然后使用这些描述作为提示来指导恢复。该过程可以在测试期间自动化，从而消除了人工干预的需要。

其次，我们按照PixART[12]的方法，对所有的训练图像进行文本标注，加强文本控制在我们模型训练过程中的作用。

（3）Negative-Quality Samples and Prompt

无分类器引导(CFG)[30]提供了另一种控制方式，通过使用负面提示为模型指定不需要的内容。我们可以使用这个特征来指定模型不产生低质量的图像。

具体来说，在扩散的每一步，我们将使用正提示pos 和负提示neg 进行两次预测，并将这两个结果的融合作为最终输出:

其中，为带适配器的扩散模型，为时间步长为 t 的噪声方差，为超参数。

然而，在我们的训练数据中缺少负质量的样本和提示可能导致微调后的 SUPIR 在理解负提示方面失败。因此，在采样期间使用负质量提示可能会引入伪影，示例参见图 4。

为了解决这个问题，我们使用SDXL生成了100K的图像，对应于低质量的提示。我们反直觉地将这些低质量图像添加到训练数据中，以确保所提出的SUPIR模型可以学习到负质量概念。

3.3. Restoration-Guided Sampling

强 generative prior 是双刃剑：生成能力太强，会破坏复原的忠实性（fidelity）。

如何在采样阶段，把 diffusion 拉回到 LQ 图像附近，而不是重新训练模型？

为了解决这一问题，我们对 EDM 采样方法[41]进行了改进，提出了一种恢复引导采样方法。我们希望在每个扩散步骤中有选择地引导预测结果接近 LQ 图像。我们在预测输出和LQ潜在之间进行加权插值，作为恢复引导输出。

本文利用了 diffusion 一个已知但常被忽略的性质：diffusion 的时间语义：早期（大，大，大），主要生成低频 / 结构；后期（小，小，小），主要生成高频 / 纹理。

RGS：早期：强约束，拉近 LQ，保结构；后期，弱约束，放手生成，保细节。

的物理意义：当大，预测被强行拉向 LQ latent；当小，几乎不干预模型预测。

= fidelity <-> realism 的旋钮，

小：始终较大，输出 ≈ LQ → 保真但糊
大：约束很弱，输出 ≈ generative prior → 锐但可能假

实验选择= 4，是一个折中点。

二、实验

数据集

对于训练，整体训练数据包括2000万张带有文字描述的高质量图像、70K张人脸图像和100K张负质量样本，以及相应的提示符。为了实现更大的批处理大小，我们在训练期间将图像裁剪成512×512补丁。我们使用合成退化模型训练我们的模型，遵循Real-ESRGAN[78]使用的设置，唯一的区别是我们将生成的LQ图像的大小调整为512×512进行训练。

实验设置

我们使用学习率为0.00001的AdamW优化器[53]。训练过程持续10天，在64个Nvidia A6000 GPU 上进行，批量大小为256个。对于测试，超参数T=100，=7.5，=4。我们的方法能够处理大小为1024×1024的图像。我们将输入图像的短边调整为1024，并裁剪1024×1024子图像进行测试，然后在恢复后将其调整回原始大小。除非另有说明，否则不会手动提供提示-处理将完全自动进行。