🚀 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}
}