标签:
- 文本生成图像
- KOALA模型
数据集:
- Ejafa/ye-pop
KOALA-Lightning-1B 模型卡片
概述
KOALA-Lightning系列模型
模型 |
链接 |
koala-lightning-1b |
https://huggingface.co/etri-vilab/koala-lightning-1b |
koala-lightning-700m |
https://huggingface.co/etri-vilab/koala-lightning-700m |
摘要
一句话总结
我们通过压缩SDXL的U-Net结构并蒸馏其知识,提出了快速文本生成图像模型KOALA。KOALA-Lightning-700M在NVIDIA 4090显卡上仅需0.66秒即可生成1024x1024图像,比SDXL快4倍以上。在资源受限场景下,KOALA-700M可作为SDM与SDXL之间的高性价比替代方案。
完整摘要
随着文本生成图像模型规模扩大,其对昂贵大显存GPU的依赖导致推理成本剧增,加之训练数据获取受限,使模型复现面临挑战。本研究旨在降低推理成本,探索仅使用公开数据集和开源模型能实现的生成能力边界。基于业界标准模型Stable Diffusion XL(SDXL),我们提出构建高效T2I模型的三大实践:(1) 知识蒸馏:发现自注意力机制是蒸馏SDXL生成能力至高效U-Net的关键;(2) 数据策略:少量高分辨率带丰富标注的图像比大量低分辨率简短标注图像更有效;(3) 教师模型:分步蒸馏教师模型可减少去噪步骤。基于这些发现,我们构建了KOALA-Turbo和-Lightning两种高效文本生成图像模型,其压缩U-Net(1B/700M参数)体积仅为SDXL U-Net的54%和69%。其中KOALA-Lightning-700M在保持满意生成质量前提下,速度比SDXL快4倍。与SDXL不同,我们的模型可在8GB显存的消费级GPU(如3060Ti)上生成1024px高清图像。相信KOALA模型将为资源受限环境中的学术研究者和普通用户提供高性价比的SDXL替代方案。
以下1024x1024样本由KOALA-Lightning-700M在NVIDIA 4090显卡上通过10步去噪生成,耗时仅0.66秒。
架构设计
通过减少残差块和Transformer块,我们设计了两种压缩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 |
Transformer块配置 |
[1,1,1,1] |
[0,2,10] |
[0,2,6] |
[0,2,5] |
中间块 |
✓ |
✓ |
✓ |
✗ |
延迟时间 |
1.131秒 |
3.133秒 |
1.604秒 |
1.257秒 |
- "Tx"表示Transformer模块,"CKPT"指训练好的模型文件
- 延迟测试采用FP16精度,25步去噪,NVIDIA 4090显卡(24GB显存)
- SDM-v2.0使用768x768分辨率,SDXL与KOALA模型使用1024x1024分辨率
不同GPU上的延迟与显存占用对比
我们在NVIDIA 3060Ti(8GB)、2080Ti(11GB)和4090(24GB)多款消费级GPU上测试了各模型性能。"OOM"表示显存不足。值得注意的是,SDXL系列无法在8GB显存的3060Ti上运行,而我们的KOALA模型可适配全部GPU类型。
核心特性
- 高效U-Net架构:相比SDXL,KOALA模型U-Net体积分别减小54%和69%
- 自注意力知识蒸馏:通过重点蒸馏自注意力特征保持图像生成质量
模型信息