定制属于你自己的 Stable Diffusion

我们很幸运生活在AI时代。人工智能迅速发展，现在变得更加易于获取。大多数开源模型都经过了充分的训练，拥有来自广泛类别的合理数量的数据，因此我们可以生成许多不同类别的图像。如果模型从未学习过我们想要生成的特定主题怎么办？例如，我想要生成我今天早上刚做的蛋糕的图像。尽管模型可以生成无数种蛋糕图像，但模型不知道如何生成我刚刚制作的蛋糕的图像。

尽管从头开始训练对个人来说几乎是不可能的，但由于不断发展的模型微调技术，将新主题纳入模型可以轻松实现。自2022年以来，Stable Diffusion作为最新趋势的一代模型，为图像生成开启了一个全新的时代。在本文中，我想分享如何使用现代技术LoRA逐步微调Stable Diffusion模型的指南。整个练习在Kaggle笔记本上进行，这样每个人都可以在没有适当设备的情况下完成这个练习。

在我们开始实现之前，让我们快速了解一下Stable Diffusion的概念和微调技术LoRA。如果你已经熟悉背后的知识，请随时跳到实施部分。

Stable Diffusion

图1. Stable Diffusion中的反向扩散

Stable Diffusion是一种可以从文本和图像生成逼真照片的生成性AI模型。扩散模型在潜在空间而不是图像空间中施展魔法，这使得图像生成对公众来说更加可行，因为所需的计算能力更低。人们甚至可以在CPU机器上的小憩期间生成一个不存在的图像。

有许多在线的稳定扩散平台供公众和开发者使用，如HuggingFace¹和Stable Diffusion Online²。我们可以在没有机器学习知识的情况下生成图像。

LoRA（低秩适应）

图2. Stable Diffusion网络

LoRA是一种微调方法，它在稳定扩散网络的交叉注意力层中添加了额外的权重。交叉注意力层在图2中由黄色块表示。交叉注意力层融合了图像ZT和文本τθ的中间信息。向交叉注意力层添加权重扩展了文本提示和图像之间的相关性。这就是LoRA如何向稳定扩散模型引入新知识。

修改后的权重被分解为更小的（低秩）矩阵。这些较小的矩阵携带的参数更少，并单独存储。这就是为什么LoRA权重文件的尺寸可控，但必须与底层的Stable Diffusion模型一起使用。

LoRA的权重大小通常在2到200 MB之间。这是共享和管理模型的优势之一。与其以几GB的整个模型共享，更有效的方法是以更小的尺寸共享LoRA权重。

环境

这个练习是在Kaggle笔记本中进行的，环境设置如下：

• NVIDIA Tesla P100 (RAM=16G)

• diffusers-0.30.0.dev0

• Pytorch 2.1.2

• Python 3.10

Kaggle笔记本已经安装了Python和PyTorch。我们需要安装的唯一库是来自HuggingFace的Diffusers。我们可以使用以下命令从源安装Diffusers。从源安装确保了库的最新版本。路径/kaggle/working是Kaggle运行时的默认工作空间。人们可以调整路径以适应自己的开发环境。

%cd /kaggle/working/!pip install accelerate!git clone https://github.com/huggingface/diffusers%cd /kaggle/working/diffusers!pip install /kaggle/working/diffusers/. -q%cd /kaggle/working/diffusers/examples/text_to_image!pip install -r /kaggle/working/diffusers/examples/text_to_image/requirements.txt -q!pip install --upgrade peft transformers xformers bitsandbytes -q1.2.3.4.5.6.7.

实验在这个练习中，我想向模型介绍蜡笔小新雕像。让我们在以下各节中经历步骤。

第1步：数据收集

图3. 训练样本

首先，我们需要收集一些主题的图像。在这个练习中，我们收集了25张蜡笔小新雕像的照片作为我们的训练数据。一些样本可以在图3中找到。以下是我遵循的数据收集推荐提示：

• 统一图像大小：例如，512x512。

• 高图像质量：图片以12M分辨率拍摄。

• 保持主题在图像中心。

此外，训练集中尽可能多地覆盖多样性也很重要，以便模型学习一般特征，如角色的眼睛、发型和身材。换句话说，我们希望模型学习什么是蜡笔小新雕像，而不是特定着装风格（例如红衬衫和黄色短裤）和姿势（例如站立）。

我已经在5、15、25的培训规模上进行了实验。尽管每一轮之间只有10张图像的差异，但模型的泛化程度有了显著提高。

第2步：基线评估

from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")prompt = "A crayon shin-chan figurine."image = pipeline(prompt, num_inference_steps=100).images[0]image.save("baseline.png")1.2.3.4.5.

在开始训练之前，我们想检查一下预训练模型是否已经知道我们想要微调的特定雕像，使用上述Python脚本。

• AutoPipelineForText2Image类用于加载Stable Diffusion预训练模型stable-diffusion-v1–5。

• 接下来，我们将文本提示设置为“A crayon shin-chan figurine”。我保持文本提示简洁，以便文本信息对模型来说直接了当。

• 将提示和推理次数传递给AutoPipelineForText2Image的实例并进行生成过程。

图4. 基线评估结果

