SDXL Trainer

一个魔改 kohya-ss/sd-scripts 的 SDXL 微调的项目。

特性

在 kohya-ss 的基础上，对训练脚本进行了优化，主要包括：

1. 重写了数据集类： a. 实现了数据库读取，从 json 文件中读取需要的所有参数； b. 实现了多 GPU 并行缓存潜变量； c. 实现了异步潜变量缓存写入，通过异步 IO 来提高缓存潜变量时的 GPU 利用率，从而加速缓存； d. 实现了缓存文件有效性检查； e. 支持多线程数据加载； f. 支持仅使用缓存文件，而不需要图像文件的训练； g. 支持非一次性的缓存读取，而是仅在开始时检查缓存的有效性，随训练逐步读取，降低了从运行脚本到开始训练的等待时间；
2. 支持单独缓存潜变量：仅加载 VAE 来缓存潜变量，而非在训练开始时加载所有组件后再缓存。
3. 更美观的控制台日志
4. 更安全的模型保护机制：在训练出错/人为打断训练时自动保存模型文件，防止训练进度丢失。
5. 更灵活的配置文件：使用配置文件来配置训练参数，支持多种自定义参数。
6. 更自由的训练策略定制：提供训练策略自定义方式，支持自定义训练策略，包括自定义每个数据的重复次数、自定义标注处理方式等。

使用方法

安装

git clone https://github.com/Eugeoter/sdxl-trainer.git
cd sdxl-trainer
pip install -r requirements.txt

训练

准备数据

准备好图像（或潜变量缓存）文件和同名的标注文件（或整一个元数据 json 文件），将图像（或潜变量缓存）文件放置在同一文件夹 image_dir 内。

训练器会加载配置中 image_dirs 和 metadata_files 所指向的数据。其中，image_dirs 是一个字符串列表，包含了所有的图像文件夹路径。metadata_files 是一个字符串列表，包含了所有的元数据文件路径。

数据加载方式有两种，一种是加载图像和与图像同名的 txt 文本文件作为标注，另一种是直接从元数据文件中加载标注。具体使用哪一种方式取决于 metadata_files 是否为空。

• 当 metadata_files 为空时，训练器会加载 image_dirs 中的所有图像文件，并读取与图像同名的 txt 文件作为标注。
• 当 metadata_files 不为空时，训练器会从 metadata_files 中的元数据文件中读取每个数据的信息，包括了图像路径、标注、原始尺寸等等。该方式无需为每个图像文件创建一个 txt 文件作为标注，但需要额外的元数据文件。 当某个元数据中的图像路径不存在时，训练器会尝试从 image_dirs 中查找文件名（不包含路径）相同的图像文件。

无论以哪种加载数据，请确保所加载的数据中没有去除后缀后重名的图像文件，否则会导致数据加载错误。

您可以直接使用由缓存潜变量生成的 npz 缓存文件代替原始图像文件。

以下几种数据组织方式均可：

a. 图像 + 图像同名 txt 标注文件：

  image_dir
  ├── image1.jpg
  ├── image1.txt
  ├── image2.jpg
  ├── image2.txt
  ├── ...

b. 缓存潜变量 + 潜变量同名 txt 标注文件：

  latent_dir
  ├── image1.npz
  ├── image1.txt
  ├── image2.npz
  ├── image2.txt
  ├── ...

元数据文件

元数据文件是一个 json 文件，包含了一些数据的各种信息，如图像路径、标注、图像尺寸、美学评分等。训练时，可以用一个或多个元数据文件来描述数据集。

使用元数据文件替代传统 txt 标注能简化数据集管理，不再需要为每个图像文件创建一个 txt 文件作为标注。其能够存储更多种类的数据信息。

制作元数据文件可使用 Waifuset 项目 或 kohya-ss。是时候摆脱 txt 标注文件了。

元数据 json 文件格式如下：

{
	"image1": {
			"image_path": ".../image1.jpg", // 图像路径
			"caption": "1girl, solo, ...", // 标注
      ... // 其他元数据
    },
  "image2": {
      "image_path": ".../image2.jpg",
      "caption": "1girl, solo, ...",
      ...
    },
  ...
}

c. 图像 + 元数据 json 文件：

  image_dir
  ├── image1.jpg
  ├── image2.jpg
  ├── ...
  metadata.json

d. 缓存潜变量 + 元数据 json 文件：

  latent_dir
  ├── image1.npz
  ├── image2.npz
  ├── ...
  metadata.json

配置参数

在 configs/train_config.py 内配置参数（或自行新建）后，执行 accelerate launch sdxl_train.py --config configs/train_config.py 即可开始训练。

其中，--config 参数为配置文件路径，configs/train_config.py 为默认配置文件，您需要根据自己的需求修改配置文件，或指定您自己修改的配置文件路径。

完整的参数介绍请参考 docs/CONFIG.md。

多 GPU 训练

将原本参数中的 accelerate launch sdxl_train.py --config configs/train_config.py 更换为 accelerate launch --num_processes=4 --multi_gpu --gpu_ids=0,1,2,3 sdxl_train.py --config configs/train_config.py 即可进行多 GPU 训练，其中 --num_processes 为进程数，--gpu_ids 为 GPU 编号。

缓存潜变量

建议单独提前缓存潜变量，以加速训练。

缓存潜变量通过单独的脚本 cache_latents.py 运行。您可使用与训练相同的配置文件执行 accelerate launch cache_latents.py --config configs/train_config.py 缓存潜变量，缓存后的潜变量文件将被储存为 {图像文件名}.npz 的格式。您也可以用类似的方式进行多 GPU 缓存。