使用Dreambooth LoRA微调SDXL 0.9

本文将介绍如何通过LoRA对Stable Diffusion XL 0.9进行Dreambooth微调。DreamBooth是一种仅使用几张图像(大约3-5张)来个性化文本到图像模型的方法。

本教程基于通过LoRA进行Unet微调，而不是进行全部的训练。LoRA是在LoRA: Low-Rank Adaptation of Large Language Models中引入的一种参数高效的微调技术。

本文基于diffusers包，至少需要0.18.2或更高版本。

基于GeForce RTX 4090 GPU (24GB)的本地实验，VRAM消耗如下:

512分辨率- 11GB用于训练，19GB保存检查点
1024分辨率- 17GB的训练，19GB时保存检查点

环境设置

建议创建一个新的虚拟环境，下面是我们需要使用的python包

Pytorch

 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

Diffusers

 pip install git+https://github.com/huggingface/diffusers

LoRA进行SDXL 0.9 Dreambooth微调需要0.19.0.dev0及以上版本的diffusers

还有一些其他的依赖包：

 pip install invisible_watermark transformers accelerate safetensors

然后就是进行配置，在终端上执行如下命令配置accelerate:

 accelerate config

使用以下设置在单个GPU上进行混合精度的训练:

 ----------------------------------------------------------------------------------------------------------------------------
 In which compute environment are you running?
 This machine
 ----------------------------------------------------------------------------------------------------------------------------
 Which type of machine are you using?
 No distributed training
 
 Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:
 no
 
 Do you wish to optimize your script with torch dynamo?[yes/NO]:
 no
 
 Do you want to use DeepSpeed? [yes/NO]: 
 no
 
 What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:
 all
 ----------------------------------------------------------------------------------------------------------------------------
 Do you wish to use FP16 or BF16 (mixed precision)?
 fp16

或者，使用以下命令使用默认值

 accelerate config default

数据集

我们这里将介绍Dreambooth微调所需的最简单配置。对于数据集的准备，只需收集一些相同主题或风格的图像，并将其放在一个目录中。

比如下面的文件夹结构:

 data/xxx.png
 data/xxy.png
 ...
 data/xxz.png
 data/yyz.png

这里要确保所有的训练图像都是相同的大小。如果大小不同，需要先调整大小。建议使用1024 * 1024作为图像分辨率。

我们这里使用dog示例数据集通过LoRA测试Dreambooth微调。这个数据集可以直接从网站下载，以下Python脚本可以将其下载到本地：

 from huggingface_hub import snapshot_download
 
 local_dir = "./data"
 snapshot_download(
     "diffusers/dog-example",
     local_dir=local_dir, repo_type="dataset",
     ignore_patterns=".gitattributes",
 )

微调训练

在官方库下载train_dreambooth_lora_sdxl.py训练脚本。将该文件放在工作目录中。

如果你使用的是旧版本的diffusers，它将由于版本不匹配而报告错误。但是你可以通过在脚本中找到check_min_version函数并注释它来轻松解决这个问题，如下所示:

 # check_min_version("0.19.0.dev0")

虽然可以用，但是还是建议使用官方的推荐版本。

如果全部设置正确，那么可以通过LoRA进行Dreambooth微调的训练命令:

 accelerate launch train_dreambooth_lora_sdxl.py \
   --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-0.9" \
   --instance_data_dir=data \
   --output_dir=output \
   --mixed_precision="fp16" \
   --instance_prompt="a photo of zwc dog" \
   --resolution=1024 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=4 \
   --learning_rate=1e-4 \
   --lr_scheduler="constant" \
   --lr_warmup_steps=0 \
   --checkpointing_steps=500 \
   --max_train_steps=1000 \
   --seed="0" \
   --checkpoints_total_limit=5

对于Windows用户，将所有\符号替换为^符号。因为转义符不同

简单的介绍一些参数：

instance_prompt：带有指定实例标识符的提示符。
resolution：输入图像的分辨率，训练/验证数据集中的所有图像都将调整为此大小。默认值是512，将其设置为1024，因为它是用于SDXL训练的分辨率。
train_batch_size：训练数据加载器的批处理大小(每个设备)。减少批处理大小，防止训练过程中出现内存不足错误。
num_train_steps：训练步数。建议设置为N × 100，其中N表示训练图像的个数。
checkpointing_steps：每X次更新时保存训练状态的检查点。默认为500。将其设置为更高的值以减少保存的检查点数量，因为模型需要保存到磁盘，所以频繁的保存会降低训练速度。
checkpoints_total_limit：限制保存的检查点的数量。将删除/删除旧的检查点。

在第一次运行是，程序会下载Stable Diffusion模型并将其保存在本地缓存文件夹中，如果网不好的话这里会很慢。在随后的运行中，它将重用相同的缓存数据。

请注意SDXL 0.9权重需要登录HuggingFace并接受许可。然后，通过HuggingFace -cli命令登录，并使用从HuggingFace设置中获取的API令牌。

默认情况下，每个checkpointing_steps脚本只保存一次LoRA权重和一些检查点文件。

最后我们的结果如下：

 |- output
 |  |- checkpoint-500
 |  |- checkpoint-1000
 |  |- checkpoint-1500
 |  |- checkpoint-2000
 |- data
 |- train_dreambooth_lora_sdxl.py

上面的每个checkpoint文件夹包含以下文件:

optimizer.bin
pytorch_lora_weights.bin
random_states_0.pkl
scaler.pt
scheduler.bin

pytorch_lora_weights.bin文件可以直接用于推理。

推理

创建一个名为inference.py的新Python文件:

 from diffusers import DiffusionPipeline
 import torch
 #初始化，加载所需的LoRA权重
 pipe = DiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-0.9",
     torch_dtype=torch.float16,
     variant="fp16",
     use_safetensors=True
 )
 # load LoRA weight
 pipe.unet.load_attn_procs("data/checkpoint-2000/pytorch_lora_weights.bin", use_safetensors=False)
 pipe.enable_model_cpu_offload()
 
 refiner = DiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-refiner-0.9",
     torch_dtype=torch.float16,
     variant="fp16",
     use_safetensors=True
 )
 refiner.enable_model_cpu_offload()
 
 #推理和保存文件
 seed = 12345
 n_steps = 50
 prompt = "a photo of zwc dog in a bucket"
 
 generator = torch.Generator(device="cuda").manual_seed(seed)
 latent_image = pipe(prompt=prompt, num_inference_steps=n_steps, generator=generator, output_type="latent").images[0]
 image = refiner(prompt=prompt, num_inference_steps=n_steps, generator=generator, image=latent_image).images[0]
 image.save("image.jpg")

然后我们可以执行如下命令:

 python inference.py

结果展示

以下是我做的一个快速测试，使用16张具有各种情绪的chibi 人物图像作为训练数据集。

分辨率1024 × 1024 、duoduo 作为实例提示

大约花了4个小时的训练，下面的输出示例:

总结

使用我们上面的代码可以随意使用不同的数据集和训练配置进行实验，以获得所需的结果。

本文首先简要介绍了Dreambooth和LoRA背后的基本概念。然后介绍了通过pip install进行安装的过程。还探讨了数据集的准备。然后整理了训练命令，并对一些常用的训练参数进行了详细的说明。并使用代码加载新训练的LoRA权重，根据输入提示生成相应的图像。最后展示了一个在本地进行的训练的简单实验。

本文使用的主要库：

https://github.com/huggingface/diffusers

https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_sdxl.md

https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl

标签：深度学习 LoRA Stable Diffusion