AIGC——BrushNet使用双分支扩散进行即插即用图像修复

1.概述

图像修复是计算机视觉中的经典问题之一，其目的是恢复图像中的遮蔽区域，使其具有合理且自然的内容。现有的工作采用传统的图像修复技术，如生成对抗网络或 GANS，以及变分自动编码器或 VAE，通常需要辅助的手工设计功能，但同时不能提供令人满意的结果。在过去的几年中，基于扩散的方法由于其卓越的高质量图像生成能力、输出多样性和细粒度控制而在计算机视觉界受到欢迎。最初尝试采用扩散模型进行文本引导图像修复，通过对预先训练的扩散模型中的遮蔽区域进行采样，并对给定图像中的未遮蔽区域进行采样，从而修改了标准去噪策略。尽管这些方法在简单的图像修复任务中取得了令人满意的性能，但它们在处理复杂的蒙版形状、文本提示和图像内容时遇到了困难，导致整体缺乏连贯性。这些方法中观察到的缺乏一致性主要是由于它们对掩模边界和未掩模图像区域上下文的感知知识有限。

尽管过去几年这些模型取得了进步、研究和开发，但图像修复仍然是计算机视觉开发人员的主要障碍。当前用于图像修复任务的扩散模型的适应涉及修改采样策略，或者特定于修复的扩散模型的开发经常遭受图像质量下降和语义不一致的困扰。为了应对这些挑战，并为图像修复模型铺平道路，在本文中，我们将讨论 BrushNet，这是一种新颖的即插即用双分支工程框架，可将像素级掩模图像特征嵌入到任何预先训练的扩散中模型，从而保证一致性，并增强图像修复任务的结果。 BrushNet 框架引入了一种新颖的范例，在该范例下，该框架将图像特征和潜在噪声划分为单独的分支。图像特征和噪声潜伏的划分极大地减少了模型的学习负担，并有助于以分层方式细致地合并基本的屏蔽图像信息。除了 BrushNet 框架之外，我们还将讨论 BrushBench 和 BrushData，它们分别促进基于分割的性能评估和图像修复训练。

2. BrushNet：使用双分支扩散进行图像修复

图像修复是一种试图在保持图像整体连贯性的同时恢复图像任务区域的方法，一直是计算机视觉领域长期存在的问题，多年来一直困扰着开发人员和研究人员。图像修复可应用于各种计算机视觉任务，包括图像编辑和虚拟试穿。最近，稳定扩散和稳定扩散 1.5 等扩散模型已表现出生成高质量图像的卓越能力，并且它们为用户提供了控制语义和结构控制的灵活性。扩散模型的巨大潜力促使研究人员采用扩散模型来执行与输入文本提示一致的高质量图像修复任务。

传统的基于扩散的文本引导修复框架采用的方法可以分为两类：采样策略修改和专用修复模型。采样策略修改方法通过从预先训练的扩散模型中采样屏蔽区域来修改标准去噪过程，并在每个去噪步骤中从给定图像复制粘贴未屏蔽区域。尽管采样策略修改方法可以在任意扩散模型中实现，但它们通常会导致不连贯的修复结果，因为它们对掩模边界和未掩模图像区域上下文的感知知识有限。另一方面，专用修复模型通过扩展基础扩散模型的输入通道的维度来微调专门设计的图像修复模型，以合并损坏的图像和掩模。虽然专用的修复模型使扩散模型能够通过专门的形状感知和内容感知模型生成更令人满意的结果，但它可能是也可能不是图像修复模型的最佳架构设计。

如下图所示，专用修复模型在早期阶段融合了蒙版图像潜伏、噪声潜伏、文本和蒙版。这种专用修复模型的架构设计很容易影响屏蔽图像特征，并且由于文本影响而导致UNet架构中的后续层无法获得纯粹的屏蔽图像特征。此外，在单个分支中处理生成和条件给 UNet 架构带来了进一步的负担，并且由于这些方法还需要在扩散主干的不同变体中进行微调，因此这些方法通常非常耗时且可转移性有限。

看起来，添加一个专门用于提取蒙版图像特征的附加分支可能是上述问题的充分解决方案，但是，现有框架在直接应用于修复时通常会导致提取和插入不充分的信息。因此，与专用修复模型相比，ControlNet 等现有框架的结果并不令人满意。为了以最有效的方式解决这个问题，BrushNet 框架在原始扩散网络中引入了一个额外的分支，从而为图像修复任务创建了一个更合适的架构。 BrushNet框架的设计和架构可以概括为三点。

1. BrushNet 框架不是随机初始化卷积层，而是实现了 VAE 编码器来处理屏蔽图像。因此，BrushNet 框架能够更有效地提取图像特征以适应 UNet 分布。
2. BrishNet 框架逐渐将完整的 UNet 特征逐层合并到预训练的 UNet 架构中，这是一种能够实现密集的每像素控制的分层方法。
3. BrushNet 框架从 UNet 组件中删除了文本交叉注意力，以确保在附加分支中考虑纯图像信息。此外，BrushNet 模型还建议实施模糊混合策略，以获得更好的一致性以及图像未遮蔽区域的更高范围的可控性。

3. BrushNet：方法和架构

下图为我们简单介绍了BrushNet框架。

可以看出，该框架采用双分支策略进行掩模图像引导插入，并使用模糊掩模的混合操作来确保更好地保留未掩模区域。值得注意的是，BrushNet框架能够调整添加的尺度以实现灵活的控制。对于给定的蒙版图像输入和蒙版，BrushNet 模型输出修复图像。该模型首先对掩模进行下采样以适应潜在空间的大小，然后将掩模图像作为 VAE 编码器的输入，以对齐潜在空间的分布。然后，该模型将屏蔽图像潜在图像、噪声潜在图像和下采样掩模连接起来，并将其用作输入。然后将模型提取的特征添加到零卷积块之后的预训练 UNet 层。去噪后，模型将蒙版图像和生成的图像与模糊蒙版混合。

