文章鏈接：https://arxiv.org/pdf/2409.15278
github鏈接：https://github.com/AFeng-x/PixWizard

亮點直擊

• 任務統一：針對視覺任務的多樣性，本項目基于生成框架將多種任務轉化為圖像到圖像的翻譯問題，并通過后處理將生成的可視化效果轉化為所需格式，從而簡化表示形式的挑戰。
• 數據構建：構建了一個包含3000萬條數據的全面訓練集，主要支持圖像生成、編輯、修復、定位和密集預測等五大功能，旨在整合視覺領域的任務和數據多樣性。
• 架構設計：采用基于流的Diffusion Transformer（DiT）作為基礎模型，強調其靈活性和穩定性。通過動態分區和填充方案增強模型對不同分辨率圖像的處理能力，并引入結構感知和語義感知指導，以支持多模態指令的實現。

總結速覽

解決的問題

統一多個視覺任務，包括圖像生成、圖像翻譯、圖像編輯和圖像修復等，并通過自由語言指令完成各類任務。

提出的方案

• 設計了一個多功能的圖像到圖像視覺助手PixWizard，整合各種視覺任務到一個統一的（圖像-指令）到圖像的生成框架。
• 構建了一個全面的Omni Pixel-to-Pixel Instruction-Tuning Dataset，采用詳細的自然語言指令模板進行訓練。

應用的技術

• 基于Diffusion Transformers (DiT)作為基礎模型，擴展其功能以支持靈活的任意分辨率機制。
• 引入結構感知和語義感知的指導，以有效融合輸入圖像的信息。

達到的效果

• PixWizard展示了在多種分辨率圖像上的卓越生成，編輯和翻譯能力。
• 具備良好的泛化能力，能夠處理未見過的任務和人類指令。

Omni Pixel-to-Pixel Instruction-Tuning數據集

為了使我們的視覺助手具備全面的圖像生成、編輯和翻譯功能，我們構建了一個多任務、多模態的視覺指令調優訓練數據集。該數據集涵蓋七個主要領域，共計3000萬條數據。根我們所知，這個數據集是最具多樣性且易于使用的圖像-指令-圖像三元組數據集。它由開源和內部數據集構成，并在MLLM和人工審查的協助下進行了過濾，如圖1所示：

PIXWIZARD

PixWizard是一個多功能的圖像到圖像模型。本節將從模型架構（如下圖2所示）和訓練策略的角度介紹PixWizard框架的細節。

基于流的條件指令調優

架構

文本編碼器：首先，我們在PixWizard中使用Gemma-2B作為文本嵌入器，對文本提示進行編碼。然而，在多任務學習中，僅依賴文本指令不足以準確引導模型執行用戶命令。為更有效地指導生成過程，我們進一步結合了CLIP文本編碼器。通過對CLIP文本嵌入應用全局平均池化，我們獲得了粗粒度的文本表示，然后通過多層感知機（MLP）生成任務嵌入（embedding）。該嵌入隨后通過調制機制與時間步嵌入相加，集成到PixWizard Block中。如圖3所示，這種方法自適應地在潛在空間中聚類相似的任務指令，同時將不同任務的指令分離，從而幫助模型生成過程朝正確的任務方向發展。

結構感知指導

為了有效捕捉輸入圖像條件的整體結構特征，我們首先使用來自SDXL的變分自編碼器（VAE）對圖像進行編碼。接下來，我們將圖像潛在向量與噪聲潛在向量在通道維度上連接，并向Patch Embeder添加額外的輸入通道，這些新通道的權重最初設置為零。

語義感知指導

任意分辨率

PixWizard繼承了（Zhuo et al., 2024）提出的動態分區和填充方案，使模型能夠在微調和推理過程中處理任意分辨率和縱橫比的圖像。然而，在實踐中，不同任務所需的分辨率可能有顯著差異。為了支持更靈活的任意分辨率處理，并盡可能保留每個圖像的原始分辨率，我們使用[5122, 7682, 10242]作為分辨率中心，以生成一組候選補丁分區。在訓練過程中，我們將具有相似分辨率的數據項分組到同一桶中，以確保每批次內最長和最短序列的長度差不會過大，因此來實現最小化填充token的使用，提高訓練效率。在推理過程中，我們還可以通過結合NTK感知縮放RoPE和夾心歸一化，來達到卓越的分辨率外推能力。

兩階段訓練和數據平衡策略

為了釋放模型的潛力并提高其在數據量較少的任務上的性能，我們提出了一種兩階段訓練和數據平衡策略。

