基于Diffusers的Stable Diffusion XL DreamBooth训练教程

苗恋蔷Samson

于 2025-06-01 09:02:34 发布

阅读量358

点赞数 4

本文链接：https://blog.csdn.net/gitblog_00479/article/details/148360441

版权

基于Diffusers的Stable Diffusion XL DreamBooth训练教程

diffusers Diffusers：在PyTorch中用于图像和音频生成的最先进扩散模型。项目地址: https://gitcode.com/gh_mirrors/di/diffusers

概述

DreamBooth是一种创新的文本到图像模型个性化方法，它仅需3-5张主题图片就能实现对Stable Diffusion等模型的定制化训练。本文将详细介绍如何使用Diffusers库中的train_dreambooth_lora_sdxl.py脚本来实现Stable Diffusion XL(SDXL)模型的DreamBooth训练。

技术背景

DreamBooth原理

DreamBooth通过在少量样本图像上微调预训练模型，使模型能够学习特定主题的视觉特征。其核心思想是将一个稀有词(如"sks")与目标主题关联，从而在生成时通过提示词控制输出。

LoRA技术

当前实现采用LoRA(Low-Rank Adaptation)技术进行参数高效微调：

仅微调SDXL的UNet部分
通过低秩矩阵分解减少可训练参数
保持原始模型权重不变，仅添加少量适配层

环境准备

依赖安装

建议创建新的Python虚拟环境并按以下步骤安装依赖：

pip install -e .  # 从源码安装diffusers
cd examples/dreambooth
pip install -r requirements_sdxl.txt

关键依赖包括：

PEFT(>=0.6.0)：用于LoRA训练的后端库
Accelerate：分布式训练支持
xformers(可选)：内存优化注意力机制
bitsandbytes(可选)：8位优化器

加速配置

初始化训练环境：

accelerate config

或使用默认配置：

accelerate config default

实战训练

数据准备

以狗的图像为例，首先下载示例数据集：

from huggingface_hub import snapshot_download
local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir, 
    repo_type="dataset",
    ignore_patterns=".gitattributes",
)

基础训练命令

accelerate launch train_dreambooth_lora_sdxl.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" \
  --instance_data_dir="dog" \
  --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
  --output_dir="lora-trained-xl" \
  --mixed_precision="fp16" \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --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=25 \
  --seed="0" \
  --push_to_hub

关键参数说明

resolution=1024：SDXL标准分辨率
gradient_accumulation_steps=4：梯度累积解决显存限制
pretrained_vae_model_name_or_path：指定优化后的VAE模型
validation_prompt：训练过程中验证生成效果

低显存优化方案

对于16GB以下显存的GPU，可添加以下参数：

--enable_xformers_memory_efficient_attention \
--gradient_checkpointing \
--use_8bit_adam \

推理生成

训练完成后，使用以下代码进行推理：

from diffusers import DiffusionPipeline
import torch

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", 
    torch_dtype=torch.float16
)
pipe.load_lora_weights("lora-trained-xl")
image = pipe("A picture of a sks dog in a bucket").images[0]

使用Refiner增强

SDXL提供专门的Refiner模型提升输出质量：

refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0", 
    torch_dtype=torch.float16
)
latents = pipe(...).images
refined_image = refiner(prompt=prompt, image=latents).images[0]