StableDiffusion模型LoRA微调

StableDiffusion是由StabilityAI、CompVis与Runway合作开发并开源的文本生成图像的模型。他可以直接用于文本生成图像的任务，也可以作为基础模型进行微调，从而从数据集上学习到新的风格，或是用于完成新的任务。本文将介绍通过在PAI完成LoRA微调StableDiffusion模型。

背景介绍

LoRA（Low-Rank Adaption of Large Language Model）是由微软提出的高效微调大语言模型的方法，他通过冻结原始模型参数，在模型上新增低秩矩阵作为可训练参数的方式微调模型。研究者发现，通过在Transformer块的Attention层上添加LoRA低秩矩阵对模型进行微调，能够获得与全参数微调水平相近的模型。相比于全参数的微调，LoRA有以下优点：

• 训练的参数量小，计算资源消耗低，训练速度更快。

• 对于计算资源/显存的要求更低，支持用户在消费级/中低端的GPU卡对大模型进行微调。

• 冻结了原始模型参数，在训练过程中不容易发生灾难性遗忘。

• 产出的模型较小，存储的成本较低，仅需推理时和原始的模型一同使用进行推理。

后续有开发者，将其引入到StableDiffsion模型的微调中，取得了不错的效果。HuggingFace提供的Diffusers库支持用户使用扩散模型进行训练或是推理，他支持用户使用LoRA微调扩散模型，并提供了相应的训练代码，支持文生图，以及DreamBooth的LoRA训练。

当前示例，我们将基于Diffusers库提供的训练代码和文档，在PAI完成StableDiffusion v1.5模型的LoRA微调训练。

准备工作

安装PAI Python SDK

安装PAI Python SDK，用于提交训练任务到PAI。

# 安装PAI Python SDK
!python -m pip install --upgrade alipai

SDK需要配置访问阿里云服务需要的AccessKey，以及当前使用的工作空间和OSS Bucket。在PAI SDK安装之后，通过在命令行终端 中执行以下命令，按照引导配置密钥、工作空间等信息。

# 以下命令，请在 命令行终端 中执行.

python -m pai.toolkit.config

我们可以执行以下代码，验证配置是否成功。

import pai
from pai.session import get_default_session

print(pai.__version__)

sess = get_default_session()

# 配置成功之后，我们可以拿到工作空间的信息
assert sess.workspace_name is not None
assert sess.oss_bucket is not None

获取PAI提供的StableDiffusion模型

PAI的公共模型仓库提供了StableDiffusion v1.5模型，用户可以通过以下代码获取模型的信息，用于后续的微调训练。

from pai.session import get_default_session
from pai.libs.alibabacloud_aiworkspace20210204.models import ListModelsRequest

sess = get_default_session()

# 获取PAI提供的StableDiffusion模型信息
resp = sess._acs_workspace_client.list_models(
    request=ListModelsRequest(
        provider="pai",
        model_name="stable_diffusion_v1.5",
    )
)
model = resp.body.models[0].latest_version

# StableDiffusion 模型的OSS URI（公共读）
print(f"StableDiffusion ModelUri: {model.uri}")

LoRA TextToImage微调训练

通过LoRA训练StableDiffusion模型，可以快速，低成本地获得一个能够生成指定风格的模型。在以下示例中，我们将使用一个Demo的图像文本数据集，对StableDiffusion模型进行LoRA微调。

准备训练数据

当前示例准备了一个简单的文本图片数据集在train-data目录下，包含训练的图片以及相应的标注文件(metadata.jsonl)。

!ls -lh train-data/
!cat train-data/metadata.jsonl

我们需要将数据上传到OSS Bucket上，供训练作业使用。

from pai.common.oss_utils import upload

train_data_uri = upload("./train-data/", "stable_diffusion_demo/text2image/train-data/")
print(train_data_uri)

Diffuerser提供的训练脚本默认使用ImageFolder格式的数据集，用户可以参考以上的格式准备数据，更加详细的介绍可以见HuggingFace datasets的ImageFolder数据集文档。

准备训练作业脚本

