文章鏈接：https://arxiv.org/pdf/2504.16915 
項(xiàng)目鏈接：https://mc-e.github.io/project/DreamO

亮點(diǎn)直擊

• 提出了DreamO，一個(gè)統(tǒng)一的圖像定制框架。它通過在預(yù)訓(xùn)練的 DiT 模型上訓(xùn)練少量額外參數(shù)，實(shí)現(xiàn)了各種復(fù)雜和多條件的定制任務(wù)。
• 基于擴(kuò)散模型內(nèi)部的表示對(duì)應(yīng)關(guān)系，本文設(shè)計(jì)了一種特征路由約束，以增強(qiáng)一致性保真度，并在多條件場景中實(shí)現(xiàn)有效的解耦。
• 引入了一種漸進(jìn)式訓(xùn)練策略，以促進(jìn)多任務(wù)和復(fù)雜任務(wù)設(shè)置的收斂。本文還設(shè)計(jì)了一種占位符策略，以建立文本描述與條件圖像之間的對(duì)應(yīng)關(guān)系。
• 大量實(shí)驗(yàn)表明，本文的方法不僅在廣泛的圖像定制場景中實(shí)現(xiàn)了高質(zhì)量的結(jié)果，而且在適應(yīng)多條件場景時(shí)表現(xiàn)出強(qiáng)大的靈活性。

總結(jié)速覽

解決的問題

• 任務(wù)局限性：現(xiàn)有圖像定制方法多為特定任務(wù)設(shè)計(jì)（如身份保持、風(fēng)格遷移、虛擬試穿等），缺乏通用性，難以整合多類型控制條件。
• 多條件交互困難：現(xiàn)有統(tǒng)一框架（如UniControl、OminiControl）對(duì)多條件的組合控制能力不足，存在冗余或交互僵化問題。
• 訓(xùn)練成本高：部分方案（如OmniGen、UniReal）依賴大語言模型或視頻預(yù)訓(xùn)練，需數(shù)百GPU資源，成本高昂。

提出的方案

• 統(tǒng)一框架設(shè)計(jì)：基于擴(kuò)散Transformer（DiT）構(gòu)建通用架構(gòu)，通過統(tǒng)一序列處理多類型輸入（如身份、物體、風(fēng)格等）。
• 特征路由約束：在訓(xùn)練中引入特征路由機(jī)制，精準(zhǔn)查詢參考圖像信息，提升內(nèi)容保真度并解耦不同控制條件。
• 占位符策略：將條件與生成結(jié)果中的特定位置關(guān)聯(lián)，實(shí)現(xiàn)條件放置的靈活控制。
• 漸進(jìn)式訓(xùn)練策略：

• 初期：簡單任務(wù)訓(xùn)練，建立基礎(chǔ)一致性。
• 中期：全規(guī)模訓(xùn)練，綜合提升定制能力。
• 后期：質(zhì)量對(duì)齊階段，修正低質(zhì)量數(shù)據(jù)偏差。

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

• 擴(kuò)散Transformer（DiT）：作為主干模型，統(tǒng)一處理多模態(tài)輸入條件。
• 統(tǒng)一序列編碼：將不同條件（如草圖、深度圖、參考圖像）轉(zhuǎn)換為序列化輸入。
• 特征路由約束：通過注意力機(jī)制優(yōu)化條件特征的檢索與融合。
• 大規(guī)模多任務(wù)數(shù)據(jù)集：構(gòu)建涵蓋身份、物體、風(fēng)格等多樣化任務(wù)的數(shù)據(jù)集。

達(dá)到的效果

• 多功能支持：單模型支持身份保持、物體外觀遷移、虛擬試穿、風(fēng)格化等多種任務(wù)。
• 高質(zhì)量生成：通過漸進(jìn)訓(xùn)練和質(zhì)量對(duì)齊，生成結(jié)果保真度高且無明顯偏差。
• 靈活控制：占位符策略實(shí)現(xiàn)多條件的精確位置控制，支持復(fù)雜場景生成（如多物體插入）。
• 高效訓(xùn)練：僅需輕微訓(xùn)練成本（對(duì)比OmniGen等方案），即可實(shí)現(xiàn)多條件交互與統(tǒng)一部署。

