近年來，基于擴(kuò)散模型（Diffusion Models）的圖像生成模型層出不窮，展現(xiàn)出令人驚艷的生成效果。然而，現(xiàn)有相關(guān)研究模型代碼框架存在過度碎片化的問題，缺乏統(tǒng)一的框架體系，導(dǎo)致出現(xiàn)「遷移難」、「門檻高」、「質(zhì)量差」的代碼實(shí)現(xiàn)難題。

為此，中山大學(xué)人機(jī)物智能融合實(shí)驗(yàn)室（HCP Lab）構(gòu)建了 HCP-Diffusion 框架，系統(tǒng)化地實(shí)現(xiàn)了模型微調(diào)、個性化訓(xùn)練、推理優(yōu)化、圖像編輯等基于 Diffusion 模型的相關(guān)算法，結(jié)構(gòu)如圖 1 所示。

圖 1 HCP-Diffusion 框架結(jié)構(gòu)圖，通過統(tǒng)一框架統(tǒng)一現(xiàn)有 diffusion 相關(guān)方法，提供多種模塊化的訓(xùn)練與推理優(yōu)化方法。

HCP-Diffusion 通過格式統(tǒng)一的配置文件調(diào)配各個組件和算法，大幅提高了框架的靈活性和可擴(kuò)展性。開發(fā)者像搭積木一樣組合算法，而無需重復(fù)實(shí)現(xiàn)代碼細(xì)節(jié)。

比如，基于 HCP-Diffusion，我們可以通過簡單地修改配置文件即可完成 LoRA，DreamArtist，ControlNet 等多種常見算法的部署與組合。這不僅降低了創(chuàng)新的門檻，也使得框架可以兼容各類定制化設(shè)計(jì)。

• HCP-Diffusion 代碼工具：https://github.com/7eu7d7/HCP-Diffusion
• HCP-Diffusion 圖形界面：https://github.com/7eu7d7/HCP-Diffusion-webui

HCP-Diffusion：功能模塊介紹

框架特色

HCP-Diffusion 通過將目前主流的 diffusion 訓(xùn)練算法框架模塊化，實(shí)現(xiàn)了框架的通用性，主要特色如下：

• 統(tǒng)一架構(gòu)：搭建 Diffusion 系列模型統(tǒng)一代碼框架
• 算子插件：支持?jǐn)?shù)據(jù)、訓(xùn)練、推理、性能優(yōu)化等算子算法，如 deepspeed, colossal-AI 和 offload 等加速優(yōu)化
• 一鍵配置：Diffusion 系列模型可通過高靈活度地修改配置文件即可完成模型實(shí)現(xiàn)
• 一鍵訓(xùn)練：提供 Web UI，一鍵訓(xùn)練、推理

數(shù)據(jù)模塊

HCP-Diffusion 支持定義多個并行數(shù)據(jù)集，每個數(shù)據(jù)集可采用不同的圖像尺寸與標(biāo)注格式，每次訓(xùn)練迭代會從每個數(shù)據(jù)集中各抽取一個 batch 進(jìn)行訓(xùn)練，如圖 2 所示。此外，每個數(shù)據(jù)集可配置多種數(shù)據(jù)源，支持 txt、json、yaml 等標(biāo)注格式或自定義標(biāo)注格式，具有高度靈活的數(shù)據(jù)預(yù)處理與加載機(jī)制。

圖 2 數(shù)據(jù)集結(jié)構(gòu)示意圖

數(shù)據(jù)集處理部分提供帶自動聚類的 aspect ratio bucket，支持處理圖像尺寸各異的數(shù)據(jù)集。用戶無需對數(shù)據(jù)集尺寸做額外處理和對齊，框架會根據(jù)寬高比或分辨率自動選擇最優(yōu)的分組方式。該技術(shù)大幅降低數(shù)據(jù)處理的門檻，優(yōu)化用戶體驗(yàn)，使開發(fā)者更專注于算法本身的創(chuàng)新。

而對于圖像數(shù)據(jù)的預(yù)處理，框架也兼容 torch vision, albumentations 等多種圖像處理庫。用戶可以根據(jù)需要在配置文件中直接配置預(yù)處理方式，或是在此基礎(chǔ)上拓展自定義的圖像處理方法。

圖 3 數(shù)據(jù)集配置文件示例

HCP-Diffusion 在文本標(biāo)注方面，設(shè)計(jì)了靈活且清晰的 prompt 模板規(guī)范，可支持復(fù)雜多樣的訓(xùn)練方法與數(shù)據(jù)標(biāo)注。其對應(yīng)用上述配置文件 source 目錄下的 word_names，里面可自定義下圖大括號中的特殊字符對應(yīng)的嵌入詞向量與類別描述，以與 DreamBooth、DreamArtist 等模型兼容。

圖 4 prompt 模板

并且對于文本標(biāo)注，也提供了按句擦除 (TagDropout) 或按句打亂 (TagShuffle) 等多種文本增強(qiáng)方法，可以減少圖像與文本數(shù)據(jù)間的過擬合問題，使生成的圖像更多樣化。

模型框架模塊

