Stable Diffusion XL之使用Stable Diffusion XL训练自己的AI绘画模型

一 SDXL训练基本步骤

Stable Diffusion XL系列模型的训练过程主要分成以下几个步骤，

1. 训练集制作：数据质量评估，标签梳理，数据清洗，数据标注，标签清洗，数据增强等。
2. 训练文件配置：预训练模型选择，训练环境配置，训练步数设置，其他超参数设置等。
3. 模型训练：运行SDXL模型/LoRA模型训练脚本，使用TensorBoard监控模型训练等。
4. 模型测试：将训练好的自训练SDXL模型/LoRA模型用于效果评估与消融实验

二 从0到1上手使用Stable Diffusion XL训练自己的AI绘画模型

下面是SDXL的训练资源
链接：https://pan.quark.cn/s/5664c0de758d
提取码：LRZV

2.1 配置训练环境与训练文件

1. 首先进入SDXL-Train项目中，安装SDXL训练所需的依赖库

cd SDXL-Train
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple some-package

# accelerate库的版本需要重新检查一遍，需要安装accelerate==0.16.0版本才能兼容SDXL的训练
pip install accelerate==0.16.0 -i https://pypi.tuna.tsinghua.edu.cn/simple some-package

# 这里推荐大家安装2.0.1版本的Pytorch，能够兼容SDXL训练的全部流程
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
或者conda
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 -c pytorch -c nvidia

在完成上述的依赖库安装后，我们需要确认一下目前的Python、PyTroch、CUDA以及cuDNN的版本是否兼容，在命令行输入以下命令：

# Python版本推荐3.8或者3.9，两个版本皆可
>>>python
Python 3.9
# 加载PyTroch
>>>import torch
# 查看PyTorch版本
>>>print(torch.version)
2.0.1
# 查看CUDA版本
>>> print(torch.version.cuda)
11.8
# 查看cuDNN版本
>>> print(torch.backends.cudnn.version())
8500
# 查看PyTroch、CUDA以及cuDNN的版本是否兼容，True代表兼容
>>> print(torch.cuda.is_available())
True

2. 安装和验证好所有SDXL训练所需的依赖库后，我们还需要设置一下SDXL的训练环境，我们主要是用accelerate库的能力，accelerate库能让PyTorch的训练和推理变得更加高效简洁。我们只需在命令行输入以下命令，并对每个设置逐一进行填写即可：

# 输入以下命令，开始对每个设置进行填写
accelerate config

# 开始进行训练环境参数的配置
In which compute environment are you running? # 选择This machine，即本机
This machine

# 选择单卡或是多卡训练，如果是多卡，则选择multi-GPU，若是单卡，则选择No distributed training                                                                                                               
Which type of machine are you using?                                                                                        
multi-GPU

# 几台机器用于训练，一般选择1台。注意这里是指几台机器，不是几张GPU卡                                                                                                                  
How many different machines will you use (use more than 1 for multi-node training)? [1]: 1

# torch dynamo，DeepSpeed，FullyShardedDataParallel，Megatron-LM等环境参数，不需要配置                          
Do you wish to optimize your script with torch dynamo?[yes/NO]: # 输入回车即可                                                           
Do you want to use DeepSpeed? [yes/NO]:  # 输入回车即可                                                                                  
Do you want to use FullyShardedDataParallel? [yes/NO]:    # 输入回车即可                                                                 
Do you want to use Megatron-LM ? [yes/NO]:       # 输入回车即可

# 选择多少张卡投入训练                                                                          
How many GPU(s) should be used for distributed training? [1]:2

# 设置投入训练的GPU卡id，如果是全部的GPU都投入训练，则输入all即可。
What GPU(s) (by id) should be used for training on this machine as a comma-seperated list? [all]:all 

# 训练精度，可以选择fp16
Do you wish to use FP16 or BF16 (mixed precision)? 
fp16                             

# 完成配置后，配置文件default_config.yaml会保存在/root/.cache/huggingface/accelerate下                                                                                   
accelerate configuration saved at /root/.cache/huggingface/accelerate/default_config.yaml完成上述的流程后，接下来我们就可以进行SDXL训练数据的制作和训练脚本的配置流程了！

后续进行SDXL与SDXL LoRA模型训练的时候，只需要加载对应的default_config.yaml配置文件即可，具体的调用方法，本文后续的章节会进行详细讲解。

还有一点需要注意的是，我们进行SDXL模型的训练时，SDXL的CLIP Text Encoder会调用CLIP-ViT-bigG-14-laion2B-39B-b160k和clip-vit-large-patch14两个配置文件。一般情况下SDXL模型会从huggingface上将配置文件下载到~/.cache/huggingface/目录中，但是由于网络原因很可能会下载失败，从而导致训练的失败。

所以为了让大家能更方便的训练SDXL模型，已经将CLIP-ViT-bigG-14-laion2B-39B-b160k和clip-vit-large-patch14这两个配置文件放入SDXL-Train项目的utils_json文件夹中，并且已经为大家配置好依赖路径，大家只要使用SDXL-Train项目便无需做任何修改。如果大家想要修改CLIP-ViT-bigG-14-laion2B-39B-b160k和clip-vit-large-patch14这两个依赖文件夹的调用路径，大家可以找到SDXL-Train/library/sdxl_train_util.py脚本中的第122行，将"utils_json/"部分修改成自己的本地自定义路径比如“/本地路径/utils_json/”即可。

完成上述的流程后，接下来我们就可以进行SDXL训练数据的制作和训练脚本的配置流程了！

