定制属于你自己的 Stable Diffusion

点击下方卡片,关注“小白玩转Python”公众号

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

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

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

Stable Diffusion

1a3e9267eca133eca8092e113e11388b.png图1. Stable Diffusion中的反向扩散

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

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

LoRA(低秩适应)

5cdf0b341bac9acdcfcfd6327753e223.png

图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 -q

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

第1步:数据收集

4adaaebda30094f72544c19257289ca6.png图3. 训练样本

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

  1. 统一图像大小:例如,512x512。

  2. 高图像质量:图片以12M分辨率拍摄。

  3. 保持主题在图像中心。

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

我已经在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")

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

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

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

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

230e371758b3434deb54c972c4de026c.png图4. 基线评估结果

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

第3步:微调

9a547ae4e7d9d0f2ec97547171c0c22f.png图5. 蜡笔

定义一个新的标识符

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

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

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

训练

9c06c72a001e321404452da2d37a33ad.png图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"

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

  • 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")

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

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

a2413f71e583de3efbcdc103b5f58591.png图7. 使用提示生成的图像:“沙滩上的aawxbc雕像,戴着帽子”

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

14f046cadcbeadbf6f0eac883a2a0988.png图8. 使用提示生成的图像:“森林里的aawxbc雕像,拿着杯子。”

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

总结

2b4e1c20a205d2dd94e5c8d77a568177.png图9. 微调概述

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

参考文献 [1] https://huggingface.co/spaces/stabilityai/stable-diffusion

[2] https://stablediffusionweb.com/


·  END  ·

🌟 想要变身计算机视觉小能手?快来「小白玩转Python」公众号!

回复Python视觉实战项目,解锁31个超有趣的视觉项目大礼包!🎁

c67b87838a85855c56fafd406662ffb8.png

本文仅供学习交流使用,如有侵权请联系作者删除

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值