Dreambooth——代码实现_dreambooth代码复现

一、引言

        之前读了有关Dreambooth的论文（具体可以看我另外一篇博客Dreambooth），针对实现特定主题的图像生成，我还是挺感兴趣的哈哈，虽然这篇论文的作者并没有给出代码(主要是由于他是基于Imagen实现的)，但是我发现GitHub上面还是有大佬们实现了stable-diffusion的版本，于是我尝试下载下来跑一跑学习一下哈哈，在调试的过程中，我遇到了很多问题，说实话对于我这种科研废物来说调试也是个小难题，这篇博客主要是记录自己整个部署的过程以及遇到的一些问题。

        这篇文章主要是讲述我自己在本地环境运行的一个过程，如果你的GPU不是3090这样的有着24GB显存的显卡的话，那是训练不了的，当然如果你很想尝试，可以去用一些网上的共享GPU资源，具体的就不多介绍了，而且这个项目暂时好像还没有实现多GPU训练。

二、下载所需资源

Tips:这里我们还是需要提前安装好Git、python环境、以及pytorch的配置当然还有诸如Vscode、Pycharm等工具。

1.下载项目

我们先从GitHub上下载整个项目：JoePenna/Dreambooth-Stable-Diffusion: Implementation of Dreambooth (https://arxiv.org/abs/2208.12242) by way of Textual Inversion (https://arxiv.org/abs/2208.01618) for Stable Diffusion (https://arxiv.org/abs/2112.10752). Tweaks focused on training faces, objects, and styles. (github.com)

2.安装依赖包

# 创建虚拟环境(根据自己的版本来选择最后的python版本)
conda create -n Dreambooth python=3.10
# 安装好后激活虚拟环境，并且cd进入项目所在文件夹中
conda activate Dreambooth
# 安装所需依赖包
pip install -r requirements.txt

结合我所遇到的问题，这里有几个需要注意的点：

1. 如果你之前有安装过Stable-diffusion，有对应的虚拟环境，我建议是重新创建一个新的虚拟环境，如果用同一个虚拟环境，后面你在运行的时候会造成访问冲突等一些不知名错误。

2. 可能在安装依赖的时候，那个clip和taming-transformers还是下载不下来，那我们就直接访问GitHub上的链接自己下载下来安装，然后自己建一个src文件夹放进去。

3.安装openai/clip-vit-large-patch14

我在上面两部做好之后，我就开始运行了，然后它提示我clip-vit-large-patch14报错，具体错误好像就是网络连接不上然后下载不下来。

同样，我们用最简单暴力的方法，去官网给他下载下来openai/clip-vit-large-patch14 at main (huggingface.co)

下载好后我们自己创建文件夹openai/clip-vit-large-patch14然后把文件都放进去（当然你也可以不叫这个文件名，只需要在ldm/modules/encoders/modules.py中把文件名改掉就行。）

4.下载模型权重 

Dreambooth实际上就是对原有的stable-diffusion进行微调，所以我们需要原始的stable-diffusion模型权重，大家可以自行去下载，我这里使用的是sd-v1-4.ckpt。

CompVis/stable-diffusion-v-1-4-original at main (huggingface.co)

 三、训练

1. 准备素材

Dreambooth主要要用到两个图片素材，一个是用于正则化的素材，主要是为了防止模型在训练过程中产生“语言漂移的现象”（详细内容可以阅读一下论文）。另一个则是要对模型进行微调的输入图片。注：图片的尺寸需要一致，我这里是将它们都裁剪成了512*512的大小。

一个一个裁剪比较麻烦，这里推荐一个批量裁剪的网址：在线图片固定比例批量裁剪工具 - UU在线工具 (uutool.cn)

准备好图片后，我们把他们放入项目中： 

2. 开始训练

现在，我们就可以开始训练了,下面给出了我第一次进行微调的命令范例（运行main.py脚本）：

python "main.py" --project_name "Dreambooth_dog" --training_model "sd-v1-4.ckpt" --regularization_images "images/regularization/regularization_dog" --training_images "images/samples/samples_dog" --max_training_steps 2000 --class_word "dog" --token "ddd" --flip_p 0 --learning_rate 1.0e-06 --save_every_x_steps 500

如果一切顺利的话，就会出现如下图的状况，那么模型就开始进行训练啦（我用实验室的3090训练了2000steps大概用了一个多小时吧，可以玩几局王者等一会哈哈哈）

 四、运行

训练完成后，你就可以用微调后的模型进行个性化图像生成啦，它的模型默认保存的地址是trained_model文件夹。

有了模型之后，我们运行脚本来实现个性化图像生成（stable_txt2img.py）：

python scripts/stable_txt2img.py --ddim_eta 0.0 --n_samples 8 --n_iter 1 --scale 10.0 --ddim_steps 50  --ckpt "trained_models/2023-12-08T14-33-04_Dreambooth_dog_02000_steps_5_training_images_ddd_token_dog_class_word.ckpt" --prompt "a ddd dog in Tiananmen Square, Beijing"

这是我输入的五张图片：

我的输出（a ddd dog in Tiananmen Square, Beijing，ddd是我给这只狗设的token）：

Tips:在运行的过程中，由于整个脚本的实现默认是在GPU0上运行的，但是我运行的时候，实验室GPU0的显卡被占用了，于是我用了另外一张显卡，如果你也遇到这样的问题可以在脚本文件的开头加上下面的代码：

torch.cuda.set_device(5) #指定你用的GPU卡号，比如我这里用的就是5号卡

五、部分参数解释

这个项目的作者在参数方面做了很好的解释，为了让整篇文章完整一点这里就搬运过来简单写写了哈哈哈，当然所有的训练脚本所用的参数在项目文件dreambooth_helpers/arguments.py里都可以看到。