small-stable-diffusion-v0

Small Stable Diffusion V0

最近，我们发布了diffusion-deploy，以加快在GPU(\~4倍速度，基于TensorRT)和CPU(\~12倍速度，基于IntelOpenVINO)上的推理。

与这个 repo 集成后，small-stable-diffusion可以**在CPU上仅用10秒生成图像**。

这里我们给出了一个small-stable-diffusion, 和stablte diffusion相似的图片质量，但是模型小了近一半

一些样例:

代码使用范例

Training

Initialization

这个模型是由stable-diffusion v1-4初始化的。由于模型结构与stable-diffusion不一样，参数数量较少，不能直接利用stable-diffusion的参数。因此，small-stable-diffusion设置 layers per block=1 ，并选择原始stable-diffusion中每个块的第一层来初始化小模型。

Training Procedure

初始化后，模型在8xA100 GPUS中被训练了1100k步。训练过程由三个阶段组成。第一阶段是一个简单的预训练阶段。在后两个阶段，利用原始的stable-diffusion作为教师模型来蒸馏小模型。在所有阶段，只有unet中的参数被训练，其他参数被冻结。

• *硬件:** 8 x A100-80GB GPUs
• *优化器:** AdamW

• *阶段 1** - 预训练模型的unet部分。
• **Steps**: 500,000
• **Batch:** batch size=8, GPUs=8, Gradient Accumulations=2. Total batch size=128
• **Learning rate:** 预热10,000步到1e-5，然后保持稳定
• *阶段 2** - 使用stable-diffusionv1-4作为教师来提炼模型。除了地面真实，这个阶段的训练还使用教师模型产生的软标签（ pred noise ）
• **Steps**: 400,000
• **Batch:** batch size=8, GPUs=8, Gradient Accumulations=2. Total batch size=128
• **Learning rate:** 预热5,000步到1e-5，然后保持稳定
• **Soft label weight:** 0.5
• **Hard label weight:** 0.5
• *阶段 3** - 使用stable-diffusionv1-5作为老师来蒸馏模型。 用到了 Knowledge Distillation of Transformer-based Language Models Revisited 中的一些技术，包括除了软标签之外的基于相似性的层匹配。
• **Steps**: 200,000
• **Batch:** batch size=8, GPUs=8, Gradient Accumulations=2. Total batch size=128
• **Learning rate:** 预热5,000步到1e-5，然后保持稳定
• **Softlabel weight:** 0.5
• **Hard label weight:** 0.5

Training Data

我们使用以下数据集来训练模型。

• LAION-2B en aesthetic
• LAION-Art
• LAION-HD

Citation