🚀 KOALA-Lightning-1B 模型卡片
KOALA-Lightning-1B 是基於潛在擴散的文本到圖像生成模型,通過自注意力知識蒸餾方法訓練,能有效降低推理成本,在圖像生成效率和質量上表現出色。
🚀 快速開始
模型信息
屬性 |
詳情 |
模型類型 |
基於潛在擴散的文本到圖像生成模型 |
訓練數據 |
LAION-POP 數據集的一個子集 |
訓練輪次 |
500K,批次大小為 128 |
訓練 GPU |
4 x NVIDIA A100 (80GB) |
教師模型 |
SDXL-Lightning |
推理代碼示例
使用 🤗Diffusers 庫進行推理,去噪步驟為 25:
import torch
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
pipe = StableDiffusionXLPipeline.from_pretrained("etri-vilab/koala-lightning-1b", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
pipe.scheduler = EulerDiscreteScheduler.from_config(
pipe.scheduler.config, timestep_spacing="trailing"
)
prompt = "A portrait painting of a Golden Retriever like Leonard da Vinci"
negative = "worst quality, low quality, illustration, low resolution"
image = pipe(prompt=prompt, negative_prompt=negative, guidance_scale=3.5, num_inference_steps=10).images[0]
✨ 主要特性
- 高效的 U-Net 架構:KOALA 模型採用簡化的 U-Net 架構,與前身 Stable Diffusion XL (SDXL) 相比,模型大小分別最多減少 54% 和 69%。
- 基於自注意力的知識蒸餾:KOALA 的核心技術專注於自注意力特徵的蒸餾,這對保持圖像生成質量至關重要。
📦 安裝指南
文檔未提供具體安裝步驟,可參考 Diffusers 庫 的安裝說明。
💻 使用示例
基礎用法
上述推理代碼即為基礎用法示例,可根據需求修改提示詞 prompt
和負提示詞 negative
來生成不同的圖像。
高級用法
文檔未提供高級用法的代碼示例,可根據實際需求調整模型參數,如 guidance_scale
、num_inference_steps
等,以獲得不同效果的圖像。
📚 詳細文檔
模型概述
KOALA 模型由 ETRI 視覺智能實驗室 開發,開發者包括 Youngwan Lee、Kwanyong Park 等。該模型使用與 SDXL-Base-1.0 相同的文本編碼器,僅用壓縮後的 U-Net 替換去噪 U-Net。
模型架構
有兩種壓縮 U-Net 類型,KOALA-1B 和 KOALA-700M,通過減少殘差塊和變壓器塊實現。
U-Net 對比
U-Net |
SDM-v2.0 |
SDXL-Base-1.0 |
KOALA-1B |
KOALA-700M |
參數數量 |
865M |
2,567M |
1,161M |
782M |
檢查點大小 |
3.46GB |
10.3GB |
4.4GB |
3.0GB |
變壓器塊 |
[1, 1, 1, 1] |
[0, 2, 10] |
[0, 2, 6] |
[0, 2, 5] |
中間塊 |
✓ |
✓ |
✓ |
✗ |
延遲 |
1.131s |
3.133s |
1.604s |
1.257s |
- Tx 表示變壓器塊,CKPT 表示訓練後的檢查點文件。
- 延遲在 NVIDIA 4090 GPU (24GB) 上以 FP16 精度和 25 個去噪步驟測量。
- SDM-v2.0 使用 768x768 分辨率,而 SDXL 和 KOALA 模型使用 1024x1024 分辨率。
不同 GPU 上的延遲和內存使用對比
在多種消費級 GPU 上測量了 SDXL-Turbo 和 KOALA-Turbo 模型在 512x512 分辨率下的推理時間,以及其他模型在 1024x1024 分辨率下的推理時間。'OOM' 表示內存不足。需要注意的是,SDXL 模型無法在 8GB VRAM 的 3060Ti 上運行,而 我們的 KOALA 模型可以在所有類型的 GPU 上運行。
模型用途
直接使用
該模型僅用於研究目的,可能的研究領域和任務包括:
- 藝術品生成,以及在設計和其他藝術過程中的應用。
- 教育或創意工具中的應用。
- 生成模型的研究。
- 對有潛在生成有害內容的模型進行安全部署。
- 探索和理解生成模型的侷限性和偏差。
超出範圍的使用
該模型並非用於真實或準確地表示人物或事件,因此使用該模型生成此類內容超出了其能力範圍。
模型侷限性和偏差
- 文本渲染:模型在圖像中渲染長而清晰的文本時面臨挑戰。
- 複雜提示:KOALA 有時難以處理涉及多個屬性的複雜提示。
- 數據集依賴:當前的侷限性部分歸因於訓練數據集 (LAION-aesthetics-V2 6+) 的特性。
🔧 技術細節
訓練方法
使用 基於自注意力的知識蒸餾 方法進行訓練,探索如何有效地將 SDXL 的生成能力蒸餾到高效的 U-Net 中,並發現自注意力是最關鍵的部分。
數據選擇
儘管樣本較少,但具有豐富描述的高分辨率圖像比大量具有簡短描述的低分辨率圖像更為關鍵。
教師模型
分步蒸餾的教師模型允許文本到圖像模型減少噪聲步驟。
📄 許可證
文檔未提供許可證信息。
📖 引用
@misc{Lee@koala,
title={KOALA: Empirical Lessons Toward Memory-Efficient and Fast Diffusion Models for Text-to-Image Synthesis},
author={Youngwan Lee and Kwanyong Park and Yoorhim Cho and Yong-Ju Lee and Sung Ju Hwang},
year={2023},
eprint={2312.04005},
archivePrefix={arXiv},
primaryClass={cs.CV}
}