拋棄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計算流的統一注意力控制和圖像塊合并機制。
應用的技術：
- Diffusion Transformers (DiT) 框架
- DPM-Solver反演算法
- 統一注意力控制和圖像塊合并機制
達到的效果：
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生成未來