方法

概述

路由約束

受UniPortrait和AnyStory的啟發(fā)，本文在DiT架構(gòu)中設(shè)計(jì)了適用于通用圖像定制任務(wù)的路由約束。如前面圖2所示，在條件引導(dǎo)框架中，條件圖像與生成結(jié)果之間存在交叉注意力：

為可視化 DiT 框架中的圖像間交叉注意力圖，進(jìn)行了簡化實(shí)驗(yàn)。具體而言，使用 Subject200k 數(shù)據(jù)集在圖 2 的可控框架上訓(xùn)練主體驅(qū)動(dòng)生成任務(wù)。下圖 3 第一行展示了生成結(jié)果及內(nèi)部交叉注意力圖，本文選擇 Flux 雙塊結(jié)構(gòu)中的三個(gè)層級(jí)（第 2、8、14 層）進(jìn)行注意力可視化?？梢杂^察到：在 DiT 框架下，圖像間注意力呈現(xiàn)空間模式，雖粒度較粗但能聚焦目標(biāo)物體。

為約束圖像間注意力集中于特定主體，采用 MSE 損失優(yōu)化 DiT 中條件圖像與生成結(jié)果間的注意力關(guān)系：

如上圖3第二行所示，經(jīng)過路由約束訓(xùn)練后，條件圖像的注意力明顯聚焦于目標(biāo)主體，生成結(jié)果在細(xì)節(jié)上與參考圖像的一致性也有所提升。除了增強(qiáng)一致性外，該策略還有助于多參考情況下的解耦。更多細(xì)節(jié)見消融實(shí)驗(yàn)（下圖11）。下圖4展示了占位符的效果，它可以控制生成結(jié)果中條件的位置。

訓(xùn)練數(shù)據(jù)構(gòu)建

本工作旨在實(shí)現(xiàn)廣義的圖像定制。本文收集了一個(gè)覆蓋廣泛任務(wù)的訓(xùn)練數(shù)據(jù)集。

身份配對(duì)數(shù)據(jù)

由于高質(zhì)量的身份配對(duì)數(shù)據(jù)難以從互聯(lián)網(wǎng)獲取，我們采用開源ID定制方法PuLID進(jìn)行數(shù)據(jù)集構(gòu)建，因其展現(xiàn)出高面部相似性和強(qiáng)提示對(duì)齊。具體而言，我們?yōu)镻uLID-SDXL提供參考人臉圖像和描述目標(biāo)風(fēng)格的文本提示，使其控制生成肖像的風(fēng)格，從而得到（參考人臉，提示，風(fēng)格化人臉）形式的訓(xùn)練對(duì)。對(duì)于寫實(shí)場景，我們使用PuLID-FLUX生成同一身份的兩張圖像作為相互參考。

主體驅(qū)動(dòng)數(shù)據(jù)

對(duì)于單主體驅(qū)動(dòng)的圖像定制，本文使用開源數(shù)據(jù)集Subject200K作為訓(xùn)練數(shù)據(jù)。為補(bǔ)充角色相關(guān)條件的缺失，我們通過檢索和人工篩選收集了一組成對(duì)的角色相關(guān)數(shù)據(jù)。對(duì)于多主體驅(qū)動(dòng)的圖像定制，我們?cè)赟ubject200K數(shù)據(jù)集上通過圖像拼接構(gòu)建雙列圖像。此時(shí)，文本提示變?yōu)椋骸吧梢粡堧p列圖像。左側(cè)為{prompt1}，右側(cè)為{prompt2}”，其中prompt1和prompt2分別對(duì)應(yīng)兩幅拼接圖像的描述。為增強(qiáng)處理復(fù)雜場景（尤其是涉及同一類型的多個(gè)參考因素）的能力，我們?cè)谕ㄟ^圖像拼接構(gòu)建多主體圖像時(shí)，會(huì)按一定比例從同一類別中選擇主體。本文還使用開源數(shù)據(jù)集X2I-subject進(jìn)行多主體驅(qū)動(dòng)訓(xùn)練。

