稳定扩散的体积缩小 50%,速度更快 🚀
知识蒸馏过的小型版本稳定扩散模型。非官方实现,描述请见 BK-SDM。
这些通过蒸馏训练的模型在生成图像质量方面与完整版稳定扩散模型相似,同时速度更快且体积更小。
本仓库的组成部分:
- data.py 包含下载训练数据的脚本。
- distill_training.py 使用论文中描述的方法训练 U-net。这可能需要根据所需训练的模型类型 (sd_small/sd_tiny)、批大小、超参数等进行额外配置。基础训练代码来自 Huggingface 🤗 diffusers 库。
可以通过使用标准的 diffusers 脚本来进行 LoRA 训练和从检查点进行训练。
训练详情:
知识蒸馏训练神经网络类似于教师一步步指导学生(一个较松散的例子)。大模型在大量数据上进行训练,然后在较小数据集上训练小模型,目标是模仿大模型的输出,同时对数据集进行经典训练。
在知识蒸馏训练中,我们使用了 SG161222/Realistic_Vision_V4.0's U-net 作为教师模型,并使用了 recastai/LAION-art-EN-improved-captions 的子集作为训练数据。
最终的训练损失是教师 U-net 预测的噪声与学生 U-net 预测的噪声之间的均方误差 (MSE) 损失、实际添加的噪声与预测的噪声之间的 MSE 损失,以及每个块之后教师和学生 U-net 预测之间的 MSE 损失之和的总和。
总损失:
任务损失(即添加噪声与实际噪声之间的 MSE 损失):
知识蒸馏输出损失(即教师 U-net 与学生 U-net 最终输出之间的 MSE 损失):
特征级别知识蒸馏损失(即 U-net 中每个块的输出之间的 MSE 损失):
以下是我们用于训练的设置:
lr = 1e-5
scheduler = "cosine"
batch_size = 32
output_weight = 0.5 # 最终损失公式中的 Lambda Out
feature_weight = 0.5 # 最终损失公式中的 Lambda Feat
参数:
普通稳定扩散 U-net:
参数数量: 859,520,964
SD_Small U-net:
参数数量: 579,384,964
SD_Tiny U-net:
参数数量: 323,384,964
(使用 torchinfo 报告的模型参数)
使用方法
import torch
from diffusers import DiffusionPipeline
from diffusers import DPMSolverMultistepScheduler
from torch import Generator
path = 'segmind/small-sd' # 对应模型类型的路径
# 在下方插入您的提示。
prompt = "年轻漂亮的 18 岁白人女孩肖像,穿着高领毛衣,(杰作,极其细致的皮肤,逼真,浓重阴影,戏剧性和电影般的灯光,主光,辅助光),清晰对焦,BREAK 史诗现实主义"
# 在下方插入负面提示。我们建议使用此负面提示以获得最佳效果。
negative_prompt = "(变形的虹膜,变形的瞳孔,半逼真,cgi,3d,渲染,草图,卡通,绘画,动漫:1.4),文字,特写,裁剪,框外,最差质量,低质量,jpeg 压缩痕迹,丑陋,重复,病态,残缺,多余手指,畸形手,画得差的手,画得差的脸,突变,畸形,模糊,脱水,解剖结构差,比例不良,额外肢体,克隆脸,畸形,恶心比例,畸形肢体,缺少手臂,缺少腿,额外手臂,额外腿,融合的手指,过多的手指,长脖子"
torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True
# 下面的代码将在 GPU 上运行,请将设备的所有地方传递为 cpu,并将 'dtype' 设置为 torch.float32 以进行 CPU 推理。
with torch.inference_mode():
gen = Generator("cuda")
gen.manual_seed(1674753452)
pipe = DiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False)
pipe.to('cuda')
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.unet.to(device='cuda', dtype=torch.float16, memory_format=torch.channels_last)
img = pipe(prompt=prompt, negative_prompt=negative_prompt, width=512, height=512, num_inference_steps=25, guidance_scale=7, num_images_per_prompt=1, generator=gen).images[0]
img.save("image.png")
训练模型:
知识蒸馏训练的说明类似于 diffusers 文本到图像微调脚本,但增加了一些额外的参数:
--distill_level
: "sd_small" 或 "sd_tiny" 之一,具体取决于要训练的模型类型。
--output_weight
: 表示输出级别 KD 损失的缩放比例的浮点数。
--feature-weight
: 表示特征级别 KD 损失的缩放比例的浮点数。
--prepare_unet
: "True" 或 "False" 之一;表示是否在蒸馏之前将 U-net 准备为较小的尺寸。如果是从全尺寸稳定扩散检查点进行蒸馏,请设置为 "True"。
此外,snr_gamma
已被移除。
我们建议使用标准的稳定扩散模型进行蒸馏训练,因为脚本已针对这些架构进行配置。
示例:
export MODEL_NAME="SG161222/Realistic_Vision_V4.0"
export DATASET_NAME="fantasyfish/laion-art"
accelerate launch --mixed_precision="fp16" distill_training.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name=$DATASET_NAME \
--resolution=512 --center_crop --random_flip \
--train_batch_size
加速启动 --mixed_precision="fp16" checkpoint_training.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name=$DATASET_NAME \
--use_ema \
--resolution=512 --center_crop --random_flip \
--train_batch_size=1 \
--gradient_accumulation_steps=4 \
--gradient_checkpointing \
--max_train_steps=15000 \
--learning_rate=1e-05 \
--max_grad_norm=1 \
--lr_scheduler="constant" --lr_warmup_steps=0 \
--output_dir="sd-laion-art"
进行LoRA训练:
export MODEL_NAME="segmind/small-sd"
export DATASET_NAME="fantasyfish/laion-art"
加速启动 --mixed_precision="fp16" lora_training.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--dataset_name=$DATASET_NAME --caption_column="text" \
--resolution=512 --random_flip \
--train_batch_size=1 \
--num_train_epochs=100 --checkpointing_steps=5000 \
--learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
--seed=42 \
--output_dir="fantasyfish/laion-art" \
--validation_prompt="A man in a suit" --report_to="wandb"
--use_peft \
--lora_r=4 --lora_alpha=32 \
--lora_text_encoder_r=4 --lora_text_encoder_alpha=32
后两个脚本来自🤗 diffusers github。
预训练检查点:
- 训练好的“sd-small”模型版本可在这个Huggingface 🤗 仓库找到
- 训练好的“sd-tiny”模型版本可在这个Huggingface 🤗 仓库找到
- 在人像图片上微调的“sd-tiny模型”版本可在这个Huggingface 🤗 仓库找到
以下是一些生成的示例:
在人像图片上微调的SD-tiny模型
以下是使用在肖像图像上微调的sd-tiny模型生成的一些图片。
一名28岁男子的特写肖像照,穿着粗糙的皮夹克,留着胡茬,眼角周围有明显的笑纹,在柔和的黄金时段灯光下拍摄。
RAW照片,(特写:1.2),一名35岁女子的肖像,化妆淡薄,展现了她的雀斑,在一个葱郁的植物园里,神态宁静,被柔和的斑驳阳光照亮。
一张高质量的40岁欧洲男子的面部肖像照片,戴着眼镜,展示了他额头上的细纹和特征,他的盐胡胡子增添了一丝精致感,微弱的阳光透过浓密的树叶微微透出。
一名48岁女子的黑白照片,从侧面拍摄,突出了她优雅的侧脸和脸颊上的细纹,在一个风大的阴天海滩上捕捉到了她宁静的瞬间。
一张高质量的30岁亚洲女子的特写肖像照片,穿着传统服饰,强调了她瓷器般无瑕的皮肤质感,展现了她传统头饰的细节。
模型链接 -> Huggingface 🤗 仓库
LoRA训练的SD-tiny模型在抽象图像上生成的结果
以下是使用LoRA训练的sd-tiny模型,在抽象概念上生成的一些图片。
模型链接 -> Huggingface 🤗 仓库
在NVIDIA A100 80GB上推理速度的比较:
优点
- 推理速度提高至多100%
- VRAM占用降低至多30%
- 更快的dreambooth和LoRA训练
限制
- 蒸馏模型处于早期阶段,输出结果可能尚未达到生产质量。
- 这些模型可能不是最佳的通用模型。它们最好用于特定概念/风格上的微调或LoRA训练。
- 蒸馏模型在可组合性或多概念处理上仍然表现不佳。
研究路线图
- SDXL蒸馏模型和代码。
- 进一步微调的SD-1.5基础模型,以提高可组合性和通用性。
- 在训练/微调过程中应用Flash Attention-2以加快速度。
- 应用TensorRT和/或AITemplate以进一步加速。
- 在蒸馏过程中考虑量化感知训练(QAT)。
致谢:
- 如果没有Nota AI的论文关于压缩模型的工作,这项工作是不可能完成的。我们对他们在该领域的研究工作表示感谢。
贡献者:
由contrib.rocks制作。
引用
@article{kim2023architectural,
title={On Architectural Compression of Text-to-Image Diffusion Models},
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}
}
@article{Kim_2023_ICMLW,
title={BK-SDM: Architecturally Compressed Stable Diffusion for Efficient Text-to-Image Generation},
author={Kim, Bo-Kyeong and Song, Hyoung-Kyu and Castells, Thibault and Choi, Shinkook},
journal={ICML Workshop on Efficient Systems for Foundation Models (ES-FoMo)},
year={2023},
url={https://openreview.net/forum?id=bOVydU0XKC}
}