我们将使用Diffusers库提供的训练作业脚本(train_text_to_image_lora.py)完成LoRA训练。执行以下代码，我们将代码下载到本地，用于后续提交训练任务。

!mkdir -p train_lora

# code source: https://github.com/huggingface/diffusers/blob/v0.17.1/examples/text_to_image/train_text_to_image_lora.py
!wget -P train_lora https://raw.githubusercontent.com/huggingface/diffusers/v0.17.1/examples/text_to_image/train_text_to_image_lora.py

我们提交的训练作业将使用PAI提供的PyTorch 1.12的GPU镜像运行，我们需要准备一个requirements.txt文件在训练代码目录下，以安装一些额外的依赖包。

训练脚本目录提交到PAI上执行训练时，目录下的requirements.txt文件将被安装到作业环境中。

%%writefile train_lora/requirements.txt

diffusers>=0.17.1

# source: https://github.com/huggingface/diffusers/blob/v0.17.1/examples/text_to_image/requirements.txt
accelerate>=0.16.0,<=0.18.0
torchvision
transformers>=4.25.1,<5.0.0
datasets
ftfy
tensorboard
Jinja2

提交训练作业

Diffuers提供的训练脚本，需要使用Accelerate工具启动，并通过命令行参数的方式，传递超参，预训练模型路径/ID，以及训练数据集地址。PAI的训练作业，支持通过环境变量的方式获取输入输出的数据/模型路径，以及训练作业超参。以下脚本中，我们通过环境变量的方式，传递超参、输入输出路径给到训练脚本。

from pai.image import retrieve

# 使用PAI提供的PyTorch 1.12 GPU镜像
image_uri = retrieve(
    "PyTorch",
    "1.12",
    accelerator_type="GPU",
).image_uri

print(image_uri)


# 训练作业启动命令，通过环境变量的方式获取:
# a）输入输出的模型/数据路径
# b）训练任务的超参数
command = """accelerate launch --mixed_precision="fp16"  train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$PAI_INPUT_PRETRAINED_MODEL  \
  --train_data_dir=$PAI_INPUT_TRAIN_DATA \
  --output_dir=$PAI_OUTPUT_MODEL \
  --logging_dir=$PAI_OUTPUT_TENSORBOARD \
  --dataloader_num_workers=8 \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=$PAI_HPS_TRAIN_BATCH_SIZE \
  --gradient_accumulation_steps=$PAI_HPS_GRADIENT_ACCUMULATION_STEPS \
  --max_train_steps=$PAI_HPS_MAX_TRAIN_STEPS \
  --learning_rate=$PAI_HPS_LEARNING_RATE \
  --checkpointing_steps=$PAI_HPS_CHECKPOINTING_STEPS \
  --max_grad_norm=1 \
  --lr_scheduler="cosine" --lr_warmup_steps=0 \
  --validation_prompt="$PAI_HPS_VALIDATION_PROMPT" \
  --validation_epochs=$PAI_HPS_VALIDATION_EPOCHS \
  --seed=$PAI_HPS_SEED"""


# 训练作业超参
hps = {
    "validation_prompt": "a photo of cat in a bucket",  # 验证模型的Prompt
    "validation_epochs": 1,  # 每隔50个epoch验证一次
    "max_train_steps": 10,  # 最大训练步数
    "learning_rate": 1e-4,  # 学习率
    "train_batch_size": 2,  # 训练batch size
    "gradient_accumulation_steps": 1,  # 梯度累积步数
    "checkpointing_steps": 5,  # 每隔100个step保存一次模型
    "seed": 1337,  # 随机种子
}

以下代码中，我们使用Estimator类，指定训练作业使用的镜像，训练作业超参，输入数据路径等，将LoRA训练作业提交到PAI执行。

对于使用SDK提交训练作业的详细介绍，用户可以参考文档: 提交训练作业。

from pai.estimator import Estimator
from pai.image import retrieve

# 使用PAI提供的PyTorch 1.12 GPU镜像
image_uri = retrieve(
    "PyTorch",
    "1.12",
    accelerator_type="GPU",
).image_uri

