DiT控制新紀(jì)元!"即插即控",Tiamat AI重磅開(kāi)源EasyControl:任意比例出圖+推理速度狂飆
文章鏈接:https://arxiv.org/abs/2503.07027
項(xiàng)目鏈接:https://github.com/Xiaojiu-z/EasyControl
亮點(diǎn)直擊
- 提出了EasyControl,一種面向DiT模型的條件生成新范式。EasyControl中,每個(gè)條件通過(guò)獨(dú)立的條件分支進(jìn)行處理,該分支通過(guò)條件注入LoRA模塊從預(yù)訓(xùn)練DiT模型適配而來(lái)。此設(shè)計(jì)實(shí)現(xiàn)了與定制模型的無(wú)縫集成,支持靈活的條件注入與多條件高效融合。
- 高效性:框架通過(guò)兩項(xiàng)關(guān)鍵創(chuàng)新實(shí)現(xiàn)高效計(jì)算。位置感知訓(xùn)練范式將輸入條件標(biāo)準(zhǔn)化為固定分辨率,確保適應(yīng)性與計(jì)算效率;因果注意力機(jī)制與KV緩存技術(shù)的結(jié)合,首次在條件生成任務(wù)中成功應(yīng)用KV緩存,顯著降低延遲并提升整體效率。
- 靈活性:EasyControl支持生成不同分辨率與長(zhǎng)寬比的圖像,通過(guò)平衡高質(zhì)量生成與多樣化需求,確??鐖?chǎng)景的魯棒性能。
總結(jié)速覽
解決的問(wèn)題
- 計(jì)算效率瓶頸
DiT架構(gòu)中,自注意力機(jī)制因圖像token的引入導(dǎo)致計(jì)算復(fù)雜度呈平方級(jí)增長(zhǎng),增加推理延遲,限制實(shí)際應(yīng)用擴(kuò)展。
- 多條件協(xié)同控制困難
現(xiàn)有方法在單條件訓(xùn)練范式下難以實(shí)現(xiàn)多條件穩(wěn)定協(xié)同,潛在空間中的條件信號(hào)表征沖突導(dǎo)致生成質(zhì)量下降,尤其在零樣本多條件組合場(chǎng)景中表現(xiàn)不佳。
- 模型適配性不足
當(dāng)前參數(shù)高效微調(diào)方法(如LoRA)與社區(qū)定制模型存在參數(shù)沖突,導(dǎo)致風(fēng)格遷移時(shí)特征退化,模塊缺乏真正的即插即用特性。
提出的方案
- 輕量級(jí)條件注入模塊(Condition Injection LoRA Module)
通過(guò)隔離處理?xiàng)l件信號(hào),以并行分支機(jī)制注入預(yù)訓(xùn)練模型,僅對(duì)條件分支token應(yīng)用低秩投影,凍結(jié)文本和噪聲分支權(quán)重,實(shí)現(xiàn)與定制模型的無(wú)縫兼容。
- 位置感知訓(xùn)練范式(Position-Aware Training Paradigm)
標(biāo)準(zhǔn)化輸入條件的分辨率,結(jié)合位置感知插值技術(shù),保持條件token與噪聲token的空間一致性,支持任意長(zhǎng)寬比和多分辨率生成。
- 因果注意力機(jī)制與KV緩存(Causal Attention + KV Cache)
在初始擴(kuò)散步(t=0)預(yù)計(jì)算條件特征的Key-Value對(duì)并緩存,后續(xù)時(shí)間步(t≥1)直接復(fù)用,顯著減少重復(fù)計(jì)算。
應(yīng)用的技術(shù)
- 低秩自適應(yīng)(LoRA)
條件分支采用低秩矩陣分解,避免修改基礎(chǔ)模型權(quán)重,實(shí)現(xiàn)高效參數(shù)微調(diào)。
- 分辨率歸一化與位置感知插值
將輸入條件縮放到固定分辨率以減少序列長(zhǎng)度,通過(guò)插值技術(shù)保留空間信息。
- 因果注意力與KV緩存優(yōu)化
將傳統(tǒng)全注意力替換為因果注意力,結(jié)合KV緩存技術(shù)復(fù)用條件特征,降低計(jì)算開(kāi)銷。
達(dá)到的效果
- 高效推理
KV緩存技術(shù)減少約30%的推理延遲,支持實(shí)時(shí)生成;分辨率歸一化降低輸入序列長(zhǎng)度,提升計(jì)算效率。
- 靈活控制
零樣本多條件泛化能力:即使僅訓(xùn)練單條件數(shù)據(jù),仍可和諧融合多條件(如空間控制+主體驅(qū)動(dòng))。
- 廣泛兼容性
即插即用設(shè)計(jì)兼容社區(qū)定制模型(如風(fēng)格化DiT),無(wú)需重新訓(xùn)練即可適配多樣任務(wù)。
- 高質(zhì)量生成
在虛擬試穿、圖像編輯、多分辨率生成等任務(wù)中,生成質(zhì)量顯著優(yōu)于基線方法(如ControlNet for DiT)。
效果可視化對(duì)比
多條件生成設(shè)置下與身份自定義方法比較
空間控制生成可視化對(duì)比
不同分辨率生成設(shè)置下與基線方法可視化比較
主題控件生成可視化
方法
本節(jié)將詳細(xì)介紹EasyControl的技術(shù)實(shí)現(xiàn),方法整體框架如下圖2所示。EasyControl基于FLUX.1開(kāi)發(fā)平臺(tái)構(gòu)建,包含以下核心組件:條件注入LoRA模塊、因果注意力機(jī)制、位置感知訓(xùn)練范式以及推理KV緩存。
條件注入LoRA模塊
為高效融合條件信號(hào)并保持預(yù)訓(xùn)練模型的泛化能力,在FLUX架構(gòu)中擴(kuò)展了獨(dú)立的條件分支。與傳統(tǒng)添加獨(dú)立控制模塊的方法不同,本方案通過(guò)以下方式實(shí)現(xiàn)條件信息的無(wú)縫集成,同時(shí)避免冗余參數(shù)與計(jì)算開(kāi)銷:
通過(guò)僅在條件分支應(yīng)用基于LoRA的自適應(yīng)機(jī)制,確保條件信號(hào)能高效注入模型,同時(shí)不破壞預(yù)訓(xùn)練的文本和噪聲表征。這種定向自適應(yīng)使模型能靈活整合條件信息,同時(shí)保持原始特征空間的完整性,從而實(shí)現(xiàn)更具可控性和高保真度的圖像生成。
EasyControl 中的因果注意力機(jī)制
條件因果注意力
該機(jī)制強(qiáng)制兩條規(guī)則:
- 各條件分支內(nèi)部進(jìn)行條件內(nèi)計(jì)算
- 采用注意力掩碼防止訓(xùn)練期間條件token查詢?nèi)ピ耄ㄎ谋?amp;噪聲)token
形式化定義單條件訓(xùn)練時(shí)的輸入序列為:
該設(shè)計(jì)通過(guò)阻斷條件分支到去噪分支(噪聲&文本)的單向注意力,同時(shí)允許去噪分支token自由聚合條件信號(hào)。通過(guò)嚴(yán)格隔離條件到去噪的查詢操作,該方案在推理時(shí)可實(shí)現(xiàn)各分支解耦的KV緩存狀態(tài),從而減少冗余計(jì)算并顯著提升圖像生成效率。
互因果注意力
本文模型僅使用單條件輸入訓(xùn)練,每個(gè)條件token學(xué)習(xí)與去噪token的優(yōu)化交互。在多條件推理時(shí),雖然所有條件都與去噪token正常交互,但由于未訓(xùn)練的跨條件token交互會(huì)導(dǎo)致條件間干擾(見(jiàn)下圖5)。
該機(jī)制通過(guò)以下形式化定義實(shí)現(xiàn)多條件推理時(shí)的輸入序列:
位置感知訓(xùn)練范式
雖然這種基礎(chǔ)下采樣方法對(duì)主體條件(如人臉圖像)有效,但會(huì)破壞空間條件(如Canny邊緣圖)的幾何對(duì)齊性,限制模型生成任意分辨率圖像的能力。為此,引入兩種定制策略:
- 位置感知插值(PAI):針對(duì)空間條件,在縮放時(shí)保持像素級(jí)對(duì)齊;
- 位置編碼偏移策略(詳見(jiàn)附錄B):針對(duì)主體條件,在高度維度施加固定位移。
位置感知插值
為保持條件token與噪聲token的空間一致性,提出位置感知插值(PAI)策略,在條件信號(hào)縮放過(guò)程中對(duì)位置編碼進(jìn)行插值。該方法確保模型能精確捕捉控制條件與生成圖像像素間的空間關(guān)系。
原始圖像中的位置編碼序列表示為:
而調(diào)整大小后的圖像的插值序列為:
這確保了調(diào)整大小后的圖像中空間關(guān)系的保留。
損失函數(shù)
本文損失函數(shù)使用流匹配損失。其數(shù)學(xué)表達(dá)式如下:
通過(guò)KV緩存實(shí)現(xiàn)高效推理
通過(guò)利用因果注意機(jī)制,本文框架將條件分支隔離為一個(gè)與去噪時(shí)間步無(wú)關(guān)的計(jì)算模塊。這種獨(dú)特設(shè)計(jì)使得在推理過(guò)程中能夠新穎地應(yīng)用KV緩存技術(shù)。由于條件分支的計(jì)算與去噪時(shí)間步無(wú)關(guān),在初始時(shí)間步只需預(yù)計(jì)算并存儲(chǔ)所有條件特征的鍵值(KV)對(duì)一次。這些緩存的KV對(duì)在所有后續(xù)時(shí)間步中重復(fù)使用,消除了相同條件特征的冗余重新計(jì)算。這種方法通過(guò)避免N次重新計(jì)算(針對(duì)N個(gè)去噪步驟)來(lái)減少推理延遲,同時(shí)保持生成質(zhì)量和模型靈活性。
實(shí)驗(yàn)
本節(jié)首先描述EasyControl的實(shí)現(xiàn)細(xì)節(jié),然后概述評(píng)估指標(biāo)。接下來(lái),展示實(shí)驗(yàn)結(jié)果,包括定性和定量分析,以及消融實(shí)驗(yàn)。
實(shí)現(xiàn)細(xì)節(jié)
采用FLUX.1 dev作為預(yù)訓(xùn)練的DiT。對(duì)于每個(gè)空間或主題條件訓(xùn)練,我們使用4個(gè)A100 GPU(80GB),每個(gè)GPU的批量大小為1,學(xué)習(xí)率為1e-4,訓(xùn)練100,000步。在推理期間,應(yīng)用流匹配采樣進(jìn)行25個(gè)采樣步驟。
實(shí)驗(yàn)設(shè)置
視覺(jué)比較: 我們?cè)u(píng)估以下設(shè)置:(1) 單條件生成,(2) 使用定制模型的單條件適應(yīng),(3) 多條件集成(如下圖3和下圖4所示),以及(4) 分辨率適應(yīng)性。定量比較:我們?cè)u(píng)估以下方面:(1) 單條件和雙條件生成下的推理時(shí)間和模型參數(shù)數(shù)量(以評(píng)估效率,如下表1所示),(2) 使用面部+OpenPose作為多條件的可控性、生成質(zhì)量和文本一致性,以及(3) 單條件設(shè)置下的可控性、生成質(zhì)量和文本一致性。
比較方法: 對(duì)于單條件,與Controlnet、OminiControl和Uni-ControlNet進(jìn)行比較。對(duì)于多條件設(shè)置,評(píng)估本文方法與幾個(gè)即插即用基線方法,包括Controlnet+IP-Adapter、Controlnet+Redux和Uni-Controlnet。還比較了與ControlNet集成的幾種ID定制方法[15, 35, 71]。
實(shí)驗(yàn)結(jié)果
定性比較
上圖3 (a)比較了不同方法在單控制條件下的性能。在Canny控制下,Uni-ControlNet和ControlNet表現(xiàn)出顏色不一致,導(dǎo)致與輸入文本偏離。在深度控制下,Uni-ControlNet未能生成連貫的圖像,而ControlNet和OmniControl引入了偽影,例如狗和沙發(fā)的融合。在OpenPose控制下,本文方法保留了文本渲染,而其他方法則削弱或失去了這種能力。在主題控制下,IP-Adapter和Uni-ControlNet未能與參考對(duì)齊??傮w而言,本文方法確保了文本一致性和在不同控制條件下的高質(zhì)量生成。
上圖3 (b)比較了不同方法在四個(gè)定制模型上生成圖像的即插即用能力。最左列顯示了來(lái)自LoRA微調(diào)的Flux.1 Dev模型的原始文本到圖像(T2I)結(jié)果。ControlNet和OmniControl都犧牲了風(fēng)格化,并遭受質(zhì)量下降。相比之下,本文方法展示了在不失去可控性的情況下最小化風(fēng)格化損失的能力,體現(xiàn)了我們方法的即插即用能力。
上圖4展示了不同方法在多條件控制下的視覺(jué)比較。對(duì)于OpenPose和面部控制,本文方法在身份一致性和可控性方面表現(xiàn)優(yōu)異。相比之下,其他方法在控制條件之間表現(xiàn)出沖突。雖然ControlNet和IP-Adapter的組合保持了可控性,但卻損害了身份一致性。
ControlNet+Redux和Uni-ControlNet未能同時(shí)保持身份一致性和可控性,這在主題-深度控制場(chǎng)景中也有所體現(xiàn)(右側(cè)第三/第四行)。對(duì)于OpenPose-Canny和Depth-Canny組合,本文方法和Uni-ControlNet都生成了符合控制條件的圖像。然而,Uni-ControlNet難以與文本輸入對(duì)齊,并產(chǎn)生質(zhì)量較低的圖像。Multi-ControlNet未能同時(shí)滿足兩個(gè)條件。這些結(jié)果展示了我們方法在無(wú)縫整合多種條件方面的靈活性。
定量比較
上表1展示了在單個(gè)A100 GPU上進(jìn)行20次采樣步驟時(shí)各種算法的推理時(shí)間和相應(yīng)的模型參數(shù)數(shù)量。在單條件設(shè)置下,本文完整模型實(shí)現(xiàn)了最佳性能,推理時(shí)間為16.3秒,比沒(méi)有位置感知訓(xùn)練范式(PATP)和KV緩存的版本減少了58%。值得注意的是,本文方法在保持最小參數(shù)數(shù)量15M的同時(shí)實(shí)現(xiàn)了這一效率,明顯低于ControlNet的3B參數(shù)。對(duì)于雙條件任務(wù),本文完整模型實(shí)現(xiàn)了18.3秒的推理時(shí)間,比沒(méi)有PATP和KV緩存的版本快75%。這一性能與ControlNet+IPA(16.8秒)競(jìng)爭(zhēng),同時(shí)保持了更小的模型大小(30M參數(shù)相比于ControlNet+IPA的4B)。結(jié)果突出了我們提出的PATP和KV緩存機(jī)制在提高推理效率方面的有效性,而不影響模型的緊湊性。
消融研究
在本文消融研究中,我們分析了去除各個(gè)模塊的影響。首先,將條件注入LoRA(CIL)替換為標(biāo)準(zhǔn)LoRA結(jié)構(gòu)(W.O. CIL)允許單條件控制,但無(wú)法以零樣本方式推廣到多條件控制。對(duì)于位置感知訓(xùn)練范式(PATP),我們訓(xùn)練了一個(gè)沒(méi)有PATP的模型,其中控制信號(hào)和噪聲固定在512×512分辨率,同時(shí)保持其他訓(xùn)練設(shè)置不變。該模型在生成高分辨率(例如1024×1024)或非正方形縱橫比(例如1024×768)圖像時(shí)表現(xiàn)出偽影和質(zhì)量下降。相比之下,基于PATP的訓(xùn)練有效地緩解了這些問(wèn)題。對(duì)于因果注意,去除因果互注意(CMA)仍允許圖像生成,這得益于注意力的自適應(yīng)性質(zhì)。然而,條件之間的沖突降低了控制精度,導(dǎo)致諸如在多控制場(chǎng)景中改變?nèi)梭w姿勢(shì)和移動(dòng)物體位置(例如月亮)等偏差。當(dāng)所有模塊一起使用時(shí),本文方法實(shí)現(xiàn)了最高的可控性、生成質(zhì)量以及對(duì)不同分辨率和縱橫比的適應(yīng)性。
結(jié)論
EasyControl,一個(gè)高效且靈活的統(tǒng)一條件引導(dǎo)擴(kuò)散模型框架。本文框架利用了三個(gè)關(guān)鍵創(chuàng)新:(1) 一個(gè)輕量級(jí)的條件注入LoRA模塊,能夠無(wú)縫整合多樣的條件信號(hào)而不改變核心模型的功能。(2) 一個(gè)位置感知訓(xùn)練范式,確保對(duì)各種分辨率和縱橫比的適應(yīng)性。(3) 一個(gè)新穎的因果注意機(jī)制結(jié)合KV緩存技術(shù),顯著提高了效率。這些組件共同解決了可控圖像生成中的效率和靈活性挑戰(zhàn)。EasyControl在廣泛的視覺(jué)任務(wù)中實(shí)現(xiàn)了強(qiáng)大的可控性和高質(zhì)量結(jié)果。廣泛的實(shí)驗(yàn)展示了其處理復(fù)雜的多條件場(chǎng)景的能力,同時(shí)擴(kuò)展到多樣的分辨率和縱橫比。本文框架為條件圖像生成提供了一個(gè)強(qiáng)大且可適應(yīng)的解決方案。
本文轉(zhuǎn)自AI生成未來(lái) ,作者:AI生成未來(lái)
原文鏈接:??https://mp.weixin.qq.com/s/uO55WYUYM08kLjMz7QxmyQ??
