Dreambooth-Stable-Diffusion：基于 Stable Diffusion 的个性化图像生成工具

本文介绍 Dreambooth-Stable-Diffusion，一个基于 Stable Diffusion 实现 Google Dreambooth 论文的开源项目。该项目允许用户仅用几张图片就能微调文本到图像的扩散模型，生成个性化的图像内容。

项目简介

Dreambooth-Stable-Diffusion 是 Google Dreambooth 论文的开源实现，基于 Stable Diffusion 模型。原始 Dreambooth 基于 Imagen 文本到图像模型，但 Imagen 模型和预训练权重并未公开。为了让更多人能够使用少量样本微调文本到图像模型，作者将 Dreambooth 的思想移植到了 Stable Diffusion 上。

该项目在 GitHub 上已获得 7.8k stars，主要使用 Jupyter Notebook 和 Python 编写，采用 MIT 许可证开源。

核心特性

完整模型微调：与 Textual Inversion 仅优化词嵌入不同，Dreambooth 微调整个扩散模型
低样本学习：仅需几张目标图片即可训练个性化模型
保持可编辑性：在学习新概念的同时保留模型的编辑能力
正则化支持：使用正则化图像防止过拟合
GPU 内存优化：支持梯度检查点技术减少显存占用

技术栈

PyTorch：深度学习框架
PyTorch Lightning：训练流程管理
Stable Diffusion：基础扩散模型
Jupyter Notebook：交互式开发和演示

安装指南

前置要求

Python 3.8+
CUDA 支持的 GPU（推荐使用 A6000 或同等配置）
至少 16GB 显存

环境配置

克隆仓库并设置环境：

bash

git clone https://github.com/XavierXiao/Dreambooth-Stable-Diffusion.git
cd Dreambooth-Stable-Diffusion

按照 Textual Inversion 或原始 Stable Diffusion 的说明设置 ldm 环境。
下载预训练的 Stable Diffusion 模型权重（推荐 sd-v1-4-full-ema.ckpt）从 HuggingFace 获取。

快速开始

1. 准备正则化图像

Dreambooth 需要正则化图像来防止过拟合。使用以下命令生成：

bash

python scripts/stable_txt2img.py \
  --ddim_eta 0.0 \
  --n_samples 8 \
  --n_iter 1 \
  --scale 10.0 \
  --ddim_steps 50 \
  --ckpt /path/to/sd-v1-4-full-ema.ckpt \
  --prompt "a photo of a <class>"

建议生成 100-200 张正则化图像以获得更好的效果。

2. 训练模型

bash

python main.py \
  --base configs/stable-diffusion/v1-finetune_unfrozen.yaml \
  -t \
  --actual_resume /path/to/sd-v1-4-full-ema.ckpt \
  -n <job_name> \
  --gpus 0, \
  --data_root /root/to/training/images \
  --reg_data_root /root/to/regularization/images \
  --class_word <xxx>

训练参数说明：

--data_root：训练图像目录（包含目标对象的图片）
--reg_data_root：正则化图像目录
--class_word：类别词（如 dog、container 等）
默认学习率：1.0e-6
默认正则化权重：1.0
训练步数：800 步（在 500 步和 800 步各保存一个检查点）

3. 生成图像

训练完成后，使用以下命令生成个性化图像：

bash

python scripts/stable_txt2img.py \
  --ddim_eta 0.0 \
  --n_samples 8 \
  --n_iter 1 \
  --scale 10.0 \
  --ddim_steps 100 \
  --ckpt /path/to/saved/checkpoint \
  --prompt "photo of a sks <class>"

其中 sks 是默认的标识符（identifier），可以在 ldm/data/personalized.py 中修改。

使用示例

示例：训练自定义容器

假设你有几张特定容器的照片，想要生成该容器在不同场景下的图像：

准备训练数据：收集 3-5 张目标容器的照片

生成正则化图像：

bash

python scripts/stable_txt2img.py --prompt "photo of a container" --n_samples 200 ...

训练：

bash

python main.py --class_word container --data_root ./my_container ...

生成场景图像：
- "photo of a sks container on the beach" - 海滩上的容器
- "photo of a sks container on the moon" - 月球上的容器
- "photo of a red sks container" - 红色的容器
- "a dog on top of sks container" - 容器上的狗

注意事项

标识符选择：默认使用 sks 作为标识符，如需更改请修改 ldm/data/personalized.py 第 10 行
正则化图像数量：建议使用 100-200 张正则化图像以获得最佳效果
学习率调整：默认 1.0e-6 比论文中的 1.0e-5 效果更好，能保持更好的可编辑性
显存优化：代码已启用梯度检查点，可在配置文件中调整
** unrealistic 正则化图像**：当生成"man"或"woman"等类别时，如果生成的正则化图像不够真实，可以在线搜索多样化的真实图像作为替代

项目链接

GitHub 仓库：https://github.com/XavierXiao/Dreambooth-Stable-Diffusion
Dreambooth 论文：https://arxiv.org/abs/2208.12242
Stable Diffusion：https://github.com/CompVis/stable-diffusion
Textual Inversion：https://github.com/rinongal/textual_inversion

字节笔记本

Dreambooth-Stable-Diffusion：基于 Stable Diffusion 的个性化图像生成工具

项目简介

核心特性

技术栈

安装指南

前置要求

环境配置

快速开始

1. 准备正则化图像

2. 训练模型

3. 生成图像

使用示例

示例：训练自定义容器

注意事项

项目链接