print(image_uri)


# 训练作业启动命令，通过环境变量的方式获取:
# a）输入输出的模型/数据路径
# b）训练任务的超参数

command = """accelerate launch --mixed_precision="fp16"  train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$PAI_INPUT_PRETRAINED_MODEL  \
  --train_data_dir=$PAI_INPUT_TRAIN_DATA \
  --output_dir=$PAI_OUTPUT_MODEL \
  --logging_dir=$PAI_OUTPUT_TENSORBOARD \
  --dataloader_num_workers=8 \
  --resolution=512 --center_crop --random_flip \
  --train_batch_size=$PAI_HPS_TRAIN_BATCH_SIZE \
  --gradient_accumulation_steps=$PAI_HPS_GRADIENT_ACCUMULATION_STEPS \
  --max_train_steps=$PAI_HPS_MAX_TRAIN_STEPS \
  --learning_rate=$PAI_HPS_LEARNING_RATE \
  --checkpointing_steps=$PAI_HPS_CHECKPOINTING_STEPS \
  --max_grad_norm=1 \
  --lr_scheduler="cosine" --lr_warmup_steps=0 \
  --validation_prompt="$PAI_HPS_VALIDATION_PROMPT" \
  --validation_epochs=$PAI_HPS_VALIDATION_EPOCHS \
  --seed=$PAI_HPS_SEED"""


# 训练作业超参
hps = {
    "validation_prompt": "a photo of cat in a bucket",  # 验证模型的Prompt
    "validation_epochs": 1,  # 每隔50个epoch验证一次
    "max_train_steps": 10,  # 最大训练步数
    "learning_rate": 1e-4,  # 学习率
    "train_batch_size": 2,  # 训练batch size
    "gradient_accumulation_steps": 1,  # 梯度累积步数
    "checkpointing_steps": 5,  # 每隔100个step保存一次模型
    "seed": 1337,  # 随机种子
}


est = Estimator(
    image_uri=image_uri,  # 训练作业使用的镜像
    source_dir="train_lora",  # 训练代码路径，代码会被上传，并准备到训练作业环境中
    command=command,  # 训练任务启动命令
    instance_type="ecs.gn6i-c4g1.xlarge",  # 4 vCPU, 16 GiB 内存, 1 x NVIDIA T4 GPU
    base_job_name="sd_lora_t2i_",  # 训练作业名称前缀
    hyperparameters=hps,  # 作业超参，训练命令和脚本可以通过 `PAI_HPS_{HP_NAME_UPPER_CASE}` 环境变量，或是读取`/ml/input/config/hpyerparameters.json`文件获取
)

使用inputs参数指定准备到训练作业环境的模型和数据，提交训练作业。

inputs参数是一个字典，Key是输入的名称，Value是输入数据的存储路径（例如OSS URI)。相应的数据会被准备到作业执行环境中（通过挂载的方式），训练作业脚本，能够通过环境变量PAI_INPUT_{KeyUpperCase}获取到输入数据的路径，通过读取本地文件的方式读取预训练模型和数据。

print("Input PreTrainedModel: ", model.uri)
print("Input TrainData: ", train_data_uri)


# 提交训练作业
est.fit(
    inputs={
        "pretrained_model": model.uri,
        "train_data": train_data_uri,
    },
    wait=True,
)

在启动命令中，我们使用--output_dir=$PAI_OUTPUT_MODEL，让训练脚本将模型写出到指定的输出目录中。对应的模型数据会被保存到用户的OSS Bucket中，我们可以通过est.model_data()获得输出的模型的OSS URI。

import os
from pai.common.oss_utils import download

print("OutputModel Path: ", est.model_data())
lora_weight_uri = os.path.join(est.model_data(), "pytorch_lora_weight.bin")
lora_model_path = download(oss_path=lora_weight_uri, local_path="./lora_model/")

以上训练获得LoRA模型，可以使用diffuser的推理pipeline加载使用：

# StableDiffusionPipeline加载LoRA模型