3.1 蒙版图像引导

BrushNet 框架使用附加分支将屏蔽图像特征插入到预先训练的扩散网络中，从而将屏蔽图像的特征提取与图像生成过程明确分开。输入是通过连接掩模图像潜在图像、噪声潜在图像和下采样掩模形成的。更具体地说，噪声潜伏在当前生成过程中为图像生成提供信息，并帮助框架增强掩模图像特征的语义连贯性。然后，BrushNet 框架使用变分自动编码器从蒙版图像中提取隐藏的蒙版图像。此外，该框架采用三次插值对掩模进行下采样，以尝试确保掩模尺寸与掩模图像潜伏和噪声潜伏对齐。为了处理屏蔽图像特征，BrushNet 框架实现了预训练扩散模型的克隆，并排除了扩散模型的交叉注意层。原因是扩散模型的预训练权重作为提取掩模图像特征的强先验，并且排除交叉注意层确保模型仅考虑附加分支内的纯图像信息。 BrushNet框架将特征逐层插入到冻结扩散模型中，从而实现分层密集的逐像素控制，并且还采用零卷积层在可训练的BrushNet模型和锁定模型之间建立连接，确保有害噪声在初始训练阶段对可训练副本中的隐藏状态没有影响。

3.2 混合操作

如前所述，在潜在空间中进行混合操作会调整掩模的大小，这通常会导致一些不准确的情况，而 BrushNet 框架在调整掩模的大小以匹配潜在空间的大小时也会遇到类似的问题。此外，值得注意的是，变分自动编码器中的编码和解码操作具有固有的有限操作，并且可能无法确保完整的图像重建。为了确保框架重建未遮蔽区域的完全一致的图像，现有的工作已经实现了不同的技术，例如从原始图像复制未遮蔽区域。尽管该方法有效，但它通常会导致最终结果的生成缺乏语义一致性。另一方面，采用潜在混合操作等其他方法在保留未屏蔽区域中所需的信息方面面临困难。

3.3 灵活的控制

BrushNet 框架的架构设计使其成为各种预训练扩散模型固有的即插即用集成的合适选择，并实现灵活的保存规模。由于 BrishNet 框架不会改变预训练扩散模型的权重，因此开发人员可以灵活地将其作为即插即用组件与微调扩散模型集成，从而轻松采用预训练模型并进行实验。此外，开发人员还可以选择通过将 BrushNet 模型的特征合并到具有给定权重 w 的冻结扩散模型中来控制未遮蔽区域的保留规模，该权重决定了 BrushNet 框架对保留规模的影响，为开发人员提供了调整所需保存水平的能力。最后，BrushNet框架允许用户调整模糊比例，并决定是否执行模糊操作，从而轻松定制未遮罩区域的保留比例，为图像修复过程的灵活调整和细粒度控制留出了空间。

4.BrushNet：实施和结果

为了分析其结果，BrushNet 框架提出了 BrushBench，这是一个基于主体的图像修复数据集，包含 600 多个图像，每个图像都附有人工注释的掩模和标题注释。基准数据集中的图像在自然图像和人工图像之间均匀分布，也保证了不同类别之间的均匀分布，从而能够对不同类别进行公平评估。为了进一步增强修复任务的分析，BrushNet 框架根据所使用的方法将数据集分为两个不同的部分：基于分割的方法和画笔蒙版。

4.1 定量比较

下表将 BrushNet 框架与 BrushBench 数据集上现有的基于扩散的图像修复模型进行了比较，以稳定扩散作为基本模型。

可以看出，BrushNet 框架在屏蔽区域保留、文本对齐和图像质量方面表现出了卓越的效率。此外，稳定扩散修复、高清画家、PowerPaint 等模型在图像内部修复任务中表现出强大的性能，尽管它们无法复制其在外部修复任务中的性能，特别是在文本对齐和图像质量方面。总体而言，BrushNet 框架提供了最强的结果。

此外，下表将 BrushNet 框架与 EditBench 数据集上现有的基于扩散的图像修复模型进行了比较，其性能与 BrushBench 数据集上观察到的性能相当。结果表明 BrushNet 框架在具有不同掩模类型的各种图像修复任务中提供了强大的性能。

4.2 定性比较

下图对 BrushNet 框架与现有图像修复方法进行了定性比较，结果涵盖了不同修复任务中的人工智能和自然图像，包括随机掩模修复、分割掩模内部修复和分割掩模外部修复。

可以看出，BrushNet框架在未掩模区域和相干区域的一致性方面取得了显着的效果，并且由于双分支解耦方法的实施，成功实现了背景信息的感知。此外，预训练扩散模型的未触及分支还提供了更好地覆盖动漫和绘画等不同数据域的优势，从而在不同场景下获得更好的性能。

5.最后的思考

在本文中，我们讨论了 BrushNet，这是一种新颖的即插即用双分支工程框架，它将像素级掩模图像特征嵌入到任何预先训练的扩散模型中，从而保证一致性并增强图像修复任务的结果。 BrushNet 框架引入了一种新颖的范例，在该范例下，该框架将图像特征和潜在噪声划分为单独的分支。图像特征和噪声潜伏的划分极大地减少了模型的学习负担，并有助于以分层方式细致地合并基本的屏蔽图像信息。除了 BrushNet 框架之外，我们还将讨论 BrushBench 和 BrushData，它们分别促进基于分割的性能评估和图像修复训练。