HCP-Diffusion 通過將目前主流的 diffusion 訓(xùn)練算法框架模塊化，實(shí)現(xiàn)了框架的通用性。具體而言，Image Encoder，Image Decoder 完成圖像的編解碼，Noise Generator 產(chǎn)生前向過程的噪聲，Diffusion Model 實(shí)現(xiàn)擴(kuò)散過程，Condition Encoder 對生成條件進(jìn)行編碼，Adapter 微調(diào)模型與下游任務(wù)對齊，positive 與 negative 雙通道代表正負(fù)條件對圖像的控制生成。

圖 5 模型結(jié)構(gòu)示例配置 (模型插件，自定義單詞等)

如圖 5 所示，HCP-Diffusion 在配置文件中通過簡易的組合，即可實(shí)現(xiàn) LoRA、ControlNet、DreamArtist 等多種主流訓(xùn)練算法。同時支持對上述算法進(jìn)行組合，例如 LoRA 和 Textual Inversion 同時訓(xùn)練，為 LoRA 綁定專有觸發(fā)詞等。此外，通過插件模塊，可以輕松自定義任意插件，業(yè)已兼容目前所有主流方法接入。通過上述的模塊化，HCP-Diffusion 實(shí)現(xiàn)了對任意主流算法的框架搭建，降低了開發(fā)門檻，促進(jìn)了模型的協(xié)同創(chuàng)新。

HCP-Diffusion 將 LoRA、ControlNet 等各種 Adapter 類算法統(tǒng)一抽象為模型插件，通過定義一些通用的模型插件基類，可以將所有這類算法統(tǒng)一對待，降低用戶使用成本和開發(fā)成本，將所有 Adapter 類算法統(tǒng)一。

框架提供四種類型的插件，可以輕松支持目前所有主流算法:

+ SinglePluginBlock: 單層插件，根據(jù)該層輸入改變輸出，比如 lora 系列。支持正則表達(dá)式 (re: 前綴) 定義插入層， 不支持 pre_hook: 前綴。

+ PluginBlock: 輸入層和輸出層都只有一個，比如定義殘差連接。支持正則表達(dá)式 (re: 前綴) 定義插入層， 輸入輸出層都支持 pre_hook: 前綴。

+ MultiPluginBlock: 輸入層和輸出層都可以有多個，比如 controlnet。不支持正則表達(dá)式 (re: 前綴)， 輸入輸出層都支持 pre_hook: 前綴。

+ WrapPluginBlock: 替換原有模型的某個層，將原有模型的層作為該類的一個對象。支持正則表達(dá)式 (re: 前綴) 定義替換層，不支持 pre_hook: 前綴。

訓(xùn)練、推理模塊

圖 6 自定義優(yōu)化器配置

HCP-Diffusion 中的配置文件支持定義 python 對象，運(yùn)行時自動實(shí)例化。該設(shè)計(jì)使得開發(fā)者可以輕松接入任何 pip 可安裝的自定義模塊，例如自定義優(yōu)化器，損失函數(shù)，噪聲采樣器等，無需修改框架代碼，如上圖所示。配置文件結(jié)構(gòu)清晰，易于理解，可復(fù)現(xiàn)性強(qiáng)，有助于平滑連接學(xué)術(shù)研究和工程部署。

加速優(yōu)化支持

HCP-Diffusion 支持 Accelerate、DeepSpeed、Colossal-AI 等多種訓(xùn)練優(yōu)化框架，可以顯著減少訓(xùn)練時的顯存占用，加快訓(xùn)練速度。支持 EMA 操作，可以進(jìn)一步提高模型的生成效果和泛化性。在推理階段，支持模型 offload 和 VAE tiling 等操作，最低僅需 1GB 顯存即可完成圖像生成。

圖 7 模塊化配置文件

通過上述簡單的文件配置，即可無需耗費(fèi)大量精力查找相關(guān)框架資源完成模型的配置，如上圖所示。HCP-Diffusion 模塊化的設(shè)計(jì)方式，將模型方法定義，訓(xùn)練邏輯，推理邏輯等完全分離，配置模型時無需考慮訓(xùn)練與推理部分的邏輯，幫助用戶更好的聚焦于方法本身。同時，HCP-Diffusion 已經(jīng)提供大多數(shù)主流算法的框架配置樣例，只需對其中部分參數(shù)進(jìn)行修改，就可以實(shí)現(xiàn)部署。

HCP-Diffusion：Web UI 圖像界面

除了可直接修改配置文件，HCP-Diffusion 已提供了對應(yīng)的 Web UI 圖像界面，包含圖像生成，模型訓(xùn)練等多個模塊，以提升用戶體驗(yàn)，大幅降低框架的學(xué)習(xí)門檻，加速算法從理論到實(shí)踐的轉(zhuǎn)化。

圖 8 HCP-Diffusion Web UI 圖像界面

實(shí)驗(yàn)室簡介

中山大學(xué)人機(jī)物智能融合實(shí)驗(yàn)室 (HCP Lab) 由林倞教授于 2010 年創(chuàng)辦，近年來在多模態(tài)內(nèi)容理解、因果及認(rèn)知推理、具身學(xué)習(xí)等方面取得豐富學(xué)術(shù)成果，數(shù)次獲得國內(nèi)外科技獎項(xiàng)及最佳論文獎，并致力于打造產(chǎn)品級的AI技術(shù)及平臺。實(shí)驗(yàn)室網(wǎng)站：http://www.sysu-hcp.net