import torch
from diffusers import StableDiffusionPipeline

# 加载基础模型
model_id_or_path = "<SdModelId_Or_Path>"
pipe = StableDiffusionPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)

# 加载LoRA模型
pipe.unet.load_attn_procs(lora_model_path)

# 使用推理pipeline
image = pipe(
    "A pokemon with blue eyes.", num_inference_steps=25, guidance_scale=7.5,
    cross_attention_kwargs={"scale": 0.5},
).images[0]

或则用户也可以将其转为safetensor格式，在StableDiffusiuson WebUI中使用。

import torch
from safetensors.torch import save_file

# 加载模型
lora_model = torch.load(lora_model_bin_path, map_location="cpu")

# 转换为safetensor格式
save_file(lora_model, "lora.safetensors")

LoRA && DreamBooth微调训练

DreamBooth简介

DreamBooth是Google的研究人员于2022年提出的技术，支持在少量的图片上进行训练，然后将自定义的主题注入到扩散模型中。

图片来源: https://dreambooth.github.io/

直接使用少量的图片文本数据集对扩散模型进行训练容易导致过拟合，或是语言漂移。DreamBooth使用以下方式避免了模型的退化：

• 用户需要为新的主题选择一个罕见的词（标识符），模型将在训练过程中将这个词和图片的主题进行关联。

• 为了避免过拟合和语言漂移，微调过程中，使用相同类别的图片参与训练（这些图片可以由扩散模型生成）。

对于DreamBooth的详细介绍，用户可以参考DreamBooth的博客，以及HuggingFace博客对于DreamBooth的介绍。

当通过DreamBooth训练扩散模型时，用户可以选择进行普通的微调（直接微调原始的模型参数），也可以使用LoRA的方式进行微调，在以下示例中，我们将使用LoRA的方式，进行DreamBooth训练。

准备训练数据集

为了训练DreamBooth，用户需要准备特定风格的图片数据集，当前示例中，我们准备了数据集在sks-dog目录下。

通过以下代码，我们将将数据集上传到OSS上，供训练作业使用。

from pai.common.oss_utils import upload

train_data_uri = upload("sks-dog", "stable_diffusion/dreambooth/train-sks-dog/")

准备训练代码

我们使用HuggingFace Diffusers库提供的训练脚本，通过LoRA && DreamBooth方式训练扩散模型。通过以下代码，我们下载训练脚本到本地。

# 创建训练脚本保存路径
!mkdir -p train_dreambooth/

# 下载HuggingFace diffusers(v1.17.1)库提供的示例代码（因为访问GitHub的网络并不稳定，用户当出现下载失败，可以多尝试几次）
!wget https://raw.githubusercontent.com/huggingface/diffusers/v0.17.1/examples/dreambooth/train_dreambooth_lora.py -P train_dreambooth/

训练作业将使用PAI提供的PyTorch镜像运行脚本，我们需要通过以下的requirements.txt安装训练脚本依赖的库。

%%writefile train_dreambooth/requirements.txt
# %%writefile 指令会将当前内容写入到 train_dreambooth/requirements.txt 文件中

diffusers>=0.17.1

# source: https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/requirements.txt
accelerate>=0.16.0,<=0.18.0     # diffusers 提供的示例代码(v0.17.1)无法运行在accelerate>=0.18.0上.
torchvision
transformers>=4.25.1,<5.0.0
ftfy
tensorboard
Jinja2

提交训练作业

通过以下代码，我们使用PAI Python SDK，提交训练作业到PAI。

from pai.estimator import Estimator
from pai.image import retrieve

image_uri = retrieve(
    "PyTorch",
    "latest",
    accelerator_type="GPU",
).image_uri


