Logo
Project Icon

BK-SDM

高效轻量的Stable Diffusion压缩模型

Stable DiffusionAI绘图模型压缩知识蒸馏图像生成Github开源项目

BK-SDM是一种压缩版Stable Diffusion模型,通过移除U-Net中的部分模块实现轻量化。该模型采用有限数据进行蒸馏预训练,适用于SD v1和v2各版本,提供基础、小型和微型三种规模。BK-SDM在保持图像质量的同时,显著提高了推理速度,降低了计算资源需求,为高效文本到图像生成提供了新选择。

GithubGithubHuggingfaceHuggingface论文论文
介绍相关项目

块移除知识蒸馏稳定扩散

BK-SDM: 轻量、快速且经济的稳定扩散版本的官方代码库 [ArXiv] [ECCV 2024]。

BK-SDM是轻量级文本到图像(T2I)合成模型：

  • 从SD的U-Net中移除了某些残差和注意力模块。
  • 使用非常有限的数据进行蒸馏预训练，但(令人惊讶地)仍然有效。

⚡快速链接：KD预训练 | MS-COCO评估 | DreamBooth微调 | 演示

公告

模型描述

安装

conda create -n bk-sdm python=3.8
conda activate bk-sdm
git clone https://github.com/Nota-NetsPresso/BK-SDM.git
cd BK-SDM
pip install -r requirements.txt

关于我们使用的torch版本说明：

  • 在单个24GB RTX3090上进行MS-COCO评估和DreamBooth微调时使用torch 1.13.1
  • 在单个80GB A100上进行KD预训练时使用torch 2.0.1
    • 如果在A100上使用总批量大小256进行预训练导致GPU内存不足，请检查torch版本并考虑升级到torch>2.0.0。

使用🤗Diffusers的最小示例

