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

亮點直擊

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

總結速覽

解決的問題

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

提出的方案

• 統一框架設計：基于擴散Transformer（DiT）構建通用架構，通過統一序列處理多類型輸入（如身份、物體、風格等）。
• 特征路由約束：在訓練中引入特征路由機制，精準查詢參考圖像信息，提升內容保真度并解耦不同控制條件。
• 占位符策略：將條件與生成結果中的特定位置關聯，實現條件放置的靈活控制。
• 漸進式訓練策略：

• 初期：簡單任務訓練，建立基礎一致性。
• 中期：全規模訓練，綜合提升定制能力。
• 后期：質量對齊階段，修正低質量數據偏差。

應用的技術

• 擴散Transformer（DiT）：作為主干模型，統一處理多模態輸入條件。
• 統一序列編碼：將不同條件（如草圖、深度圖、參考圖像）轉換為序列化輸入。
• 特征路由約束：通過注意力機制優化條件特征的檢索與融合。
• 大規模多任務數據集：構建涵蓋身份、物體、風格等多樣化任務的數據集。

達到的效果

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

方法

概述

路由約束

受UniPortrait和AnyStory的啟發，本文在DiT架構中設計了適用于通用圖像定制任務的路由約束。如前面圖2所示，在條件引導框架中，條件圖像與生成結果之間存在交叉注意力：

為可視化 DiT 框架中的圖像間交叉注意力圖，進行了簡化實驗。具體而言，使用 Subject200k 數據集在圖 2 的可控框架上訓練主體驅動生成任務。下圖 3 第一行展示了生成結果及內部交叉注意力圖，本文選擇 Flux 雙塊結構中的三個層級（第 2、8、14 層）進行注意力可視化。可以觀察到：在 DiT 框架下，圖像間注意力呈現空間模式，雖粒度較粗但能聚焦目標物體。

為約束圖像間注意力集中于特定主體，采用 MSE 損失優化 DiT 中條件圖像與生成結果間的注意力關系：

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

訓練數據構建

本工作旨在實現廣義的圖像定制。本文收集了一個覆蓋廣泛任務的訓練數據集。

身份配對數據

由于高質量的身份配對數據難以從互聯網獲取，我們采用開源ID定制方法PuLID進行數據集構建，因其展現出高面部相似性和強提示對齊。具體而言，我們為PuLID-SDXL提供參考人臉圖像和描述目標風格的文本提示，使其控制生成肖像的風格，從而得到（參考人臉，提示，風格化人臉）形式的訓練對。對于寫實場景，我們使用PuLID-FLUX生成同一身份的兩張圖像作為相互參考。

主體驅動數據

對于單主體驅動的圖像定制，本文使用開源數據集Subject200K作為訓練數據。為補充角色相關條件的缺失，我們通過檢索和人工篩選收集了一組成對的角色相關數據。對于多主體驅動的圖像定制，我們在Subject200K數據集上通過圖像拼接構建雙列圖像。此時，文本提示變為：“生成一張雙列圖像。左側為{prompt1}，右側為{prompt2}”，其中prompt1和prompt2分別對應兩幅拼接圖像的描述。為增強處理復雜場景（尤其是涉及同一類型的多個參考因素）的能力，我們在通過圖像拼接構建多主體圖像時，會按一定比例從同一類別中選擇主體。本文還使用開源數據集X2I-subject進行多主體驅動訓練。

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

試穿數據

在該任務中，通過兩種來源構建成對的試穿數據集。部分數據直接從網絡收集模特與服裝的配對圖像；另一部分則先爬取高質量模特圖像作為真值，再通過圖像分割提取服裝并構建對應配對數據。所有收集的圖像均經過人工過濾以剔除低質量樣本。服裝區域通過分割算法提取。

風格驅動數據

本文旨在解決兩類風格遷移任務：(1) 風格參考圖像+內容文本描述控制；(2) 風格參考圖像+內容參考圖像。對于第一類任務，訓練數據需包含同一風格的參考圖像與目標圖像。如圖13所示，我們采用內部風格遷移模型（即基于SDXL的模型，使用類似InstantStyle的解耦策略），在兩種不同提示下生成相同風格但內容各異的圖像。對于第二類任務，訓練需要風格參考圖像、內容參考圖像和目標圖像。此時目標圖像與風格參考共享風格，同時保持與內容參考圖像相同的內容結構。基于第一類訓練數據，我們通過Canny-guided Flux生成與風格圖像對應的自然圖像作為內容參考。

路由掩碼提取

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

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

漸進式訓練流程

在實驗中，本文發現直接在全部數據上訓練會導致收斂困難。這主要是由于優化參數的容量有限，使得模型在復雜數據分布下難以捕獲特定任務的能力。此外，由于訓練數據圖像質量的影響，訓練后的生成質量與Flux的生成先驗存在偏差。

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

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

實驗

實現細節

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

模型能力

DreamO是一個統一框架，能夠處理多種圖像定制任務。下圖7展示了身份驅動的圖像生成，本文的方法既支持特定個體的定制，也支持組合身份控制。模型可以生成保留單個人物身份或融合多個人物特征的圖像，同時通過文本輸入精確控制其他屬性和場景細節。

下圖8展示了物體驅動的圖像定制結果，可見DreamO既支持單主體條件控制，也支持多主體組合，能夠將多樣化視覺元素整合到統一輸出中并保持語義一致性。

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

下圖10展示了風格驅動定制，風格完全由參考圖像引導，其他屬性通過文本或圖像輸入控制，實現細粒度的多模態操控。這些結果凸顯了本文方法整合異質控制信號（如身份、物體和試穿）的能力，實現更具表現力和復雜度的圖像定制。

消融研究

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

漸進式訓練：為使模型在復雜數據分布下更好收斂并修正訓練數據分布對生成質量的影響，設計了漸進式訓練策略。圖12表明，直接在全部數據集上訓練會導致次優收斂（尤其在多主體一致性等復雜任務中）；先在較小且易學的數據集（如Subject200K）上預熱再聯合訓練可改善收斂，但生成質量易受訓練數據分布影響而偏離Flux的生成先驗；通過引入圖像質量調優階段，模型能產生更高質量的生成結果。

結論

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

針對復雜數據分布下保持基礎模型（Flux）生成先驗的同時學習強任務特定能力的挑戰，設計了漸進式訓練策略。該方法使模型能逐步獲得多樣化控制能力，同時保持基礎模型的圖像質量。大量實驗表明，DreamO能夠高質量地完成各類圖像定制任務。此外，其基于LoRA的輕量級設計實現了低計算成本的高效部署。

本文轉自AI生成未來 ，作者：AI生成未來

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