稳定扩散的梦想堡垒

这是谷歌的“稳定扩散”梦境的实现。最初的梦亭是基于图像文本到图像模型。但是，模型和 Imagen 的预训练权重均不可用。为了让人们能够用一些例子来微调文本到图像模型，我在稳定扩散上实现了Dreambooth的想法。

此代码存储库基于文本反转的代码存储库。请注意，文本反转仅优化单词衰减，而梦境则微调整个扩散模型。

该实现对文本反转的官方代码库进行了最小的更改。事实上，由于懒惰，文本反转中的某些组件（例如嵌入管理器）不会被删除，尽管它们永远不会在这里使用。

更新

9 / 20 / 2022：我刚刚找到了一种减少GPU内存的方法。请记住，此代码基于文本反转，TI 的代码库具有此行，它以硬编码方式禁用渐变检查点。这是因为在 TI 中，Unet 没有经过优化。但是，在 Dreambooth 中，我们优化了 Unet，因此我们可以打开渐变检查点指向技巧，就像此处的原始 SD 存储库一样。在配置中，渐变检查点默认为 True。我已经更新了代码。

用法

制备

首先按照文本反转存储库或原始稳定扩散存储库的指令设置环境。

ldm

要微调稳定扩散模型，你需要按照其指示获得预先训练的稳定扩散模型。权重可以在拥抱脸上下载。你可以决定使用哪个版本的检查点，但我使用 。

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，以更好地与原始纸张对齐。为了适应这一点，我缩短了配置文件中reg数据集的“重复”。

在某些情况下，如果生成的正则化图像非常不切实际（当你想要生成“男人”或“女人”时就会发生这种情况），你可以在线找到一组不同的图像（男人/女人），并将它们用作正则化图像。

训练

可以通过运行以下命令来完成训练

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>

详细配置可在 中找到。特别是，默认的学习速率是我在Dreambooth论文中发现的，导致可编辑性差。该参数对应于 Dreambooth 论文中正则化的权重，默认值设置为 。

configs/stable-diffusion/v1-finetune_unfrozen.yaml

1.0e-6

1.0e-5

reg_weight

1.0

梦境需要一个占位符词，称为标识符，如论文所示。此标识符必须是词汇表中相对罕见的标记。原始论文通过在T5-XXL分词器中使用一个罕见的单词来解决这个问题。为简单起见，在这里我只使用一个随机单词并对其进行硬编码。如果要更改它，只需在此文件中进行更改即可。

[V]

sks

训练将运行 800 个步骤，两个检查点将保存在 处，一个保存为 500 步，一个保存在最后一步。通常，500步的步骤效果很好。我使用两个A6000 GPU训练模型，大约需要15分钟。

./logs/<job_name>/checkpoints

代

训练后，可以通过运行命令获得个性化样本

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

结果

在这里，我展示了一些定性结果。训练图像是从文本反转存储库中的问题中获取的，它们是大型垃圾箱的 3 个图像。正则化图像由提示 生成。正则化图像如下所示：

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