使用默认PNDM调度器和50个去噪步骤：

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("nota-ai/bk-sdm-small", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "一个装有各种花朵的金色花瓶"
image = pipe(prompt).images[0]  
    
image.save("example.png")

import torch
from diffusers import StableDiffusionPipeline, UNet2DConditionModel

pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe.unet = UNet2DConditionModel.from_pretrained("nota-ai/bk-sdm-small", subfolder="unet", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "一个装有各种花朵的金色花瓶"
image = pipe(prompt).images[0]  
    
image.save("example.png")

蒸馏预训练

我们的代码基于Diffusers 0.15.0版本的train_text_to_image.py。要访问最新版本,请使用此链接。

[可选] 用于检查可运行性的小型测试

bash scripts/get_laion_data.sh preprocessed_11k
bash scripts/kd_train_toy.sh

<详细信息> <摘要>注意</摘要>

  • 一个小型数据集(11K图像-文本对)将被下载到./data/laion_aes/preprocessed_11k(tar.gz格式1.7GB; 解压后数据文件夹1.8GB)。
  • 可以使用小型脚本来验证代码的可执行性,并找到适合你的GPU的批量大小。使用批量大小8(=4×2),训练BK-SDM-Base 20次迭代大约需要5分钟,使用22GB的GPU内存。

</详细信息>

单GPU训练BK-SDM-{Base, Small, Tiny}

bash scripts/get_laion_data.sh preprocessed_212k
bash scripts/kd_train.sh

<详细信息> <摘要>注意</摘要>

  • 包含212K(=0.22M)对的数据集将被下载到./data/laion_aes/preprocessed_212k(tar.gz格式18GB; 解压后数据文件夹20GB)。
  • 使用批量大小256(=4×64),训练BK-SDM-Base 50K次迭代大约需要300小时,使用53GB的GPU内存。使用批量大小64(=4×16),需要60小时,使用28GB的GPU内存。
  • 训练BK-SDM-{Small, Tiny}会导致GPU内存使用减少5~10%。

</详细信息>

单GPU训练BK-SDM-{Base-2M, Small-2M, Tiny-2M}

bash scripts/get_laion_data.sh preprocessed_2256k
bash scripts/kd_train_2m.sh

<详细信息> <摘要>注意</摘要>

  • 包含2256K(=2.3M)对的数据集将被下载到./data/laion_aes/preprocessed_2256k(tar.gz格式182GB; 解压后数据文件夹204GB)。
  • 除了数据集之外,kd_train_2m.sh与kd_train.sh相同;在相同的迭代次数下,训练计算保持不变。

</详细信息>

多GPU训练

bash scripts/kd_train_toy_ddp.sh

<详细信息> <摘要>注意</摘要>

  • 支持多GPU训练(样例结果:链接),尽管我们论文中的所有实验都是使用单个GPU进行的。感谢@youngwanLEE分享脚本 :)

</详细信息>

用BK-SDM压缩SD-v2

bash scripts/kd_train_v2-base-im512.sh
bash scripts/kd_train_v2-im768.sh

# 对于推理,请参见:'scripts/generate_with_trained_unet.sh'

关于训练代码的说明

<详细信息> <摘要> KD训练的关键部分 </摘要>

  • 通过调整config.json定义学生U-Net [链接]
  • 通过复制教师U-Net的权重初始化学生U-Net [链接]
  • 为特征KD定义钩子位置 [链接]
  • 定义特征和输出KD的损失 [链接]

</详细信息>

<详细信息> <摘要> 关键学习超参数 </摘要>

--unet_config_name "bk_small" # 选项: ["bk_base", "bk_small", "bk_tiny"]
--use_copy_weight_from_teacher # 使用教师权重初始化学生unet
--learning_rate 5e-05
--train_batch_size 64
--gradient_accumulation_steps 4
--lambda_sd 1.0
--lambda_kd_output 1.0
--lambda_kd_feat 1.0

</详细信息>

在MS-COCO基准测试上的评估

我们使用以下代码获得MS-COCO上的结果。使用PNDM调度器和25步去噪生成512×512图像后,我们将它们下采样到256×256以计算评分。

使用发布的模型生成(默认使用BK-SDM-Small)

在单个3090 GPU上,'(2)'每个模型需要约10小时,'(3)'需要几分钟。

  • (1) 下载 metadata.csv 和 real_im256.npz：

    bash scripts/get_mscoco_files.sh

# ./data/mscoco_val2014_30k/metadata.csv: MS-COCO验证集中的30K提示（用于'(2)'）
# ./data/mscoco_val2014_41k_full/real_im256.npz: 41K真实图像的FID统计数据（用于'(3)'）
    关于 'real_im256.npz' 的说明
    • 遵循评估协议[DALL·E, Imagen]，真实图像的FID统计数据是在MS-COCO完整验证集（41K图像）上计算的。通过'(1)'下载的预计算统计文件位于 ./data/mscoco_val2014_41k_full/real_im256.npz。
    • 此外，可以使用 python3 src/get_stat_mscoco_val2014.py 计算 real_im256.npz，该脚本会下载所有图像，将它们调整为256×256大小，并计算FID统计数据。

  • (2) 基于MS-COCO验证集中的30K提示生成512×512图像 → 将它们调整为256×256：

    python3 src/generate.py 

# python3 src/generate.py --model_id nota-ai/bk-sdm-base --save_dir ./results/bk-sdm-base
# python3 src/generate.py --model_id nota-ai/bk-sdm-tiny --save_dir ./results/bk-sdm-tiny

    [批量生成] 增加 --batch_sz（默认：1）可以加快推理速度，但会增加显存使用量。感谢 @Godofnothing 提供此功能 :)

    点击查看推理成本详情。

    • 设置：BK-SDM-Small 在 MS-COCO 30K 图像生成任务上

    • 我们在论文结果中使用了评估批次大小为1。不同的批次大小会影响随机潜在编码的采样，导致略微不同的生成分数。

      评估批次大小1248
      GPU内存4.9GB6.3GB11.3GB19.6GB
      生成时间9.4小时7.9小时7.6小时7.3小时
      FID16.9817.0117.1616.97
      IS31.6831.2031.6231.22
      CLIP 评分0.26770.26790.26770.2675

  • (3) 计算 FID、IS 和 CLIP 评分：

    bash scripts/eval_scores.sh

