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

亮點(diǎn)直擊

• 任務(wù)統(tǒng)一：針對(duì)視覺(jué)任務(wù)的多樣性，本項(xiàng)目基于生成框架將多種任務(wù)轉(zhuǎn)化為圖像到圖像的翻譯問(wèn)題，并通過(guò)后處理將生成的可視化效果轉(zhuǎn)化為所需格式，從而簡(jiǎn)化表示形式的挑戰(zhàn)。
• 數(shù)據(jù)構(gòu)建：構(gòu)建了一個(gè)包含3000萬(wàn)條數(shù)據(jù)的全面訓(xùn)練集，主要支持圖像生成、編輯、修復(fù)、定位和密集預(yù)測(cè)等五大功能，旨在整合視覺(jué)領(lǐng)域的任務(wù)和數(shù)據(jù)多樣性。
• 架構(gòu)設(shè)計(jì)：采用基于流的Diffusion Transformer（DiT）作為基礎(chǔ)模型，強(qiáng)調(diào)其靈活性和穩(wěn)定性。通過(guò)動(dòng)態(tài)分區(qū)和填充方案增強(qiáng)模型對(duì)不同分辨率圖像的處理能力，并引入結(jié)構(gòu)感知和語(yǔ)義感知指導(dǎo)，以支持多模態(tài)指令的實(shí)現(xiàn)。

總結(jié)速覽

解決的問(wèn)題

統(tǒng)一多個(gè)視覺(jué)任務(wù)，包括圖像生成、圖像翻譯、圖像編輯和圖像修復(fù)等，并通過(guò)自由語(yǔ)言指令完成各類任務(wù)。

提出的方案

• 設(shè)計(jì)了一個(gè)多功能的圖像到圖像視覺(jué)助手PixWizard，整合各種視覺(jué)任務(wù)到一個(gè)統(tǒng)一的（圖像-指令）到圖像的生成框架。
• 構(gòu)建了一個(gè)全面的Omni Pixel-to-Pixel Instruction-Tuning Dataset，采用詳細(xì)的自然語(yǔ)言指令模板進(jìn)行訓(xùn)練。

應(yīng)用的技術(shù)

• 基于Diffusion Transformers (DiT)作為基礎(chǔ)模型，擴(kuò)展其功能以支持靈活的任意分辨率機(jī)制。
• 引入結(jié)構(gòu)感知和語(yǔ)義感知的指導(dǎo)，以有效融合輸入圖像的信息。

達(dá)到的效果

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

Omni Pixel-to-Pixel Instruction-Tuning數(shù)據(jù)集

為了使我們的視覺(jué)助手具備全面的圖像生成、編輯和翻譯功能，我們構(gòu)建了一個(gè)多任務(wù)、多模態(tài)的視覺(jué)指令調(diào)優(yōu)訓(xùn)練數(shù)據(jù)集。該數(shù)據(jù)集涵蓋七個(gè)主要領(lǐng)域，共計(jì)3000萬(wàn)條數(shù)據(jù)。根我們所知，這個(gè)數(shù)據(jù)集是最具多樣性且易于使用的圖像-指令-圖像三元組數(shù)據(jù)集。它由開(kāi)源和內(nèi)部數(shù)據(jù)集構(gòu)成，并在MLLM和人工審查的協(xié)助下進(jìn)行了過(guò)濾，如圖1所示：

PIXWIZARD

PixWizard是一個(gè)多功能的圖像到圖像模型。本節(jié)將從模型架構(gòu)（如下圖2所示）和訓(xùn)練策略的角度介紹PixWizard框架的細(xì)節(jié)。

基于流的條件指令調(diào)優(yōu)

架構(gòu)

文本編碼器：首先，我們?cè)赑ixWizard中使用Gemma-2B作為文本嵌入器，對(duì)文本提示進(jìn)行編碼。然而，在多任務(wù)學(xué)習(xí)中，僅依賴文本指令不足以準(zhǔn)確引導(dǎo)模型執(zhí)行用戶命令。為更有效地指導(dǎo)生成過(guò)程，我們進(jìn)一步結(jié)合了CLIP文本編碼器。通過(guò)對(duì)CLIP文本嵌入應(yīng)用全局平均池化，我們獲得了粗粒度的文本表示，然后通過(guò)多層感知機(jī)（MLP）生成任務(wù)嵌入（embedding）。該嵌入隨后通過(guò)調(diào)制機(jī)制與時(shí)間步嵌入相加，集成到PixWizard Block中。如圖3所示，這種方法自適應(yīng)地在潛在空間中聚類相似的任務(wù)指令，同時(shí)將不同任務(wù)的指令分離，從而幫助模型生成過(guò)程朝正確的任務(wù)方向發(fā)展。