為增強(qiáng)以人類為參考的主體驅(qū)動(dòng)生成，仿照MoiveGen的流程構(gòu)建了另一個(gè)數(shù)據(jù)集。從長視頻數(shù)據(jù)集出發(fā)，通過內(nèi)容感知的場景檢測獲取多個(gè)短視頻片段。為獲取片段內(nèi)實(shí)例匹配，我們使用Mask2Former提取關(guān)鍵幀中所有人的掩碼，并在視頻片段中進(jìn)行目標(biāo)跟蹤。為獲取跨片段實(shí)例匹配，本文采用SigLip嵌入所有人體對(duì)象并進(jìn)行聚類。之后，使用GPT-4o結(jié)合圖像和彩色實(shí)例掩碼輸入生成實(shí)例 grounded caption。

試穿數(shù)據(jù)

在該任務(wù)中，通過兩種來源構(gòu)建成對(duì)的試穿數(shù)據(jù)集。部分?jǐn)?shù)據(jù)直接從網(wǎng)絡(luò)收集模特與服裝的配對(duì)圖像；另一部分則先爬取高質(zhì)量模特圖像作為真值，再通過圖像分割提取服裝并構(gòu)建對(duì)應(yīng)配對(duì)數(shù)據(jù)。所有收集的圖像均經(jīng)過人工過濾以剔除低質(zhì)量樣本。服裝區(qū)域通過分割算法提取。

風(fēng)格驅(qū)動(dòng)數(shù)據(jù)

本文旨在解決兩類風(fēng)格遷移任務(wù)：(1) 風(fēng)格參考圖像+內(nèi)容文本描述控制；(2) 風(fēng)格參考圖像+內(nèi)容參考圖像。對(duì)于第一類任務(wù)，訓(xùn)練數(shù)據(jù)需包含同一風(fēng)格的參考圖像與目標(biāo)圖像。如圖13所示，我們采用內(nèi)部風(fēng)格遷移模型（即基于SDXL的模型，使用類似InstantStyle的解耦策略），在兩種不同提示下生成相同風(fēng)格但內(nèi)容各異的圖像。對(duì)于第二類任務(wù)，訓(xùn)練需要風(fēng)格參考圖像、內(nèi)容參考圖像和目標(biāo)圖像。此時(shí)目標(biāo)圖像與風(fēng)格參考共享風(fēng)格，同時(shí)保持與內(nèi)容參考圖像相同的內(nèi)容結(jié)構(gòu)?；诘谝活愑?xùn)練數(shù)據(jù)，我們通過Canny-guided Flux生成與風(fēng)格圖像對(duì)應(yīng)的自然圖像作為內(nèi)容參考。

路由掩碼提取

為了獲取路由約束（即Eq.3）的標(biāo)簽，我們從各種類型的訓(xùn)練數(shù)據(jù)中提取目標(biāo)主體掩碼。具體來說，對(duì)于X2I-Subject數(shù)據(jù)集，使用InternVL模型提取主體描述，然后利用LISA結(jié)合描述和目標(biāo)圖像生成主體掩碼。對(duì)于Subject200K數(shù)據(jù)集，使用InternVL獲取主體名稱，然后通過LISA預(yù)測相應(yīng)的掩碼。對(duì)于人臉數(shù)據(jù)，我們直接使用LISA從目標(biāo)圖像中提取人臉掩碼。

盡管路由約束僅應(yīng)用于有限的訓(xùn)練數(shù)據(jù)（即Subject200K、人臉數(shù)據(jù)和X2I-Subject），但本文的方法在未見場景中表現(xiàn)出強(qiáng)大的泛化能力。如下圖5所示，第一行顯示在試穿任務(wù)中，條件圖像準(zhǔn)確地定位了目標(biāo)區(qū)域。第二行顯示在更復(fù)雜的條件下，模型有效地區(qū)分了與不同條件圖像相關(guān)的區(qū)域。具體來說，身份參考圖像聚焦于面部和身體，而帶有太陽鏡的圖像則直接關(guān)注太陽鏡區(qū)域，避開了面部。

