英偉達(dá)開(kāi)源4K圖像生成模型Sana,可在16G顯存電腦部署,支持ComfyUI和LoRA訓(xùn)練
英偉達(dá)開(kāi)源了一個(gè)可以直接生成 4K 圖片的模型 Sana。 Sana-0.6B 可以在 16GB 的筆記本電腦 GPU 上部署。生成 1024 × 1024 分辨率的圖像只需不到 1 秒鐘。官方已經(jīng)支持了 Comfyui,而且放出了 Lora 訓(xùn)練工具。
相關(guān)鏈接
- 論文:https://arxiv.org/pdf/2410.10629
- 主頁(yè):https://nvlabs.github.io/Sana
論文介紹
Sana是一個(gè)文本到圖像的框架,可以高效地生成分辨率高達(dá) 4096 × 4096 的圖像。Sana 可以以極快的速度合成具有強(qiáng)大文本-圖像對(duì)齊功能的高分辨率、高質(zhì)量圖像,可在筆記本電腦的 GPU 上部署。
核心設(shè)計(jì)包括:
- 深度壓縮自動(dòng)編碼器:與僅將圖像壓縮 8 倍的傳統(tǒng) AE 不同,我們訓(xùn)練了一個(gè)可以將圖像壓縮 32 倍的 AE,從而有效減少了潛在標(biāo)記的數(shù)量。 - 線性 DiT:我們用線性注意力取代了 DiT 中的所有原始注意力,這在高分辨率下效率更高,且不會(huì)犧牲質(zhì)量。
- 僅解碼器的文本編碼器:我們用現(xiàn)代僅解碼器的小型 LLM 取代了 T5 作為文本編碼器,并設(shè)計(jì)了具有上下文學(xué)習(xí)的復(fù)雜人工指令以增強(qiáng)圖像-文本對(duì)齊。 高效的訓(xùn)練和采樣:我們提出了 Flow-DPM-Solver 來(lái)減少采樣步驟,并通過(guò)高效的字幕標(biāo)記和選擇來(lái)加速收斂。
因此,Sana-0.6B 與現(xiàn)代巨型擴(kuò)散模型(例如 Flux-12B)相比極具競(jìng)爭(zhēng)力,其體積小 20 倍,測(cè)量吞吐量快 100 多倍。此外,Sana-0.6B 可以部署在 16GB 筆記本電腦 GPU 上,生成 1024 × 1024 分辨率圖像只需不到 1 秒。Sana 可以以低成本實(shí)現(xiàn)內(nèi)容創(chuàng)建。
提高效率的幾個(gè)核心設(shè)計(jì)細(xì)節(jié)
- 深度壓縮自動(dòng)編碼器: 我們引入了一種新的 深度壓縮自動(dòng)編碼器 (DC-AE),將縮放因子大幅增加到 32。與 AE-F8 相比,我們的 AE-F32 輸出的潛在標(biāo)記減少了 16 倍,這對(duì)于高效訓(xùn)練和生成超高分辨率圖像(例如 4K 分辨率)至關(guān)重要。
- 高效的線性 DiT: 我們引入了一種新的線性 DiT,取代了香草二次注意力,并將復(fù)雜度從 O(N2) 降低到O(N)。Mix-FFN 在 MLP 中使用 3×3 深度卷積,增強(qiáng)了 token 的局部信息。線性注意力實(shí)現(xiàn)與香草相當(dāng)?shù)慕Y(jié)果,將 4K 生成的延遲提高了 1.7 倍。Mix-FFN 還消除了位置編碼 (NoPE) 的需要,并且沒(méi)有質(zhì)量損失,標(biāo)志著第一個(gè)沒(méi)有位置嵌入的 DiT。
- 僅解碼器的小型 LLM 作為文本編碼器: 我們使用僅解碼器的 LLM Gemma 作為文本編碼器,以增強(qiáng)提示中的理解和推理。與 CLIP 或 T5 不同,Gemma 提供了卓越的文本理解和指令遵循能力。我們解決了訓(xùn)練不穩(wěn)定性問(wèn)題,并設(shè)計(jì)了復(fù)雜的人工指令 (CHI) 來(lái)利用 Gemma 的上下文學(xué)習(xí),改善圖像-文本對(duì)齊。
- 高效的訓(xùn)練和推理策略: 我們提出了自動(dòng)標(biāo)記和訓(xùn)練策略來(lái)提高文本與圖像的一致性。多個(gè) VLM 生成不同的重新字幕,基于 CLIPScore 的策略選擇高 CLIPScore 字幕以增強(qiáng)收斂和對(duì)齊。此外,與 Flow-Euler-Solver 相比,我們的Flow-DPM-Solver將推理步驟從 28-50 減少到 14-20,性能更佳。
總體表現(xiàn)
我們?cè)诒?1 中將 Sana 與最先進(jìn)的文本到圖像擴(kuò)散模型進(jìn)行了比較。對(duì)于 512 × 512 分辨率,Sana-0.6 的吞吐量比具有相似模型大小的 PixArt-Σ 快 5 倍,并且在 FID、Clip Score、GenEval 和 DPG-Bench 中的表現(xiàn)明顯優(yōu)于它。對(duì)于 1024 × 1024 分辨率,Sana 比大多數(shù)具有 <3B 參數(shù)的模型要強(qiáng)得多,并且在推理延遲方面表現(xiàn)出色。即使與最先進(jìn)的大型模型 FLUX-dev 相比,我們的模型也能實(shí)現(xiàn)具有競(jìng)爭(zhēng)力的性能。例如,雖然 DPG-Bench 上的準(zhǔn)確率相當(dāng),GenEval 上的準(zhǔn)確率略低,但 Sana-0.6B 的吞吐量快 39 倍,Sana-1.6B 快 23 倍。
ComfyUI 使用
作者已經(jīng)開(kāi)發(fā)了插件來(lái)將 Sana 與 ComfyUI 集成。
- 插件:https://github.com/Efficient-Large-Model/ComfyUI_ExtraModels
- GitHub:https://github.com/NVlabs/Sana/blob/main/asset/docs/ComfyUI/comfyui.md
- ComfyUI:https://github.com/NVlabs/Sana/blob/main/asset/docs/ComfyUI/comfyui.md
Sana-LoRA Dreambooth
Sana-LoRA 由 diffusers 支持。查看我們的指南來(lái)訓(xùn)練您的自定義模型。我們?cè)谙旅嬲故玖?Sana-LoRA 微調(diào)過(guò)程中的一些示例。
- 鏈接:https://github.com/NVlabs/Sana/blob/main/asset/docs/sana_lora_dreambooth.md