結(jié)構(gòu)感知指導(dǎo)

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

語(yǔ)義感知指導(dǎo)

任意分辨率

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

兩階段訓(xùn)練和數(shù)據(jù)平衡策略

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

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

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

實(shí)驗(yàn)

第一部分結(jié)果

設(shè)置

對(duì)于圖像修復(fù)，遵循之前的研究（Conde et al., 2024; Potlapalli et al., 2024），我們?cè)谟?xùn)練期間準(zhǔn)備各種修復(fù)任務(wù)的數(shù)據(jù)集。對(duì)于評(píng)估，首先選擇兩個(gè)代表性基準(zhǔn)：Rain100L用于去雨，SIDD用于去噪。此外，論文還進(jìn)一步評(píng)估了其他修復(fù)任務(wù)的性能并檢驗(yàn)zero-shot的能力（具體請(qǐng)看論文的Appendix）。

對(duì)于圖像定位，我們?cè)趃RefCOCO、RefCOCO和RefCOCO+的驗(yàn)證和測(cè)試集上評(píng)估引用分割任務(wù)。我們使用累積IoU (cIoU)作為性能指標(biāo)。

密集圖像預(yù)測(cè)任務(wù)評(píng)估三項(xiàng)視覺(jué)任務(wù)：ADE20k用于語(yǔ)義分割，NYUv2和SUNRGB-D用于單目深度估計(jì)，以及NYU-Depth v2用于表面法線估計(jì)。對(duì)于語(yǔ)義分割，我們通過(guò)識(shí)別最近鄰的RGB顏色值來(lái)分配標(biāo)簽，并使用平均交并比 (mIoU)指標(biāo)評(píng)估準(zhǔn)確性。對(duì)于單目深度估計(jì)，我們對(duì)輸出圖像在三個(gè)通道上進(jìn)行平均，并應(yīng)用訓(xùn)練期間使用的線性變換層，獲得范圍在[0,10]米內(nèi)的深度估計(jì)，深度估計(jì)使用均方根誤差 (RMSE)進(jìn)行評(píng)估。對(duì)于表面法線估計(jì)，我們從輸出圖像中恢復(fù)相應(yīng)的法向量，并使用平均角誤差來(lái)評(píng)估準(zhǔn)確性。

結(jié)果

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

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

第二部分結(jié)果（圖像編輯）

設(shè)置

我們?cè)贛agicBrush測(cè)試（Zhang et al., 2024a）和Emu Edit測(cè)試（Sheynin et al., 2024）兩個(gè)基準(zhǔn)上評(píng)估PixWizard，以評(píng)估其圖像編輯能力的有效性。為了公平比較，我們主要與以指令引導(dǎo)的圖像編輯方法進(jìn)行對(duì)比，包括InstructPix2Pix、MagicBrush、Emu Edit和UltraEdit。與Emu Edit一致，我們使用L1距離、CLIP圖像相似度、DINO相似度、CLIP文本-圖像相似度和CLIP文本-圖像方向相似度作為指標(biāo)。

結(jié)果

圖像編輯

表2展示了我們與基線的結(jié)果。結(jié)果表明，我們的模型在自動(dòng)化量化指標(biāo)上始終超越InstructPix2Pix、MagicBrush和UltraEdit，并且在性能上與最先進(jìn)的方法Emu Edit相當(dāng)。圖8提供了定性比較。我們的模型精準(zhǔn)識(shí)別編輯區(qū)域，同時(shí)保留其他像素，展現(xiàn)了對(duì)給定指令的最佳理解。

第三部分結(jié)果（圖像生成）

設(shè)置

本節(jié)重點(diǎn)評(píng)估PixWizard的生成能力，具體任務(wù)包括經(jīng)典的文本到圖像生成、可控圖像生成、圖像inpainting和圖像外推outpainting。在可控圖像生成中，我們?cè)u(píng)估PixWizard基于特定條件（Canny邊緣圖和深度圖）生成圖像的能力。我們通過(guò)比較輸入條件與生成圖像中提取的相應(yīng)特征之間的相似度來(lái)評(píng)估可控性，具體使用深度圖控制的RMSE和Canny邊緣的F1得分。此外，為了評(píng)估生成圖像的質(zhì)量及其與輸入文本的對(duì)齊程度，我們采用FID（Fréchet Inception Distance）和CLIP-Score指標(biāo)，所有實(shí)驗(yàn)在512×512的分辨率下進(jìn)行。

