抛弃UNet，首个基于DiT的图像编辑框架！DiT4Edit：多尺寸编辑质量更优 | 北大&港科大

文章链接：https://arxiv.org/pdf/2411.03286
工程链接：https://github.com/fkyyyy/DiT4Edit（待开源）

亮点直击

• 基于基于Transformer的扩散模型在图像编辑中的优势，提出了DiT4Edit，第一个不需要调优的基于扩散Transformer（DiT）的图像编辑框架。
• 为了适应基于Transformer的去噪计算机制，首先提出了统一的注意力控制机制，以实现图像编辑。引入了DPM-Solver反演和图像块合并策略，以减少推理时间。
• 大量定性和定量的实验结果证明了DiT4Edit在对象编辑、风格编辑和形状感知编辑方面的优越性能，适用于各种图像尺寸，包括512 × 512、1024 × 1024和1024 × 2048。

总结速览

解决的问题：
现有的基于UNet的高分辨率图像编辑方法在实现形状感知的对象编辑方面仍存在不足。相比之下，Diffusion Transformers (DiT) 能更有效地捕捉图像块之间的长程依赖关系，但目前在图像编辑领域的应用仍较少。

提出的方案：
本文提出了基于Diffusion Transformers的首个图像编辑框架——DiT4Edit。DiT4Edit采用DPM-Solver反演算法来获取反演隐变量，与传统UNet框架常用的DDIM反演算法相比，显著减少了迭代步数。此外，设计了适用于Transformer计算流的统一注意力控制和图像块合并机制。

应用的技术：

1. Diffusion Transformers (DiT) 框架
2. DPM-Solver反演算法
3. 统一注意力控制和图像块合并机制

达到的效果：
DiT4Edit在编辑质量和速度上优于UNet框架，尤其在高分辨率和任意尺寸的图像编辑中表现突出。大量实验验证了该框架在各种编辑场景中的强大性能，展示了Diffusion Transformers在支持图像编辑方面的潜力。

方法

基础知识：隐空间扩散模型

扩散模型架构

选择Transformer作为去噪模型的原因
与UNet结构相比，Transformer引入了全局注意力机制，使得模型能够关注图像中的更广泛区域。这种增强的可扩展性使得Transformer能够生成大尺寸（例如大于512×512）甚至任意尺寸的高质量图像。本文基于DiT的编辑框架在大尺寸图像上的编辑结果在下图1和图2中有所展示，这些是UNet框架之前未涉及的编辑任务。因此，采用了基于Transformer的去噪模型作为编辑框架，利用Transformer的能力来应对这些更复杂的编辑挑战。

基于扩散Transformer的图像编辑

如上图2所示，基于预训练的扩散Transformer图像编辑框架的流程。

最近的研究（Hong等，2024）提出了一种通过反向欧拉方法来获取方程8中高阶项近似的策略。

统一控制注意力机制：在之前的研究Prompt to Prompt (P2P)中，研究人员展示了交叉注意力层包含来自提示文本的丰富语义信息。这一发现可以通过在扩散过程中替换源图像与目标图像之间的交叉注意力图来编辑图像。具体而言，常用的两种基于文本引导的交叉注意力控制策略是交叉注意力替换和交叉注意力细化。这两种方法确保了从目标提示到源提示的信息无缝流动，从而将隐空间变量图引导到期望的方向。

Patches Merging：为了提高推理速度，借鉴了Token Merging的方法，将patch合并技术嵌入到去噪模型中。这个方法的灵感来自于这样一个观察：在transformer架构中，涉及到的patch数量明显大于UNet。计算流程如下图4所示。对于一个特征图，我们首先计算每个patch之间的相似度，然后将最相似的patch合并，从而减少由注意力机制处理的patch数量。经过注意力计算后，我们会将patch拆分回来，以保持模型下一层的输入大小。通过将patch合并技术融入我们的框架中，我们旨在简化处理过程，提高整体效率，同时不改变每一层的基本操作。

实验

实现细节