漸進(jìn)式訓(xùn)練流程

在實(shí)驗(yàn)中，本文發(fā)現(xiàn)直接在全部數(shù)據(jù)上訓(xùn)練會(huì)導(dǎo)致收斂困難。這主要是由于優(yōu)化參數(shù)的容量有限，使得模型在復(fù)雜數(shù)據(jù)分布下難以捕獲特定任務(wù)的能力。此外，由于訓(xùn)練數(shù)據(jù)圖像質(zhì)量的影響，訓(xùn)練后的生成質(zhì)量與Flux的生成先驗(yàn)存在偏差。

為了解決這些問題，本文設(shè)計(jì)了一種漸進(jìn)式訓(xùn)練策略，使模型能夠在不同任務(wù)間平滑收斂，同時(shí)糾正訓(xùn)練數(shù)據(jù)對(duì)Flux高質(zhì)量生成先驗(yàn)的影響。訓(xùn)練流程如下圖6所示。首先在主體驅(qū)動(dòng)的訓(xùn)練數(shù)據(jù)上優(yōu)化模型，以初始化模型的一致性保持能力。需要注意的是，訓(xùn)練中使用的Subject200K數(shù)據(jù)集是由基模型（即Flux）生成的，因此與模型生成空間具有相似的分布，這有助于快速收斂。由于X2I-subject數(shù)據(jù)集是通過MS-Diffusion合成的，許多訓(xùn)練樣本包含不希望的偽影和失真。因此，在這個(gè)預(yù)熱階段，還使用了雙列Subject200K圖像作為訓(xùn)練數(shù)據(jù)的一部分，以促進(jìn)多主體生成控制的快速收斂。圖6的右側(cè)部分顯示，在第一個(gè)訓(xùn)練階段后，模型獲得了初步的主體驅(qū)動(dòng)生成能力，并表現(xiàn)出強(qiáng)大的文本跟隨性能。在第二個(gè)訓(xùn)練階段，我們整合所有訓(xùn)練數(shù)據(jù)并進(jìn)行全數(shù)據(jù)微調(diào)，使模型能夠進(jìn)一步收斂到本工作中定義的所有子任務(wù)上。

在全數(shù)據(jù)訓(xùn)練的第二階段后，生成質(zhì)量受到訓(xùn)練數(shù)據(jù)的嚴(yán)重影響，尤其是低質(zhì)量的訓(xùn)練樣本。為了使生成質(zhì)量與Flux的生成先驗(yàn)重新對(duì)齊，本文設(shè)計(jì)了一個(gè)圖像質(zhì)量精煉訓(xùn)練階段。具體來說，我們利用Flux生成了約40K個(gè)訓(xùn)練樣本。在訓(xùn)練過程中，使用原始圖像作為參考來指導(dǎo)模型重建自身。為了防止復(fù)制粘貼效應(yīng)，丟棄了95%的參考圖像token。經(jīng)過短時(shí)間優(yōu)化后，生成質(zhì)量顯著提高，實(shí)現(xiàn)了與Flux生成先驗(yàn)的對(duì)齊。

實(shí)驗(yàn)

實(shí)現(xiàn)細(xì)節(jié)

本文采用Flux-1.0-dev作為基礎(chǔ)模型，設(shè)置額外LoRA的秩為128，共增加707M參數(shù)。訓(xùn)練過程使用Adam優(yōu)化器，學(xué)習(xí)率設(shè)為4e-5，在8塊NVIDIA A100 80G GPU上進(jìn)行訓(xùn)練。批大小設(shè)置為8。第一階段訓(xùn)練20K次迭代，第二階段90K次迭代，最后階段3K次迭代。

模型能力