在圖像inpainting任務(wù)中，我們遵循LDM的設(shè)置來(lái)測(cè)量FID和LPIPS，評(píng)估在40-50%圖像區(qū)域需要修復(fù)時(shí)生成樣本的質(zhì)量。對(duì)于圖像外推任務(wù)（outpainting），我們遵循MaskGIT設(shè)置，將圖像向右擴(kuò)展50%，并使用FID和Inception Score（IS）與常見(jiàn)基線進(jìn)行比較。

在文本到圖像生成任務(wù)中，使用兩種主要評(píng)估方法。首先，我們直觀展示PixWizard生成的圖像示例。此外，計(jì)算兩個(gè)自動(dòng)評(píng)估指標(biāo)：人類偏好評(píng)分（HPS）v2和MS-COCO數(shù)據(jù)集上的zero-shot FID-30K。

可控生成結(jié)果

無(wú)需為每個(gè)模型單獨(dú)訓(xùn)練，PixWizard是一個(gè)一體化解決方案，能夠處理多種條件。如表3所示，PixWizard在深度條件下實(shí)現(xiàn)了最高的可控性和最佳的圖像質(zhì)量，同時(shí)在圖像-文本對(duì)齊方面與當(dāng)前的獨(dú)立模型相當(dāng)。圖9展示了若干視覺(jué)樣本，證明了我們方法的有效性。

圖像修復(fù)結(jié)果

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

基于它強(qiáng)大的定位和圖像修復(fù)能力，PixWizard支持更精確的圖像編輯任務(wù)：

1. Remove Anything：解決對(duì)象移除問(wèn)題，允許用戶在保持視覺(jué)無(wú)縫的情況下去除特定對(duì)象。過(guò)程包括識(shí)別和移除，隨后將遮罩應(yīng)用于原始圖像，PixWizard填補(bǔ)該區(qū)域的適當(dāng)背景細(xì)節(jié)。如圖11所示。
2. Replace Anything：允許用戶在圖像中替換任何對(duì)象。該過(guò)程與Remove Anything相似，但模型在移除對(duì)象的同時(shí)，確保用指定對(duì)象替換，背景保持一致。
3. Add Anything：用戶可以將任何對(duì)象插入圖像，并自由放置。用戶為所需區(qū)域添加遮罩并提供文本提示，PixWizard利用其強(qiáng)大的圖像修復(fù)能力生成相關(guān)內(nèi)容。

圖像外推

在表3的定量比較結(jié)果中，PixWizard在圖像外推任務(wù)上超越了其他基線，達(dá)到了7.54的FID分?jǐn)?shù)和22.18的IS分?jǐn)?shù)，展現(xiàn)了最先進(jìn)的圖像生成質(zhì)量。圖10的樣本展示了PixWizard在各種場(chǎng)景和風(fēng)格中合成圖像的能力，靈活處理多個(gè)方向和縱橫比的圖像外推，并且邊際一致性更好。

文本到圖像生成

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

消融研究

我們對(duì)具有兩種不同指導(dǎo)的模型進(jìn)行比較，以及使用動(dòng)態(tài)語(yǔ)義token采樣（DSTS）和針對(duì)不同任務(wù)的兩階段訓(xùn)練和數(shù)據(jù)平衡策略的消融實(shí)驗(yàn)。

討論與結(jié)論

本研究探討了如何從任務(wù)定義、數(shù)據(jù)構(gòu)建和模型架構(gòu)三個(gè)關(guān)鍵方面構(gòu)建一個(gè)多功能的交互式圖像生成助手。我們的目標(biāo)是創(chuàng)建一個(gè)能夠精準(zhǔn)遵循自由形式用戶指令進(jìn)行圖像生成、操作和轉(zhuǎn)換的系統(tǒng)。PixWizard，消除了對(duì)特定任務(wù)設(shè)計(jì)選擇的需求，并在一系列多樣任務(wù)中實(shí)現(xiàn)了高度競(jìng)爭(zhēng)的性能，展現(xiàn)了強(qiáng)大的泛化能力。

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

未來(lái)，我們將探索使用更先進(jìn)的擴(kuò)散基礎(chǔ)模型，如SD3和FLUX，并繼續(xù)推動(dòng)這一有前景的方向，直到我們?cè)谝曈X(jué)領(lǐng)域?qū)崿F(xiàn)“GPT-4的時(shí)刻”。