Dreambooth-Stable-Diffusion 使用教程

引言

Dreambooth 是一种基于深度学习的个性化模型微调方法，能够将特定的概念（例如某个人、物品或场景）融入到生成模型中。在 Stable Diffusion 的基础上，使用 Dreambooth 可以轻松生成更具针对性和个性化的图像。

本教程将手把手教您如何使用 Dreambooth 微调 Stable Diffusion 模型。我们将包括环境配置、数据准备、模型训练、以及生成效果的展示。

一、环境准备

在开始之前，确保您具备以下条件：

1. 操作系统：Linux/Windows/MacOS
2. Python：3.8 或更高版本
3. GPU：NVIDIA GPU（推荐 12GB 显存以上）

1.1 安装依赖

克隆项目代码

# 克隆 Dreambooth 项目代码
git clone https://github.com/XavierXiao/Dreambooth-Stable-Diffusion.git
cd Dreambooth-Stable-Diffusion

安装 Python 依赖

# 创建虚拟环境
python3 -m venv dreambooth_env
source dreambooth_env/bin/activate

# 安装依赖库
pip install torch torchvision transformers accelerate diffusers
pip install -r requirements.txt

安装 GPU 支持的 PyTorch

确保安装与您的 CUDA 版本兼容的 PyTorch。

pip install torch --index-url https://download.pytorch.org/whl/cu118

二、数据准备

2.1 收集训练数据

• 收集 10-20 张目标对象的高质量图像（分辨率越高越好）。
• 图片格式为 .jpg 或 .png。
• 将所有图片保存在一个文件夹中，例如 ./training_data/。

2.2 图像标注

为每张图片指定唯一的标注词（如 photo of sks person）。标注需要与训练时的 prompt 保持一致。

三、模型训练

3.1 下载预训练模型

从 Hugging Face 下载 Stable Diffusion 的基础权重。

mkdir -p pretrained_models/stable-diffusion
cd pretrained_models/stable-diffusion
# 下载权重
wget https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/model.ckpt

3.2 启动训练

运行以下命令，开始训练模型。

python train_dreambooth.py \
  --pretrained_model_name_or_path="pretrained_models/stable-diffusion" \
  --instance_data_dir="./training_data" \
  --output_dir="./dreambooth_output" \
  --instance_prompt="photo of sks person" \
  --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

参数解释

• --pretrained_model_name_or_path：指定预训练模型路径。
• --instance_data_dir：训练数据路径。
• --output_dir：训练输出路径。
• --instance_prompt：训练图片的描述文字。
• --max_train_steps：训练步数，推荐 400-800。

训练完成后，模型权重将保存在 dreambooth_output 文件夹中。

四、生成图片

4.1 加载微调模型

将训练好的模型加载到 Stable Diffusion 中进行生成。

from diffusers import StableDiffusionPipeline
import torch

# 加载微调模型
model_path = "./dreambooth_output"
pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16).to("cuda")

# 文本生成图像
prompt = "photo of sks person in a magical forest"
image = pipeline(prompt, guidance_scale=7.5).images[0]

# 保存生成的图片
image.save("generated_image.png")

4.2 参数调整

• guidance_scale：控制生成图像与 prompt 的一致性，推荐值 7.5。
• prompt：可以根据需求修改生成描述，例如添加背景或风格信息。

五、效果展示

以下是训练和生成的效果对比：

5.1 原始图片集

目标对象的图片如下：

原始图片

5.2 生成图片

微调后的模型能够生成如下个性化图像：

生成图片

六、常见问题与解决方案

6.1 CUDA Out of Memory

问题：显存不足。

解决方案：

• 降低 train_batch_size。
• 使用 gradient_accumulation_steps 增加累计梯度。
• 使用显存优化工具，如 torch.utils.checkpoint。

6.2 生成图片质量不佳

问题：训练数据不足或训练步数太少。

解决方案：

• 增加训练数据。
• 调整 max_train_steps。

七、总结

通过本教程，您已经掌握了如何使用 Dreambooth 微调 Stable Diffusion 模型，生成个性化、高质量的图像。Dreambooth 的强大之处在于其灵活性和可扩展性，适用于多种创作场景。