DreamO是一個(gè)統(tǒng)一框架，能夠處理多種圖像定制任務(wù)。下圖7展示了身份驅(qū)動(dòng)的圖像生成，本文的方法既支持特定個(gè)體的定制，也支持組合身份控制。模型可以生成保留單個(gè)人物身份或融合多個(gè)人物特征的圖像，同時(shí)通過文本輸入精確控制其他屬性和場景細(xì)節(jié)。

下圖8展示了物體驅(qū)動(dòng)的圖像定制結(jié)果，可見DreamO既支持單主體條件控制，也支持多主體組合，能夠?qū)⒍鄻踊曈X元素整合到統(tǒng)一輸出中并保持語義一致性。

下圖9展示了試穿能力，包括特定身份的視覺試穿和任意主體的創(chuàng)意應(yīng)用。

下圖10展示了風(fēng)格驅(qū)動(dòng)定制，風(fēng)格完全由參考圖像引導(dǎo)，其他屬性通過文本或圖像輸入控制，實(shí)現(xiàn)細(xì)粒度的多模態(tài)操控。這些結(jié)果凸顯了本文方法整合異質(zhì)控制信號(hào)（如身份、物體和試穿）的能力，實(shí)現(xiàn)更具表現(xiàn)力和復(fù)雜度的圖像定制。

消融研究

路由約束：在DiT訓(xùn)練中引入路由約束以提升生成保真度并促進(jìn)多條件解耦。下圖11顯示，在單條件生成中去除該約束會(huì)導(dǎo)致參考保真度下降（如服裝顏色與參考不一致）；在多條件設(shè)置中會(huì)導(dǎo)致條件耦合（如兩個(gè)玩具特征交叉）。這些結(jié)果證實(shí)路由約束能提升不同條件的保真度和解耦性。

漸進(jìn)式訓(xùn)練：為使模型在復(fù)雜數(shù)據(jù)分布下更好收斂并修正訓(xùn)練數(shù)據(jù)分布對(duì)生成質(zhì)量的影響，設(shè)計(jì)了漸進(jìn)式訓(xùn)練策略。圖12表明，直接在全部數(shù)據(jù)集上訓(xùn)練會(huì)導(dǎo)致次優(yōu)收斂（尤其在多主體一致性等復(fù)雜任務(wù)中）；先在較小且易學(xué)的數(shù)據(jù)集（如Subject200K）上預(yù)熱再聯(lián)合訓(xùn)練可改善收斂，但生成質(zhì)量易受訓(xùn)練數(shù)據(jù)分布影響而偏離Flux的生成先驗(yàn)；通過引入圖像質(zhì)量調(diào)優(yōu)階段，模型能產(chǎn)生更高質(zhì)量的生成結(jié)果。

結(jié)論

本文提出了DreamO，一個(gè)基于預(yù)訓(xùn)練DiT框架的統(tǒng)一圖像定制系統(tǒng)，能夠處理多種控制條件（如身份、風(fēng)格、主體和試穿）。首先構(gòu)建了包含廣泛圖像定制任務(wù)的大規(guī)模訓(xùn)練數(shù)據(jù)集。通過將所有條件類型整合到DiT輸入序列，并引入特征路由約束，DreamO在確保高保真一致性的同時(shí)，有效解耦了異構(gòu)控制信號(hào)。

針對(duì)復(fù)雜數(shù)據(jù)分布下保持基礎(chǔ)模型（Flux）生成先驗(yàn)的同時(shí)學(xué)習(xí)強(qiáng)任務(wù)特定能力的挑戰(zhàn)，設(shè)計(jì)了漸進(jìn)式訓(xùn)練策略。該方法使模型能逐步獲得多樣化控制能力，同時(shí)保持基礎(chǔ)模型的圖像質(zhì)量。大量實(shí)驗(yàn)表明，DreamO能夠高質(zhì)量地完成各類圖像定制任務(wù)。此外，其基于LoRA的輕量級(jí)設(shè)計(jì)實(shí)現(xiàn)了低計(jì)算成本的高效部署。

本文轉(zhuǎn)自AI生成未來 ，作者：AI生成未來

原文鏈接:??https://mp.weixin.qq.com/s/VU8RPEq6nTsNHDPN8rIFEw??