Dreambooth 在 Stable Diffusion 上的实现

这是基于 Google 的 Dreambooth 和 Stable Diffusion 的实现。原始的 Dreambooth 是基于 Imagen 文本到图像模型的，然而 Imagen 的模型和预训练权重均未公开。为了使人们能够仅用少量示例就微调文本到图像模型，我在 Stable Diffusion 上实现了 Dreambooth 的思想。

该代码仓库基于 Textual Inversion 的代码库。需要注意的是，Textual Inversion 只优化词嵌入，而 Dreambooth 则对整个扩散模型进行微调。

该实现对 Textual Inversion 的官方代码库进行了最小的改动。事实上，由于懒惰，Textual Inversion 中的一些组件，比如嵌入管理器，并没有被删除，尽管它们在这里根本不会被使用。

更新

2022年9月20日：我刚刚找到了一种稍微减少 GPU 内存占用的方法。请记住，这段代码是基于 Textual Inversion 的，而 TI 的代码库中有一行代码这里，以硬编码的方式禁用了梯度检查点技术。这是因为在 TI 中，Unet 并没有被优化。然而，在 Dreambooth 中，我们优化了 Unet，因此可以启用梯度检查点技术，就像在原始 SD 仓库中的这里一样。在配置文件中，默认情况下梯度检查点是开启的。我已经更新了代码。

使用方法

准备工作

首先按照 Textual Inversion 仓库或原始 Stable Diffusion 仓库的说明设置 ldm 环境。

要微调一个 Stable Diffusion 模型，你需要按照他们的说明获取预训练的 Stable Diffusion 模型。权重可以在 HuggingFace 上下载。你可以选择使用哪个版本的检查点，但我使用的是 sd-v1-4-full-ema.ckpt。

我们还需要创建一组用于正则化的图片，因为 Dreambooth 的微调算法需要这些图片。算法的详细信息可以在论文中找到。需要注意的是，在原始论文中，正则化图片似乎是实时生成的。然而，在这里，我在训练之前就生成了一组正则化图片。生成正则化图片的文本提示可以是 photo of a <class>，其中 <class> 是描述你对象类别的词，比如 dog。命令如下：

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

我生成了 8 张用于正则化的图片，但更多的正则化图片可能会带来更强的正则化效果和更好的可编辑性。之后，将生成的图片分别保存为单个 .png 文件，存放在 /root/to/regularization/images 目录下。

9月9日更新 我们当然应该使用更多的正则化图片。请尝试使用 100 或 200 张，以便更好地符合原始论文的要求。为此，我在配置文件中缩短了正则化数据集的“重复”次数。

在某些情况下，如果生成的正则化图片非常不真实（当你想生成“man”或“woman”时就会发生），你可以在网上找到一组多样化的图片（关于男人/女人），并将它们用作正则化图片。

训练

训练可以通过运行以下命令来进行：

python main.py --base configs/stable-diffusion/v1-finetune_unfrozen.yaml 
                -t 
                --actual_resume /path/to/original/stable-diffusion/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>

详细的配置可以在 configs/stable-diffusion/v1-finetune_unfrozen.yaml 中找到。特别是，我将默认学习率设置为 1.0e-6，因为发现 Dreambooth 论文中使用的 1.0e-5 会导致较差的可编辑性。参数 reg_weight 对应于 Dreambooth 论文中的正则化权重，默认值为 1.0。

Dreambooth 需要一个占位符单词 [V]，称为标识符，正如论文中所述。这个标识符需要是词汇表中相对罕见的标记。原始论文通过使用 T5-XXL 分词器中的一个罕见词来实现这一点。为了简单起见，这里我直接使用了一个随机词 sks，并将其硬编码进去。如果你想更改它，只需修改这个文件即可。

训练将进行 800 步，会在 ./logs/<job_name>/checkpoints 目录下保存两个检查点，一个在 500 步时，另一个在最后一步。通常，500 步时的检查点就已经足够好用了。我使用两块 A6000 显卡进行训练，大约需要 15 分钟。

生成