脚本多次运行，图4中展示了三个生成的图像。结果显示，预训练模型知道可能是一个人角色的雕像，但不知道前缀‘crayon shin-chan’具体指的是哪个角色。生成图像中的限制验证了我们微调工作的努力，因为我们的倡议是用一个未见过的主体来微调模型。

第3步：微调

图5. 蜡笔

定义一个新的标识符

除了收集一些训练图像外，我们还需要为这个主题选择一个新的名称。每当新名称出现在文本提示中时，修改后的模型就会理解它指的是新的主题，并相应地生成内容。

我们为什么需要一个新的名称？原始名称Crayon Shin-Chan不是更具代表性吗？我们这样做的想法是，我们需要一个独特的名称，不要与已经训练的其他现有名称混在一起。在我们的例子中，内部的蜡笔暗示了与我们用于绘画的蜡笔的潜在联系。如果关联的文本提示包含Crayon Shin-Chan，生成的图像可能包括与蜡笔相关的内容。

建议使用随机名称作为新标识符。我为这个练习选择了aawxbc这个名字，它几乎不会与任何现有名称重叠。其他人也可以为新主题使用任何其他独特的名称。

训练

图6. 模型生成的类别图像

在Kaggle笔记本中，除了工作区：/kaggle/working/，另一个常用的目录是/kaggle/input/，我们在这里存储工作的输入，如检查点和训练图像。以下是基于脚本train_dreambooth_lora.py的完整训练命令。

!accelerate config default!mkdir /kaggle/working/figurine!accelerate launch --mixed_precision="fp16"  /kaggle/working/diffusers/examples/dreambooth/train_dreambooth_lora.py 
--mixed_precision="fp16" --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" --instance_data_dir="/kaggle/input/aawxbc" --instance_prompt="a photo of aawxbc figurine" --class_data_dir="/kaggle/working/figurine" --class_prompt="a photo of figurine" --with_prior_preservation 
--prior_loss_weight=1.0 --use_8bit_adam 
--lr_scheduler="constant" --gradient_checkpointing 
--enable_xformers_memory_efficient_attention 
--num_train_epochs=50 --output_dir="/kaggle/working/dreambooth-lora"1.2.3.4.5.6.7.8.9.10.11.12.13.14.15.16.17.18.

让我们经历一些关键配置：

• mixed_precision [可选]：选择平台支持的精度。选项有fp32、fp16和bf16。

• class_prompt [可选]：存储类别数据的提示和文件夹。在这个实验中，类别名称是figurine，因为我们希望扩散模型学习一种新的雕像类型。

• class_data_dir [可选]：存储模型生成的类别图像的目录。在我们的例子中，模型生成了一般的雕像图像并将其保存在这个目录下。

• instance_prompt [必须]：我们为新主题定义了一个文本提示。格式是{新名称} {类别名称}，例如aawxbc figurine、uvhhhl cat等。

• instance_data_dir [必须]：存储训练样本的目录。在这个练习中，我们将前面提到的25个样本上传到这个目录。

• num_train_epochs [可选]：训练的轮数。

通过遵循上述步骤，包括数据收集、安装库和完成训练，我们现在可以在由output_dir定义的目录中找到LoRA权重pytorch_lora_weights.safetensors。权重文件的大小为3.2 MB，相当容易管理。

第4步：推理—生成新图像

from diffusers import AutoPipelineForText2Image
pipeline = AutoPipelineForText2Image.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")pipeline.load_lora_weights("/kaggle/working/dreambooth-lora", weight_name="pytorch_lora_weights.safetensors")prompt = "A photo of aawxbc figurine."image = pipeline(prompt, num_inference_steps=180).images[0]image.save("aawxbc.png")1.2.3.4.5.6.

在成功导出LoRA权重后，我们可以进入最后也是最令人兴奋的部分：生成新图像。请注意，LoRA权重不能独立工作，因此我们需要同时加载原始的稳定扩散和LoRA权重文件。

在Python脚本中，除了我们在基线评估中使用的步骤外，我们还有一个额外的步骤，使用.load_lora_weights()函数加载LoRA模型，传入目录和LoRA权重文件的名称。

图7. 使用提示生成的图像：“沙滩上的aawxbc雕像，戴着帽子”

对于文本提示，我们在脚本中包含了标识符aawxbc和类别名称figurine。两组相应文本提示生成的图像分别在图7和图8中展示。我们可以看到，经过微调的模型可以生成蜡笔小新雕像。

图8. 使用提示生成的图像：“森林里的aawxbc雕像，拿着杯子。”

我们可以通过注释掉加载LoRA权重的行并保持文本提示不变来进一步交叉验证新模型。我们可以看到原始的稳定扩散模型无法正确解释标识符aawxbc，这与基线评估一致，并突出了LoRA模型的贡献。

总结

图9. 微调概述

工作的概述在图9中展示。新模型结合了LoRA权重，并且可以生成带有新实例名称aawxbc的蜡笔小新雕像的图像，如绿色矩形突出所示。我们所要做的就是收集足够的图像，并用不存在的标识符命名新主题。就这些。所有的艰苦工作都在代码中完成了。

参考文献：

• [1] https://huggingface.co/spaces/stabilityai/stable-diffusion

• [2] https://stablediffusionweb.com/