階段1：在第一階段，我們通過將預訓練的文本到圖像模型的權重與隨機初始化的新模塊權重結合來初始化模型的權重。首先我們選擇數據集較小的任務，并為每個數據集分配一個采樣權重，以增加其數據量。該權重決定了在單個周期內數據集的重復次數。通過這種方法，每個任務大約有20k個數據點。然后，我們隨機選擇來自其他任務的訓練樣本以匹配此規模。該訓練過程持續4個周期。

階段2：在第二階段，我們使用第一階段獲得的權重初始化模型，并將所有收集的數據組合進行進一步訓練。為了平衡任務，我們手動為每個數據集分配采樣權重，如果權重小于1.0，則隨機選擇數據子集。我們還以1:1的比例將文本到圖像訓練數據包含在內，以構成我們的第二階段訓練數據集。在此階段，總訓練數據量達到2000萬個樣本。

實驗

第一部分結果

設置

對于圖像修復，遵循之前的研究（Conde et al., 2024; Potlapalli et al., 2024），我們在訓練期間準備各種修復任務的數據集。對于評估，首先選擇兩個代表性基準：Rain100L用于去雨，SIDD用于去噪。此外，論文還進一步評估了其他修復任務的性能并檢驗zero-shot的能力（具體請看論文的Appendix）。

對于圖像定位，我們在gRefCOCO、RefCOCO和RefCOCO+的驗證和測試集上評估引用分割任務。我們使用累積IoU (cIoU)作為性能指標。

密集圖像預測任務評估三項視覺任務：ADE20k用于語義分割，NYUv2和SUNRGB-D用于單目深度估計，以及NYU-Depth v2用于表面法線估計。對于語義分割，我們通過識別最近鄰的RGB顏色值來分配標簽，并使用平均交并比 (mIoU)指標評估準確性。對于單目深度估計，我們對輸出圖像在三個通道上進行平均，并應用訓練期間使用的線性變換層，獲得范圍在[0,10]米內的深度估計，深度估計使用均方根誤差 (RMSE)進行評估。對于表面法線估計，我們從輸出圖像中恢復相應的法向量，并使用平均角誤差來評估準確性。

結果

表1展示了與近期最先進的任務特定和一體化方法的全面性能比較。結果顯示，盡管去噪和去雨數據僅占整體訓練集的一小部分，我們的方法在統一方法中表現優異，甚至超越了一些任務特定的方法。在圖像定位任務中，PixWizard在RefCOCO（驗證集）上比基于擴散的通用模型InstructDiffusion高出4.8 cIoU。然而，與其他高度專業化模型相比，仍有改進空間。此外，如圖6所示，PixWizard支持靈活的指令，能夠不僅直接在圖像上突出和可視化目標對象，還生成相應的二進制掩膜。這突顯了其在現實世界互動和實際應用中的強大性能。

對于密集預測任務,在所有三個任務中，PixWizard在與通用基線和任務特定基線的競爭中表現出色。在NYUv2測試集上的深度估計中，PixWizard在RMSE上相比Unified-IO提高了10.0%，并且與Painter和InstructCV的表現相似。在語義分割中，PixWizard在mIoU上超越Unified-IO，提升了11.05分，盡管仍落后于其他方法。此外，圖7展示了PixWizard的輸出示例。通過為同一圖像提供相應的任務特定提示，我們可以輕松生成相應的條件可視化，突顯了PixWizard的重要實用價值。

第二部分結果（圖像編輯）

設置

我們在MagicBrush測試（Zhang et al., 2024a）和Emu Edit測試（Sheynin et al., 2024）兩個基準上評估PixWizard，以評估其圖像編輯能力的有效性。為了公平比較，我們主要與以指令引導的圖像編輯方法進行對比，包括InstructPix2Pix、MagicBrush、Emu Edit和UltraEdit。與Emu Edit一致，我們使用L1距離、CLIP圖像相似度、DINO相似度、CLIP文本-圖像相似度和CLIP文本-圖像方向相似度作為指標。

結果

圖像編輯

表2展示了我們與基線的結果。結果表明，我們的模型在自動化量化指標上始終超越InstructPix2Pix、MagicBrush和UltraEdit，并且在性能上與最先進的方法Emu Edit相當。圖8提供了定性比較。我們的模型精準識別編輯區域，同時保留其他像素，展現了對給定指令的最佳理解。

第三部分結果（圖像生成）

設置

本節重點評估PixWizard的生成能力，具體任務包括經典的文本到圖像生成、可控圖像生成、圖像inpainting和圖像外推outpainting。在可控圖像生成中，我們評估PixWizard基于特定條件（Canny邊緣圖和深度圖）生成圖像的能力。我們通過比較輸入條件與生成圖像中提取的相應特征之間的相似度來評估可控性，具體使用深度圖控制的RMSE和Canny邊緣的F1得分。此外，為了評估生成圖像的質量及其與輸入文本的對齊程度，我們采用FID（Fréchet Inception Distance）和CLIP-Score指標，所有實驗在512×512的分辨率下進行。