训练完成后，可以通过运行以下命令来获得个性化样本：

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/from/training
                                 --prompt "photo of a sks <class>" 

特别地，sks 是标识符，如果你更改了标识符，就需要替换为你选择的词；而 <class> 则是训练时使用的类别词 --class_word。

结果

这里展示一些定性结果。训练图像取自 Textual Inversion 仓库中的 issue，分别是三个大型垃圾箱的图片。正则化图像由提示词 photo of a container 生成。正则化图像如下：

训练完成后，使用提示词 photo of a sks 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 生成的图像：

Dreambooth-Stable-Diffusion 快速上手指南

本指南基于 Google Dreambooth 论文思想，利用 Stable Diffusion 模型实现少量图片的微调训练。与 Textual Inversion 仅优化词嵌入不同，Dreambooth 会对整个扩散模型进行微调，从而获得更高质量的个性化生成效果。

环境准备

系统要求

• 操作系统: Linux (推荐) 或 macOS
• GPU: 建议使用显存 16GB 及以上的 NVIDIA GPU（如 A6000, A100, RTX 3090/4090）。原作者使用双卡 A6000 训练约需 15 分钟。
• Python: 3.8+

前置依赖

本项目基于 textual_inversion 和 stable-diffusion 官方代码库。

1. 克隆仓库并安装基础依赖：

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

   注：若国内下载依赖缓慢，可配置 pip 国内镜像源（如清华源）：

   pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple
   

2. 获取预训练模型权重： 你需要下载 Stable Diffusion v1 的预训练权重。推荐版本为 sd-v1-4-full-ema.ckpt。

   • 下载地址：HuggingFace - CompVis/stable-diffusion
   • 国内加速建议：若 HuggingFace 访问困难，可使用国内镜像站（如 ModelScope 或 HuggingFace 国内镜像）下载后放置到本地路径。

安装步骤与数据准备

在开始训练前，必须准备两类数据：训练图像（你的主体对象）和 正则化图像（防止过拟合）。

1. 生成正则化图像 (Regularization Images)

Dreambooth 算法需要一组描述同类物体的通用图像作为正则化项。 假设你要训练的对象类别是 dog（狗），请运行以下命令生成正则化图片（建议生成 100-200 张以获得更好效果，原示例为 8 张）：

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

请将 <class> 替换为你的类别词（如 dog, container, man 等），并将生成的图片保存至单独文件夹，例如 /root/to/regularization/images。

提示：如果生成的正则化图像质量不佳（例如训练“人”时），可以直接从网络下载多样化的同类公开图片作为正则化数据集。

2. 准备训练图像

将你的主体对象图片（通常 3-5 张即可）放入一个文件夹，例如 /root/to/training/images。

基本使用

启动训练

运行以下命令开始微调。请根据实际情况修改路径和参数。

python main.py --base configs/stable-diffusion/v1-finetune_unfrozen.yaml \
                -t \
                --actual_resume /path/to/original/stable-diffusion/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 <class>

参数说明：

• <job_name>: 自定义任务名称，日志和模型将保存在 ./logs/<job_name>/。
• <class>: 对象的类别词（需与生成正则化图像时使用的词一致）。
• --gpus: 指定使用的 GPU 编号。
• 标识符 (Identifier): 代码中硬编码了稀有词 sks 作为占位符。如需修改，请编辑 ldm/data/personalized.py 第 10 行。
• 训练步数: 默认训练 800 步，会在 500 步和结束时保存检查点。通常 500 步的效果已足够好。
• 学习率: 配置文件默认为 1.0e-6，以避免过拟合导致编辑性变差。

生成个性化图像

训练完成后，使用生成的检查点进行推理。

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/from/training \
                                 --prompt "photo of a sks <class>"

使用说明：

• --ckpt: 指向训练输出的 .ckpt 文件路径（通常在 ./logs/<job_name>/checkpoints/ 下）。
• sks: 训练时使用的唯一标识符，提示词中必须包含它来触发个性化特征。
• <class>: 替换为你的类别词。
• 你可以尝试组合新场景，例如："photo of a sks container on the moon"。