BrushNet：一种具有分解双分支扩散的即插即用图像修复模型
BrushNet：一种具有分解双分支扩散的即插即用图像修复模型
BrushNet是一种新型的即插即用图像修复模型，通过引入分解双分支扩散架构，显著提升了图像修复效果。该模型在多个基准测试中表现出色，能够生成高质量的修复结果，并保持良好的文本对齐性。
BrushNet：一种具有分解双分支扩散的即插即用图像修复模型
摘要
随着扩散模型（DM）的出现，图像修复在恢复损坏图像方面取得了重大进展。然而，现有的修复模型经常出现语义不一致和图像质量下降的问题。为了解决这些问题，BrushNet提出了一种新的范式：将蒙版图像特征和噪声潜伏物划分为单独的分支。这种设计大大减少了模型的学习负担，促进了以分层方式细致入微地合并基本的掩蔽图像信息。
BrushNet是一种新颖的即插即用双分支模型，旨在将像素级掩码图像特征嵌入到任何预训练的DM中，从而保证连贯和增强的图像修复结果。此外，研究团队还引入了BrushData和BrushBench，以促进基于分割的修复训练和性能评估。广泛的实验分析表明，BrushNet在图像质量、蒙版区域保留和文本连贯性等多个关键指标上优于现有模型。
1. 引言
图像修复旨在恢复图像的缺失区域，同时保持整体连贯性。作为一个长期存在的计算机视觉问题，它在虚拟试穿和图像编辑等领域有着广泛的应用。近年来，扩散模型在图像生成方面的突破性进展为图像修复带来了新的可能性。然而，现有的基于扩散的修复方法存在一些局限性，如语义不一致和图像质量下降等问题。
为了解决这些问题，BrushNet提出了一种创新的双分支架构。该模型通过将蒙版图像特征和噪声潜伏物分离到不同的分支中，显著提升了修复效果。此外，BrushNet还引入了BrushData和BrushBench，以促进基于分割的修复训练和性能评估。
2. 相关工作
图像修复是计算机视觉中的一个经典问题，旨在恢复图像的遮罩区域，使其具有合理和自然的内容。早期基于传统技术的方法、变分自动编码器（VAE）和生成对抗网络（GAN）通常需要辅助的手工设计功能，但效果不佳。近年来，基于扩散的方法由于其高质量生成、细粒度控制和输出多样性而广受欢迎。
BrushNet与以往的修复方法相比具有以下优势：
- 即插即用：可以轻松集成到任何预训练扩散模型中
- 内容感知：能够理解图像内容并生成连贯的修复结果
- 形状感知：能够处理复杂的掩码形状
- 灵活的未遮罩区域保留：可以通过调整参数控制未遮罩区域的保留程度
3. 预备和动机
本节将介绍扩散模型的基本原理、现有的修复方法及其局限性，以及BrushNet的设计动机。
3.1 扩散模型
略
3.2 以前的修复模型
- 抽样策略修改：这类方法通过逐渐将蒙版图像与生成的结果混合来实现修复。虽然实现简单，但在未屏蔽区域保留和生成内容对齐方面表现不佳。
- 专用修复模型：通过扩展UNet的输入通道以包含蒙版和蒙版图像输入，这些方法虽然生成结果更好，但仍存在以下问题：
- 在UNet架构的初始卷积层合并噪声潜伏层、掩蔽图像潜伏层和掩码时，后续层难以获得纯遮罩图像特征
- 将条件处理和生成合并到一个分支中会给UNet框架带来额外负担
- 需要对扩散主干的各种变体进行广泛的微调，计算密集且缺乏可转移性
3.3 设计动机
ControlNet虽然在可控图像生成方面表现出色，但其设计并不适合图像修复任务。ControlNet依赖稀疏控制和文本引导，而修复需要像素级的强约束。因此，BrushNet专门针对修复任务设计了一个新的架构。
4. 方法
BrushNet的概述如图3所示。该模型采用双分支策略进行掩码图像引导插入，使用模糊掩码进行混合操作，并支持灵活的控制。
图3：模型概述
BrushNet的输入包括噪声潜伏、掩码图像潜伏和下采样掩码。模型通过零卷积块将BrushNet特征逐层添加到预训练的UNet中。去噪后，生成的图像和掩码图像与模糊掩码混合。
4.1 掩码图像引导
BrushNet通过一个附加分支处理掩码图像特征，该分支将掩码图像的特征提取与图像生成过程显式分离。附加分支的输入包括噪声潜伏、掩码图像潜伏和下采样掩码。使用VAE从掩码图像中提取潜伏特征，以确保与预训练UNet的数据分布一致。BrushNet特征逐层插入到冻结的扩散模型中，实现密集的每像素分层控制。
4.2 混合操作
为了确保对未遮罩区域进行完全一致的图像重建，BrushNet采用模糊掩码进行混合操作。这种方法虽然可能会导致在保留掩模边界细节方面的准确性略有下降，但肉眼几乎察觉不到这种误差，并导致掩模边界的相干性显著提高。
4.3 灵活控制
BrushNet的架构设计支持与各种预训练扩散模型的无缝集成，并实现灵活的保存尺度。具体来说，BrushNet的灵活控制包括：
- 作为即插即用组件与任何社区微调扩散模型集成
- 通过调整BrushNet特征的权重来控制未遮罩区域的保存尺度
- 通过调整模糊比例和混合操作来进一步自定义未遮罩区域的保存比例
5. 实验
本节将介绍BrushNet的实验设置、评估基准和指标，以及与其他模型的比较结果。
5.1 评估基准和指标
研究团队提出了两个新的基准测试：BrushBench和EditBench。BrushBench专注于基于分割的修复，包含600张图像，每张图像都附有人工注释的掩码和标题注释。EditBench则用于随机画笔掩码的修复任务。
评估指标包括：
- 图像生成质量：使用ImageReward、HPS v2和审美评分（AS）
- 掩码区域保留：使用PSNR、LPIPS和MSE
- 文本对齐：使用CLIP相似度
5.2 实施细节
实验在NVIDIA Tesla V100上进行，使用Stable Diffusion v1.5作为基础模型。BrushNet在8个GPU上训练430,000步，大约需要3天。详细信息可以在提供的代码中找到。
5.3 定量比较
表2和表3显示了BrushNet与其他基于扩散的修复模型在BrushBench和EditBench上的定量比较结果。BrushNet在图像质量、掩码区域保留和文本对齐等多个指标上均表现出色。
表2：BrushNet与其他基于扩散的修复模型在BrushBench中的定量比较
表3：BrushNet和其他基于扩散的修复模型在EditBench上的定量比较
5.4 定性比较
与以前的图像修复方法的定性比较如图1所示。BrushNet在各种修复任务中始终显示出出色的结果，特别是在生成区域和未遮罩区域的连贯性方面。例如，在图1的第III组中，BrushNet成功识别了掩码图像中已经存在的金鱼，避免了重复生成其他鱼类。
图1：BrushNet与其他图像修复方法在各种修复任务中的性能比较
5.5 灵活的控制能力
图6展示了BrushNet与社区微调的不同扩散模型的结合能力。图7演示了BrushNet控制比例的调整，通过调整比例参数，用户可以对修复过程进行精细控制。
图6：将BrushNet集成到社区微调的扩散模型中
图7：BrushNet灵活的控制比例
5.6 消融研究
表4和表5显示了BrushNet的关键设计选择的消融研究结果。结果表明，双分支设计、使用VAE编码器处理掩码图像、逐层合并UNet特征以及去除文本交叉注意力等设计选择对模型性能有显著提升。
表4：双分支设计消融
表5：模型架构上的消融
6. 讨论
BrushNet通过引入分解双分支扩散架构，显著提升了图像修复效果。然而，该模型仍存在一些局限性，如生成质量受基础模型影响、对异常形状的掩码处理效果不佳等。未来的研究方向包括改进模型的泛化能力、处理更复杂的掩码形状等。
负面的社会影响：图像修复模型为内容创作提供了令人兴奋的机会，但它们也给个人和社会带来了潜在的风险。它们对互联网收集的训练数据的依赖会放大社会偏见，并且存在通过操纵具有冒犯性元素的人类图像来产生有说服力的错误信息的特定风险。为了解决这些问题，在使用这些模型时强调负责任的使用并建立道德准则至关重要。这也是我们未来模型发布的重点。
本文原文来自csdn