近年來，基于擴(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)、個(gè)性化訓(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)配各個(gè)組件和算法，大幅提高了框架的靈活性和可擴(kuò)展性。開發(fā)者像搭積木一樣組合算法，而無需重復(fù)實(shí)現(xiàn)代碼細(xì)節(jié)。
比如，基于 HCP-Diffusion，我們可以通過簡單地修改配置文件即可完成 LoRA，DreamArtist，Contro.NET 等多種常見算法的部署與組合。這不僅降低了創(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 支持定義多個(gè)并行數(shù)據(jù)集，每個(gè)數(shù)據(jù)集可采用不同的圖像尺寸與標(biāo)注格式，每次訓(xùn)練迭代會(huì)從每個(gè)數(shù)據(jù)集中各抽取一個(gè) batch 進(jìn)行訓(xùn)練，如圖 2 所示。此外，每個(gè)數(shù)據(jù)集可配置多種數(shù)據(jù)源，支持 txt、json、yaml 等標(biāo)注格式或自定義標(biāo)注格式，具有高度靈活的數(shù)據(jù)預(yù)處理與加載機(jī)制。
圖 2 數(shù)據(jù)集結(jié)構(gòu)示意圖
數(shù)據(jù)集處理部分提供帶自動(dòng)聚類的 aspect ratio bucket，支持處理圖像尺寸各異的數(shù)據(jù)集。用戶無需對(duì)數(shù)據(jù)集尺寸做額外處理和對(duì)齊，框架會(huì)根據(jù)寬高比或分辨率自動(dòng)選擇最優(yōu)的分組方式。該技術(shù)大幅降低數(shù)據(jù)處理的門檻，優(yōu)化用戶體驗(yàn)，使開發(fā)者更專注于算法本身的創(chuàng)新。
而對(duì)于圖像數(shù)據(jù)的預(yù)處理，框架也兼容 torch vision, albumentations 等多種圖像處理庫。用戶可以根據(jù)需要在配置文件中直接配置預(yù)處理方式，或是在此基礎(chǔ)上拓展自定義的圖像處理方法。
HCP-Diffusion 在文本標(biāo)注方面，設(shè)計(jì)了靈活且清晰的 prompt 模板規(guī)范，可支持復(fù)雜多樣的訓(xùn)練方法與數(shù)據(jù)標(biāo)注。其對(duì)應(yīng)用上述配置文件 source 目錄下的 word_names，里面可自定義下圖大括號(hào)中的特殊字符對(duì)應(yīng)的嵌入詞向量與類別描述，以與 DreamBooth、DreamArtist 等模型兼容。
并且對(duì)于文本標(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 對(duì)生成條件進(jìn)行編碼，Adapter 微調(diào)模型與下游任務(wù)對(duì)齊，positive 與 negative 雙通道代表正負(fù)條件對(duì)圖像的控制生成。
如圖 5 所示，HCP-Diffusion 在配置文件中通過簡易的組合，即可實(shí)現(xiàn) LoRA、ControlNet、DreamArtist 等多種主流訓(xùn)練算法。同時(shí)支持對(duì)上述算法進(jìn)行組合，例如 LoRA 和 Textual Inversion 同時(shí)訓(xùn)練，為 LoRA 綁定專有觸發(fā)詞等。此外，通過插件模塊，可以輕松自定義任意插件，業(yè)已兼容目前所有主流方法接入。通過上述的模塊化，HCP-Diffusion 實(shí)現(xiàn)了對(duì)任意主流算法的框架搭建，降低了開發(fā)門檻，促進(jìn)了模型的協(xié)同創(chuàng)新。
HCP-Diffusion 將 LoRA、ControlNet 等各種 Adapter 類算法統(tǒng)一抽象為模型插件，通過定義一些通用的模型插件基類，可以將所有這類算法統(tǒng)一對(duì)待，降低用戶使用成本和開發(fā)成本，將所有 Adapter 類算法統(tǒng)一。
框架提供四種類型的插件，可以輕松支持目前所有主流算法:
+ SinglePluginBlock: 單層插件，根據(jù)該層輸入改變輸出，比如 lora 系列。支持正則表達(dá)式 (re: 前綴) 定義插入層， 不支持 pre_hook: 前綴。
+ PluginBlock: 輸入層和輸出層都只有一個(gè)，比如定義殘差連接。支持正則表達(dá)式 (re: 前綴) 定義插入層， 輸入輸出層都支持 pre_hook: 前綴。
+ MultiPluginBlock: 輸入層和輸出層都可以有多個(gè)，比如 controlnet。不支持正則表達(dá)式 (re: 前綴)， 輸入輸出層都支持 pre_hook: 前綴。
+ WrAppluginBlock: 替換原有模型的某個(gè)層，將原有模型的層作為該類的一個(gè)對(duì)象。支持正則表達(dá)式 (re: 前綴) 定義替換層，不支持 pre_hook: 前綴。
訓(xùn)練、推理模塊
圖 6 自定義優(yōu)化器配置
HCP-Diffusion 中的配置文件支持定義 Python/ target=_blank class=infotextkey>Python 對(duì)象，運(yùn)行時(shí)自動(dòng)實(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)練時(shí)的顯存占用，加快訓(xùn)練速度。支持 EMA 操作，可以進(jìn)一步提高模型的生成效果和泛化性。在推理階段，支持模型 offload 和 VAE tiling 等操作，最低僅需 1GB 顯存即可完成圖像生成。
通過上述簡單的文件配置，即可無需耗費(fèi)大量精力查找相關(guān)框架資源完成模型的配置，如上圖所示。HCP-Diffusion 模塊化的設(shè)計(jì)方式，將模型方法定義，訓(xùn)練邏輯，推理邏輯等完全分離，配置模型時(shí)無需考慮訓(xùn)練與推理部分的邏輯，幫助用戶更好的聚焦于方法本身。同時(shí)，HCP-Diffusion 已經(jīng)提供大多數(shù)主流算法的框架配置樣例，只需對(duì)其中部分參數(shù)進(jìn)行修改，就可以實(shí)現(xiàn)部署。
HCP-Diffusion：Web UI 圖像界面
除了可直接修改配置文件，HCP-Diffusion 已提供了對(duì)應(yīng)的 Web UI 圖像界面，包含圖像生成，模型訓(xùn)練等多個(gè)模塊，以提升用戶體驗(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)外科技獎(jiǎng)項(xiàng)及最佳論文獎(jiǎng)，并致力于打造產(chǎn)品級(jí)的AI技術(shù)及平臺(tái)。實(shí)驗(yàn)室網(wǎng)站：http://www.sysu-hcp.net