在圖像inpainting任務中，我們遵循LDM的設置來測量FID和LPIPS，評估在40-50%圖像區域需要修復時生成樣本的質量。對于圖像外推任務（outpainting），我們遵循MaskGIT設置，將圖像向右擴展50%，并使用FID和Inception Score（IS）與常見基線進行比較。

在文本到圖像生成任務中，使用兩種主要評估方法。首先，我們直觀展示PixWizard生成的圖像示例。此外，計算兩個自動評估指標：人類偏好評分（HPS）v2和MS-COCO數據集上的zero-shot FID-30K。

可控生成結果

無需為每個模型單獨訓練，PixWizard是一個一體化解決方案，能夠處理多種條件。如表3所示，PixWizard在深度條件下實現了最高的可控性和最佳的圖像質量，同時在圖像-文本對齊方面與當前的獨立模型相當。圖9展示了若干視覺樣本，證明了我們方法的有效性。

圖像修復結果

表3中與其他圖像修復方法的比較顯示，PixWizard在FID和LPIPS指標上提升了整體圖像質量。圖10中的定性示例進一步證明了PixWizard在生成連貫內容方面的有效性。我們將這一效果歸因于PixWizard在空白畫布上“繪畫”的能力，這大大增強了它識別被遮罩區域并在其中生成連貫內容的能力。

基于它強大的定位和圖像修復能力，PixWizard支持更精確的圖像編輯任務：

1. Remove Anything：解決對象移除問題，允許用戶在保持視覺無縫的情況下去除特定對象。過程包括識別和移除，隨后將遮罩應用于原始圖像，PixWizard填補該區域的適當背景細節。如圖11所示。
2. Replace Anything：允許用戶在圖像中替換任何對象。該過程與Remove Anything相似，但模型在移除對象的同時，確保用指定對象替換，背景保持一致。
3. Add Anything：用戶可以將任何對象插入圖像，并自由放置。用戶為所需區域添加遮罩并提供文本提示，PixWizard利用其強大的圖像修復能力生成相關內容。

圖像外推

在表3的定量比較結果中，PixWizard在圖像外推任務上超越了其他基線，達到了7.54的FID分數和22.18的IS分數，展現了最先進的圖像生成質量。圖10的樣本展示了PixWizard在各種場景和風格中合成圖像的能力，靈活處理多個方向和縱橫比的圖像外推，并且邊際一致性更好。

文本到圖像生成

在表3的定量比較結果中，PixWizard在COCO數據集上測試zero-shot性能時取得了9.56的FID分數。盡管某些生成模型的FID更低，但它們專注于文本到圖像任務，且依賴顯著更多的訓練資源。此外，我們還評估了人類偏好評分（HPS v2），這一強有力的基準用以評估文本到圖像合成中的人類偏好。PixWizard表現出色，生成的圖像質量與流行的文本到圖像生成器相當。圖12提供了視覺樣本，PixWizard支持高分辨率圖像合成，最高可達1024×1024，且適用于任何分辨率和縱橫比。

消融研究

我們對具有兩種不同指導的模型進行比較，以及使用動態語義token采樣（DSTS）和針對不同任務的兩階段訓練和數據平衡策略的消融實驗。

討論與結論

本研究探討了如何從任務定義、數據構建和模型架構三個關鍵方面構建一個多功能的交互式圖像生成助手。我們的目標是創建一個能夠精準遵循自由形式用戶指令進行圖像生成、操作和轉換的系統。PixWizard，消除了對特定任務設計選擇的需求，并在一系列多樣任務中實現了高度競爭的性能，展現了強大的泛化能力。

然而，本研究仍存在一些局限性。首先，當前的模型架構尚不支持多圖像輸入條件，這在研究中是一個日益重要和有價值的方向。其次，在與專業化模型相比時，尤其是在分割和圖像定位等挑戰性任務上，我們的模型仍有改進的空間。此外，文本編碼器和基礎模型的性能也發揮著至關重要的作用。更好的文本編碼方法使模型能夠更準確地理解和執行人類指令，而更大且更穩健的模型架構直接提升了最終輸出的質量。然而值得注意的是，在PixWizard中提出的模塊和策略可以輕松應用于其他強大的文本到圖像生成器。

未來，我們將探索使用更先進的擴散基礎模型，如SD3和FLUX，并繼續推動這一有前景的方向，直到我們在視覺領域實現“GPT-4的時刻”。