一、引言
之前读了有关Dreambooth的论文(具体可以看我另外一篇博客Dreambooth),针对实现特定主题的图像生成,我还是挺感兴趣的哈哈,虽然这篇论文的作者并没有给出代码(主要是由于他是基于Imagen实现的),但是我发现GitHub上面还是有大佬们实现了stable-diffusion的版本,于是我尝试下载下来跑一跑学习一下哈哈,在调试的过程中,我遇到了很多问题,说实话对于我这种科研废物来说调试也是个小难题,这篇博客主要是记录自己整个部署的过程以及遇到的一些问题。
这篇文章主要是讲述我自己在本地环境运行的一个过程,如果你的GPU不是3090这样的有着24GB显存的显卡的话,那是训练不了的,当然如果你很想尝试,可以去用一些网上的共享GPU资源,具体的就不多介绍了,而且这个项目暂时好像还没有实现多GPU训练。
二、下载所需资源
Tips:这里我们还是需要提前安装好Git、python环境、以及pytorch的配置当然还有诸如Vscode、Pycharm等工具。
1.下载项目
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里都可以看到。
Command | Type | Example | Description |
---|---|---|---|
--config_file_path | string | "C:\\Users\\David\\Dreambooth Configs\\my-config.json" | 要使用的配置文件的路径 |
--project_name | string | "My Project Name" | 项目名称 |
--debug | bool | False | 启用调试日志记录False |
--max_training_steps | int | 3000 | 要运行的训练步骤数 |
--token | string | "ddd" | 要表示已训练模型的唯一标记。 |
--training_model | string | sd-v1-4.ckpt | 模型到训练的路径 (model.ckpt) |
--training_images | string | images/train/trian_dog | 训练图像目录的路径 |
--regularization_images | string | images/regularization/regularization_dog | 正则化图像的目录路径 |
--class_word | string | "dog" | 要训练的图像所属的类别 |
--flip_p | float | 0.0 | 翻转百分比 |
--learning_rate | float | 1.0e-06 | 自选默认值为 (0.000001)。设置学习率。 可以用科学计数法如: |
--save_every_x_steps | int | 250 | 设置每多少step保存一次模型权重 |
--gpu | int | 0 | 指定训练所用的GPU,默认为0, 暂时还没有实现多GPU训练 |