块移除知识蒸馏稳定扩散

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

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

• 从SD的U-Net中移除了某些残差和注意力模块。
  • 适用于所有SD-v1和v2版本 — 例如，v1.4、v1.5、v2.1-base、v2.1
• 使用非常有限的数据进行蒸馏预训练，但(令人惊讶地)仍然有效。

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

公告

• [2024年7月] 🎉BK-SDM被**ECCV 2024接收**！
  • 初步版本在ICCV 2023 Demo Track和ICML 2023 ES-FoMo研讨会上展示。
• [2024年4月] EdgeFusion被CVPR 2024 EDGE研讨会接收！
  • 在使用LCM增强BK-SDM-Tiny的同时，突出了改进数据和蒸馏的益处。
• [2023年12月] KOALA引入BK-SDXL基线，非常感谢！
• [2023年8月] 发布BK-SDM的Core ML权重（iPhone 14上4秒推理）。
• [2023年8月] 发布高效个性化T2I的微调代码。
• [2023年8月] 🎉发布高效通用T2I的预训练代码。
  • MODEL_CARD.md包括蒸馏预训练过程和使用不同数据量的结果。
• [2023年8月] Segmind推出他们的BK-SDM实现，非常感谢！
• [2023年8月] Hugging Face本周热门空间🔥介绍了我们的演示，非常感谢！

模型描述

• 以下BK-SDM可在🤗Hugging Face Models获取。参见MODEL_CARD.md中的压缩方法。
• 从SD-v1.4压缩：
  • BK-SDM-{Base, Small, Tiny}：使用0.22M LAION对训练，50K训练迭代。
  • BK-SDM-{Base-2M, Small-2M, Tiny-2M}：2.3M LAION对，50K训练迭代。
• 从SD-v2.1-base压缩：
  • BK-SDM-{v2-Base, v2-Small, v2-Tiny}：使用0.22M LAION对训练，50K训练迭代。

安装

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)'）

  <details> <summary> 关于 'real_im256.npz' 的说明 </summary>
  • 遵循评估协议[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统计数据。
  </details>

• (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 提供此功能 :)

  <details> <summary> 点击查看推理成本详情。 </summary>

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

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

    评估批次大小	1	2	4	8
    GPU内存	4.9GB	6.3GB	11.3GB	19.6GB
    生成时间	9.4小时	7.9小时	7.6小时	7.3小时
    FID	16.98	17.01	17.16	16.97
    IS	31.68	31.20	31.62	31.22
    CLIP 评分	0.2677	0.2679	0.2677	0.2675

  </details>

• (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

• 训练好的 U-Net 用于上述基准评估中的步骤 (2) 生成。
• 要测试特定检查点，请参考示例目录结构修改 --unet_path。

零样本 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）！ <details> <summary> [2023年8月1日] 在 Hugging Face 本周精选空间 🔥 中被推荐 </summary> <img alt="本周精选空间" img src="https://yellow-cdn.veclightyear.com/0a4dffa0/81cd9b9e-55ec-40d8-b083-886daaf7e436.png" width="100%"> </details>

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 的副本。

致谢

• 感谢 Microsoft for Startups Founders Hub 和 光州 AICA 慷慨提供 GPU 资源。
• 感谢 CompVis、Runway 和 Stability AI 在 Stable Diffusion 方面的开创性研究。
• 感谢 LAION、Diffusers、PEFT、DreamBooth、Gradio 和 Core ML Stable Diffusion 的宝贵贡献。
• 感谢会议评审员提供的建设性和富有成果的意见。

引用

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