2.2 SDXL训练数据集制作

(1) 数据筛选与清洗

首先，我们需要对数据集进行清洗，和传统深度学习时代一样，数据清洗工作依然占据了AIGC时代模型训练70%-80%左右的时间。

并且这个过程必不可少，因为数据质量决定了机器学习性能的上限，而算法和模型只是在不断逼近这个上限而已。

• 我们需要筛除分辨率较低、质量较差（比如说768*768分辨率的图片< 100kb）、存在破损以及和任务目标无关的数据，接着再去除数据里面可能包含的水印，干扰文字等污染特征。

• 同时，我们需要优先保证数据集的质量，在有质量的基础上再去增加数据集的数量与丰富度。

• 为了满足AI绘画生成图片时的尺寸适应度，我们可以对数据进行多尺度的增强，比如进行1:1，1:2，2:1，1:3，3:4，4:3，9:16，16:9等等尺寸的裁剪与缩放操作。但是切记不能在多尺度增强的时候将图片的主体特征裁剪掉（比如人脸，建筑等）。

完成上述的数据筛选与清洗工作后，我们就可以开始进行数据标注了。

数据标注可以分为自动标注和手动标注。自动标注主要依赖像BLIP（img2caption）和Waifu Diffusion 1.4（img2tag）等能够进行图片生成标签的模型，手动标注则依赖标注人员。

(2) 使用BLIP自动标注caption

我们先用BLIP对数据进行自动标注，BLIP输出的是自然语言标签，我们进入到SDXL-Train/finetune/路径下，运行以下代码即可获得自然语言标签（caption标签）：

cd SDXL-Train/finetune/
python make_captions.py "/数据路径" --caption_weights “../BLIP/model_large_caption.pth” --batch_size=8 --beam_search --min_length=5 --max_length=75 --debug --caption_extension=".caption" --max_data_loader_n_workers=2

注意：在使用BLIP进行数据标注时需要依赖bert-base-uncased模型，已经帮大家配置好了，大家只要使用SDXL-Train项目便无需做任何修改。同时，如果大家想要修改bert-base-uncased模型的调用路径，可以找到SDXL-Train/finetune/blip/blip.py脚本的第189行，将“…/bert-base-uncased”部分修改成自己的本地自定义路径比如“/本地路径/bert-base-uncased”即可。

从上面的代码可以看到，我们第一个传入的参数是训练集的路径。下面一一向大家介绍一下其余参数的意义：

--caption_weights：表示加载的本地BLIP模型，如果不传入本地模型路径，则默认从云端下载BLIP模型。
--batch_size：表示每次传入BLIP模型进行前向处理的数据数量。
--beam_search：设置为波束搜索，默认Nucleus采样。
--min_length：设置caption标签的最短长度。
--max_length：设置caption标签的最长长度。
--debug：如果设置，将会在BLIP前向处理过程中，打印所有的图片路径与caption标签内容，以供检查。
--caption_extension：设置caption标签的扩展名，一般为".caption"。
--max_data_loader_n_workers：设置大于等于2，加速数据处理。

(3) 使用Waifu Diffusion 1.4自动标注tag

接下来我们可以使用Waifu Diffusion 1.4进行自动标注，Waifu Diffusion 1.4输出的是tag关键词，这里需要注意的是，调用Waifu Diffusion 1.4模型需要安装Tensorflow库，并且需要下载特定的版本（2.10.1）或者(2.10.0)，不然运行时会报“DNN library is not found“错误。我们只需要在命令行输入以下命令即可：
如果是2.10.1建议不适用国内源，可能更新不完全，会报错，2.10.0可以使用国内源。

pip install tensorflow==2.10.1

或者

pip install tensorflow==2.10.0 -i https://pypi.tuna.tsinghua.edu.cn/simple some-package

完成上述的环境配置后，我们依然进入到SDXL-Trian/finetune/路径下，运行以下代码即可获得tag自动标注：

cd SDXL-Train/finetune/
python tag_images_by_wd14_tagger.py "/数据路径" --batch_size=8 --model_dir="../tag_models/wd-v1-4-moat-tagger-v2" --remove_underscore --general_threshold=0.35 --character_threshold=0.35 --caption_extension=".txt" --max_data_loader_n_workers=2 --debug --undesired_tags=""

从上面的代码可以看到，

我们第一个传入的参数是训练集的路径。
--batch_size：表示每次传入Waifu Diffusion 1.4模型进行前向处理的数据数量。
--model_dir：表示加载的本地Waifu Diffusion 1.4模型路径。
--remove_underscore：如果开启，会将输出tag关键词中的下划线替换为空格。
--general_threshold：设置常规tag关键词的筛选置信度。
--character_threshold：设置人物特征tag关键词的筛选置信度。
--caption_extension：设置tag关键词标签的扩展名，一般为".txt"。
-max_data_loader_n_workers：设置大于等于2，加速数据处理。
--debug：如果设置，将会在Waifu Diffusion 1.4模型前向处理过程中，打印所有的图片路径与tag关键词标签内容，以供检查。
--undesired_tags：设置不需要输出的tag关键词。

(4) 补充标注特殊tag

完成了caption和tag的自动标注之后，如果我们需要训练一些特殊标注的话，还可以进行手动的补充标注。

SDXL-Trian项目中也提供了对数据进行补充标注的代码，下面代码将其进行提炼总结，方便大家直接使用。

大家可以直接拷贝以下的代码，并按照代码中提供的注释进行参数修改，然后运行代码即可对数据集进行补充标注：

import os
# 设置为本地的数据集路径
train_data_dir