# 训练作业启动命令，通过环境变量的方式获取:
# a）输入输出的模型/数据路径
# b）训练任务的超参数
command = """accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$PAI_INPUT_PRETRAINED_MODEL  \
  --instance_data_dir=$PAI_INPUT_TRAIN_DATA \
  --output_dir=$PAI_OUTPUT_MODEL \
  --logging_dir=$PAI_OUTPUT_TENSORBOARD \
  --instance_prompt="$PAI_HPS_INSTANCE_PROMPT" \
  --resolution=512 \
  --train_batch_size=$PAI_HPS_TRAIN_BATCH_SIZE \
  --gradient_accumulation_steps=$PAI_HPS_GRADIENT_ACCUMULATION_STEPS \
  --checkpointing_steps=$PAI_HPS_CHECKPOINTING_STEPS \
  --learning_rate=$PAI_HPS_LEARNING_RATE \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=$PAI_HPS_MAX_TRAIN_STEPS \
  --validation_prompt="$PAI_HPS_VALIDATION_PROMPT" \
  --validation_epochs=$PAI_HPS_VALIDATION_EPOCHS \
  --seed="0"
  """

# 训练作业超参
hps = {
    "instance_prompt": "a photo of sks dog",  # 训练的图片数据文本使用的标注Prompt。这里的sks是我们使用的数据集的特定风格标识符。
    "validation_prompt": "a photo of sks dog in a bucket",  # 验证模型的Prompt
    # "class_prompt": "a photo of dog",  # 用于生成类别图片数据，避免模型过拟合&&语言偏移
    "validation_epochs": 50,  # 每隔50个epoch验证一次
    "max_train_steps": 500,  # 最大训练步数
    "learning_rate": 1e-4,  # 学习率
    "train_batch_size": 1,  # 训练batch size
    "gradient_accumulation_steps": 1,  # 梯度累积步数
    "checkpointing_steps": 100,  # 每隔100个step保存一次模型
}


est = Estimator(
    image_uri=image_uri,  # 训练作业使用的镜像
    source_dir="train_dreambooth",  # 训练代码路径，代码会被上传，并准备到训练作业环境中
    command=command,  # 训练任务启动命令
    instance_type="ecs.gn6i-c4g1.xlarge",  # 4 vCPU, 16 GiB 内存, 1 x NVIDIA T4 GPU
    base_job_name="sd_lora_dreambooth_",  # 训练作业名称前缀
    hyperparameters=hps,  # 作业超参，训练命令和脚本可以通过 `PAI_HPS_{HP_NAME_UPPER_CASE}` 环境变量，或是读取`/ml/input/config/hpyerparameters.json`文件获取
)

print("Input PreTrainedModel: ", model.uri)
print("Input TrainData: ", train_data_uri)

est.fit(
    inputs={
        "pretrained_model": model.uri,
        "train_data": train_data_uri,
    },
    wait=True,
)

训练任务会在输出目录下，生成一个pytorch_lora_weights.bin的模型文件，相应的文件会被上传准备到用户的OSS中，用户可以通过以下的代码，将模型文件下载到本地。

import os
import posixpath

from pai.common.oss_utils import download

# 输出模型路径
output_lora_model = posixpath.join(est.model_data(), "pytorch_lora_weights.bin")
print("OutputModel: ", output_lora_model)

model_path = download(output_lora_model, "./lora_dreambooth_model/")

获得的LoRA模型，用户可以通过Diffuser提供的API，在推理pipeline加载使用，用户可以参考diffuser的文档：DreamBooth Inference。

结语

通过当前示例，我们介绍了如何基于HuggingFace diffusers库，在PAI上完成StableDiffusion模型的LoRA微调训练。用户可以通过Diffuers库的API，直接在推理Pipeline中加载使用这些LoRA模型，也可以将模型转换Safetensors格式，用于StableDiffusionWebUI中。

除了对于LoRA的支持，Diffusers库支持对于直接对扩散模型微调，也支持包括TextInversion, ControlNet, InstructPix2Pix等方式训练扩散模型，并且提供了相应的训练脚本和教程。用户同样可以参考本示例，在PAI运行这些训练任务。

参考

• HuggingFace LoRa Tutorial: https://huggingface.co/docs/diffusers/training/lora#texttoimage

• HuggingFace LoRA Blog: https://huggingface.co/blog/lora

• Google DreamBooth Blog：https://dreambooth.github.io/