🚀 KOALA-Lightning-700M模型卡

KOALA-Lightning-700M是基于潜在扩散的文本到图像生成模型，通过自注意力知识蒸馏方法训练，有效减小模型大小，提升图像生成效率，可在多种消费级GPU上运行。

       

🚀 快速开始

KOALA-Lightning-700M是一个基于潜在扩散的文本到图像生成模型。以下是使用🤗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")

# Ensure sampler uses "trailing" timesteps and "sample" prediction type.
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的核心技术聚焦于自注意力特征的蒸馏，这对于保持图像生成质量至关重要。

📦 安装指南

文档未提供具体安装步骤，暂不展示。

💻 使用示例

基础用法

import torch
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler

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

# Ensure sampler uses "trailing" timesteps and "sample" prediction type.
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]

高级用法

文档未提供高级用法代码示例，暂不展示。

📚 详细文档

模型概述

• 由ETRI视觉智能实验室开发。
• 开发者包括Youngwan Lee、Kwanyong Park、Yoorhim Cho、Young-Ju Lee、Sung Ju Hwang。
• 模型描述：基于潜在扩散的文本到图像生成模型。KOALA模型使用与SDXL-Base-1.0相同的文本编码器，仅用压缩后的U-Net替换去噪U-Net。
• 教师模型：SDXL-Lightning。
• 训练数据集：LAION-POP数据集的一个子集。
• 训练迭代次数：500K，批量大小为128。
• 使用的GPU：4 x NVIDIA A100 (80GB)。
• 更多信息资源：查看arXiv上的KOALA报告和项目页面。

模型架构

有两种压缩U-Net类型，KOALA-1B和KOALA-700M，通过减少残差块和Transformer块来实现。

U-Net对比

• “Tx”表示Transformer块，“CKPT”表示训练后的检查点文件。
• 我们在NVIDIA 4090 GPU (24GB) 上以FP16精度和25个去噪步骤测量延迟时间。
• SDM-v2.0使用768x768分辨率，而SDXL和KOALA模型使用1024x1024分辨率。

不同GPU上的延迟和内存使用对比

我们使用多种消费级GPU（NVIDIA 3060Ti (8GB)、2080Ti (11GB) 和4090 (24GB)）测量了SDXL-Turbo和KOALA-Turbo模型在512x512分辨率下的推理时间，以及其他模型在1024x1024分辨率下的推理时间。“OOM”表示内存不足。请注意，SDXL模型无法在8GB VRAM的3060Ti上运行，而我们的KOALA模型可以在所有类型的GPU上运行。

模型用途

直接使用

该模型仅用于研究目的。可能的研究领域和任务包括：

• 艺术品生成，以及在设计和其他艺术过程中的应用。
• 在教育或创意工具中的应用。
• 生成模型的研究。
• 对有可能生成有害内容的模型进行安全部署。
• 探索和理解生成模型的局限性和偏差。

超出范围的使用

该模型并非用于生成真实或准确反映人物或事件的内容，因此使用该模型生成此类内容超出了其能力范围。

局限性和偏差

• 文本渲染：模型在图像中渲染长而清晰的文本时面临挑战。
• 复杂提示：KOALA有时在处理涉及多个属性的复杂提示时会遇到困难。
• 数据集依赖：当前的局限性部分归因于训练数据集（LAION-aesthetics-V2 6+）的特性。

🔧 技术细节

• 采用基于自注意力的知识蒸馏方法进行训练。
• 教师模型为SDXL-Lightning。
• 训练数据集是LAION-POP数据集的一个子集。
• 训练迭代次数为500K，批量大小为128。
• 使用4 x NVIDIA A100 (80GB) GPU进行训练。

📄 许可证

文档未提供许可证信息，暂不展示。

📚 引用

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