使用HuggingFace Diffusers实现LoRA高效微调Stable Diffusion模型-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01050/article/details/148549327

使用HuggingFace Diffusers实现LoRA高效微调Stable Diffusion模型

notebooks Notebooks using the Hugging Face libraries 🤗 项目地址: https://gitcode.com/gh_mirrors/note/notebooks

什么是LoRA技术

LoRA（Low-Rank Adaptation of Large Language Models）是一种高效微调大型模型的技术，它通过引入低秩分解权重矩阵对（称为更新矩阵）来优化模型微调过程。这项技术具有几个显著优势：

参数高效性：仅训练新增的低秩权重矩阵，保持预训练权重不变
防止灾难性遗忘：原始模型参数被冻结，避免了微调过程中的知识丢失问题
硬件友好：可在消费级GPU（如T4、RTX 3080等）上运行
灵活控制：通过scale参数调节新知识的融入程度

LoRA在Diffusers中的应用

目前Diffusers中LoRA主要支持UNet2DConditionalModel的注意力层。这种设计源于研究发现，仅修改语言模型的注意力层就能获得良好的性能提升。

技术实现要点

Diffusers提供了load_attn_procs()方法专门用于将LoRA权重加载到模型的注意力层。关键参数scale控制着模型适应新训练图像的程度：

scale=0：完全使用原始模型权重
scale=1：完全使用微调后的LoRA权重
0<scale<1：两种权重的线性插值

实战：使用LoRA微调Stable Diffusion

准备工作

我们以Pokémon BLIP标注数据集为例，微调stable-diffusion-v1-5模型生成自定义宝可梦图像。

首先设置环境变量：

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/pokemon"
export HUB_MODEL_ID="pokemon-lora"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

训练配置

训练脚本支持多个重要参数：

--push_to_hub：将训练结果推送到模型库
--report_to=wandb：记录训练过程到Weights & Biases
--learning_rate=1e-04：LoRA通常可以使用较高学习率

完整训练命令示例：

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=50 \
  --seed="0" \
  --push_to_hub

推理使用

训练完成后，可以加载基础模型并注入LoRA权重进行推理：

import torch
from diffusers import StableDiffusionPipeline

# 加载基础模型
model_base = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_base, torch_dtype=torch.float16)

# 注入LoRA权重
pipe.unet.load_attn_procs(model_path)
pipe.to("cuda")

# 生成图像（使用50%LoRA权重）
image = pipe(
    "A picture of a sks dog in a bucket.",
    num_inference_steps=25,
    guidance_scale=7.5,
    cross_attention_kwargs={"scale": 0.5},
).images[0]

# 保存结果
image.save("bucket-dog.png")