Diffusers代码学习-Dreambooth

在Diffusers代码中，可以使用train_dreambooth.py脚本进行模型训练。

1. 下载源代码及安装。

在运行脚本之前，请确保从以下源安装库：

git clone https://github.com/huggingface/diffusers

cd diffusers

pip install .

导航到包含train_dreambooth.py脚本的示例文件夹，并为正在使用的脚本安装所需的依赖项：

cd examples/dreambooth

pip install -r requirements.txt

注意需要修改train_dreambooth.py，将accelerator_project_config改为project_config。否则会报错returned non-zero exit status 1.

2. 准备训练数据集

同时下载好训练所用的数据集。

import os

os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

from huggingface_hub import snapshot_download

local_dir = "./dog"

snapshot_download(

    "diffusers/dog-example",

    local_dir=local_dir,

    repo_type="dataset",

    ignore_patterns=".gitattributes",

)

注意如果下载的文件中除了图片以外还有其他目录比如.huggingface目录，需要删掉。否则会报错PermissionError: [Errno 13] Permission denied

3. 初始化accelerate

然后需要用到accelerate库，可帮助在多个GPU/TPU上或以混合精度进行训练。它将根据硬件和环境自动配置训练设置。

这里不选择任何配置，使用设置默认值运行accelerate库

accelerate config default

4. 运行训练脚本

accelerate launch train_dreambooth.py 

--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5"  （原始模型文件）

--instance_data_dir="dog"  （训练数据集目录名）

--output_dir="model_dog" （训练生成模型名称）

--instance_prompt="a photo of sks dog"  （提示词）

--resolution=512 

--train_batch_size=1 

--gradient_accumulation_steps=1 

--learning_rate=5e-6 

--lr_scheduler="constant" 

--lr_warmup_steps=0 

--max_train_steps=400

由于设定的训练步骤较多，需要等一段时间才能完成训练。