标签:

• 文本到图像
• KOALA 数据集:
• Ejafa/ye-pop

       

KOALA-Lightning-700M 模型卡片

摘要

• 采用基于自注意力的知识蒸馏方法训练
• 教师模型：SDXL-Lightning
• 训练数据集：LAION-POP 子集
• 训练迭代：50万次，批次大小128
• 训练GPU：4块NVIDIA A100（80GB）

KOALA-Lightning 模型系列

摘要

一句话总结

我们通过压缩SDXL的U-Net并将其知识蒸馏到新模型中，提出了名为KOALA的快速文本生成图像模型。KOALA-Lightning-700M在NVIDIA 4090 GPU上仅需0.66秒即可生成1024x1024图像，比SDXL快4倍以上。在资源有限场景下，KOALA-700M可作为SDM与SDXL之间的高性价比替代方案。

以下1024x1024样本由KOALA-Lightning-700M在NVIDIA 4090 GPU上通过10次去噪步骤生成，耗时仅0.66秒。

架构设计

通过减少残差块和Transformer块，我们实现了KOALA-1B和KOALA-700M两种压缩U-Net架构。

U-Net对比表

• 测试环境：FP16精度，25次去噪步骤，NVIDIA 4090 GPU(24GB)
• SDM-v2.0使用768x768分辨率，SDXL和KOALA使用1024x1024分辨率

不同GPU上的延迟与显存占用对比

我们在多种消费级GPU上测试了512x512(SDXL-Turbo/KOALA-Turbo)和1024x1024(其他模型)的推理时间。"OOM"表示显存不足。值得注意的是，KOALA模型可在所有GPU类型上运行，而SDXL无法在8GB显存的3060Ti上工作。

核心特性

• 高效U-Net架构：相比SDXL，模型体积分别减小54%和69%
• 自注意力知识蒸馏：通过重点蒸馏自注意力特征保持生成质量

模型信息

• 开发团队：ETRI视觉智能实验室
• 核心开发者：李永完、朴宽勇、赵裕林、李英珠、黄成柱
• 模型类型：基于潜在扩散的文本生成图像模型，沿用SDXL文本编码器，仅替换去噪U-Net
• 更多资源：arXiv论文 | 项目主页

使用示例（🤗Diffusers库）

25步去噪推理代码：

import torch
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler

pipe = StableDiffusionXLPipeline.from_pretrained("etri-vilab/koala-lightning-700m", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# 确保采样器使用"trailing"时间步和"sample"预测类型
pipe.scheduler = EulerDiscreteScheduler.from_config(
    pipe.scheduler.config, timestep_spacing="trailing"
  )

prompt = "达芬奇风格的金毛猎犬肖像画"
negative = "低质量，插画，低分辨率"
image = pipe(prompt=prompt, negative_prompt=negative, guidance_scale=3.5, num_inference_steps=10).images[0]

应用范围

直接用途

本模型仅限研究用途，包括：

• 艺术作品生成与设计应用
• 教育/创意工具开发
• 生成模型研究
• 安全部署研究
• 模型局限性探索

非适用场景

本模型不适用于生成需要事实准确的人物/事件表征内容。

局限性

• 文本渲染：长文本可读性有待提升
• 复杂提示：对多属性组合提示处理不足
• 数据依赖：受限于LAION-aesthetics-V2 6+数据集特性

引用

@misc{Lee@koala,
    title={KOALA: 面向高效内存与快速推理的文本生成图像扩散模型实证研究}, 
    author={李永完 and 朴宽勇 and 赵裕林 and 李英珠 and 黄成柱},
    year={2023},
    eprint={2312.04005},
    archivePrefix={arXiv},
    primaryClass={cs.CV}
}