# 对于其他模型，修改脚本中的 `./results/bk-sdm-*` 路径以指定不同的模型。

[训练后] 使用训练好的 U-Net 进行生成

bash scripts/get_mscoco_files.sh
bash scripts/generate_with_trained_unet.sh

零样本 MS-COCO 256×256 30K 结果

请参阅 MODEL_CARD.md 中的结果

使用 🤗PEFT 进行 DreamBooth 微调

我们的轻量级 SD 骨干网络可用于高效的个性化生成。DreamBooth 能够根据少量图像改进文本到图像的扩散模型。DreamBooth+LoRA 可以大幅降低微调成本。

DreamBooth 数据集

数据集下载到 ./data/dreambooth/dataset [文件夹树]：30个主题 × 25个提示 × 4~6张图像。

git clone https://github.com/google/dreambooth ./data/dreambooth

DreamBooth 微调（默认使用 BK-SDM-Base）

我们的代码基于 PEFT 0.1.0 的 train_dreambooth.py。要访问最新版本，请使用此链接。

  • (1) 不使用 LoRA — 全面微调 & 在我们的论文中使用 
    bash scripts/finetune_full.sh # 学习率 1e-6
bash scripts/generate_after_full_ft.sh
  • (2) 使用 LoRA — 参数高效微调 
    bash scripts/finetune_lora.sh # 学习率 1e-4
bash scripts/generate_after_lora_ft.sh
  • 在单个 3090 GPU 上，每个主题的微调需要 10~20 分钟。

个性化生成结果

请参阅 MODEL_CARD.md 中的 DreamBooth 结果

Gradio 演示

查看我们的 Gradio 演示 和 代码（主要文件：app.py）！

[2023年8月1日] 在 Hugging Face 本周精选空间 🔥 中被推荐 本周精选空间

Core ML 权重

对于 iOS 或 macOS 应用程序，我们已将模型转换为 Core ML 格式。它们可在 🤗Hugging Face Models（nota-ai/coreml-bk-sdm）上获取，并可与 Apple 的 Core ML Stable Diffusion 库 一起使用。

  • iPhone 14 上 4 秒推理（10 步去噪）：结果

许可证

本项目及其权重受 CreativeML Open RAIL-M 许可证 约束，旨在减轻使用高度先进的机器学习系统可能产生的任何潜在负面影响。该许可证的摘要 如下：

1. 您不能使用该模型故意生产或分享非法或有害的输出或内容，
2. 我们不对您生成的输出主张任何权利，您可以自由使用它们，但需对其使用负责，使用时不应违反许可证中规定的条款，
3. 您可以重新分发权重并商业使用该模型和/或将其作为服务提供。如果您这样做，请注意您必须包含与许可证中相同的使用限制，并向所有用户分享 CreativeML OpenRAIL-M 的副本。

致谢

引用

@article{kim2023bksdm,
  title={BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion},
  author={Kim, Bo-Kyeong and Song, Hyoung-Kyu and Castells, Thibault and Choi, Shinkook},
  journal={arXiv preprint arXiv:2305.15798},
  year={2023},
  url={https://arxiv.org/abs/2305.15798}
}
相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

白日梦AI

白日梦AI提供专注于AI视频生成的多样化功能，包括文生视频、动态画面和形象生成等，帮助用户快速上手，创造专业级内容。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

讯飞绘镜

讯飞绘镜是一个支持从创意到完整视频创作的智能平台，用户可以快速生成视频素材并创作独特的音乐视频和故事。平台提供多样化的主题和精选作品，帮助用户探索创意灵感。

Project Cover

讯飞文书

讯飞文书依托讯飞星火大模型，为文书写作者提供从素材筹备到稿件撰写及审稿的全程支持。通过录音智记和以稿写稿等功能，满足事务性工作的高频需求，帮助撰稿人节省精力，提高效率，优化工作与生活。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具，简化了选题、文献检索至论文撰写的整个过程。通过简单设定，平台可快速生成高质量论文大纲和全文，配合图表、参考文献等一应俱全，同时提供开题报告和答辩PPT等增值服务，保障数据安全，有效提升写作效率和论文质量。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号