研究者们评估了DiT4Edit编辑框架与六个先前基准模型（包括Pix2Pix-Zero、PnPInversion、SDEdit、IP2P、MasaCtrl和InfEdit）的质量表现差异，所有方法均使用官方开源代码实现。

如下图5所示，在512×512和1024×1024图像上对比了本文的方案。图5的第一行展示了本文的框架在编辑真实512×512图像时，能够生成与原始内容一致的编辑图像，而现有方法通常会改变原始图像的背景或目标细节。此外，图5的第二行和第三行展示了在大规模图像和任意尺寸图像上的实验——这些任务是以前基于UNet的方法难以处理的。结果表明，本文提出的框架有效地处理了大图像中的风格和物体形状修改。相比之下，一些先进的基于UNet的方法，尽管能够进行编辑任务，但通常会导致背景和物体位置的显著变化和损坏。此外，由于UNet结构的局限性，这些方法通常只能生成512×512大小的目标图像。这些发现强调了基于transformer的扩散模型在大规模图像编辑中的巨大潜力。

定量比较

为了进行定量评估，使用了三个指标：Fréchet Inception Distance (FID)、峰值信噪比（PSNR）和CLIP，以评估我们模型与SOTA（最先进方法）在图像生成质量、背景保持和文本对齐方面的性能差异。本文比较了三种尺寸的图像：512×512、1024×1024 和1024×2048，结果详细列在下表1中。与Pix2Pix-Zero、PnPInversion、SDEdit、IP2P、MasaCtrl和InfEdit进行了性能比较。需要注意的是，由于此前没有基于DiT的编辑框架，所有对比基准都基于UNet架构。实验结果表明，DiT4Edit编辑策略在图像生成质量、背景保持和文本对齐方面优于SOTA方法。由于集成的transformer结构具备全局注意力能力，DiT4Edit框架在各种大小的编辑任务中表现出了强大的鲁棒性。生成的图像不仅显示出更高的质量，而且在背景和细节控制方面提供了更好的控制，保持了与原始图像的一致性。特别是在编辑大规模或任意尺寸的图像时，DiT4Edit相较于其他方法展示了显著的优势，充分展示了transformer架构强大的可扩展能力。同时，我们的编辑框架具有更短的推理时间，与无反演编辑方法（InfEdit）相当。

消融实验

本文进行了系列消融研究，旨在展示DPM-Solver反演和patch合并的有效性。关于patch合并的消融实验结果如前面图4和下表2所示。实现patch合并后，尽管保持了与未使用patch合并时相当的编辑质量，但大尺寸图像的编辑时间显著减少。这表明，patch合并可以显著提升图像编辑框架的整体性能。此外，DPM-Solver和DDIM的消融实验结果如图7所示。在相同推理步数（T = 30）下比较两种方法时，DPM-Solver在图像编辑质量方面始终优于DDIM。这表明本文采用的DPM-Solver反演策略能够生成更优的潜在图像，从而在较少的步中获得更好的编辑效果。

结论

本文提出了DiT4Edit，这是首个基于扩散transformer的图像编辑框架。与之前的UNet框架不同，DiT4Edit提供了更优的编辑质量，并支持多种尺寸的图像。通过利用DPM Solver反演、统一的注意力控制机制和patch合并，DiT4Edit在512×512和1024×1024尺寸图像的编辑任务中超越了UNet结构。特别是，DiT4Edit能够处理任意尺寸的图像，如1024×2048，展示了transformer在全局注意力和可扩展性方面的优势。本文的研究为基于DiT的图像编辑奠定了基础，并有助于进一步探索transformer结构在生成式AI中的潜力。

局限性
在实验中，观察到T5-tokenizer偶尔会遇到词汇分割问题，这可能导致最终编辑过程中的失败。此外，我们的模型与原始图像相比，可能会出现颜色不一致的情况。更多的编辑失败案例请参考补充材料。

本文转自AI生成未来 ，作者：AI生成未来

原文链接:​​https://mp.weixin.qq.com/s/T71qD02qJogrBa2arhTzjQ​​