Dreambooth-Stable-Diffusion - 这是 Google 的Dreambooth的一个实现,具有Stable Diffusion。

Created at: 2022-09-06 14:57:40
开发语言: Jupyter Notebook
授权协议: MIT

稳定扩散的梦想堡垒

这是谷歌的“稳定扩散”梦境的实现。最初的梦亭是基于图像文本到图像模型。但是,模型和 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