想要拥有一个属于自己的个性化多语言stable diffusion吗？试试在Altdiffusion上使用Dreambooth!

新闻 1970-01-01

435 0

简介

近日，智源研究院发布了中英双语图像生成模型AltDiffusion，并在不到一周的时间内推出升级版AltDiffusion-m9，全面支持中、英、西、法、日、韩、阿、俄、意等九种不同语言的文生图任务。

目前，AltDiffusion已经集成到Huggingface diffusers中，可以用diffusers库方便使用AltDiffusion。本文将介绍如何利用diffusers库，在AltDiffusion使用Dreambooth技术微调出一个属于你的中文（或者9种语言中的任一一种）Stable diffusion。

AltDiffusion-m9模型地址：https://huggingface.co/BAAI/AltDiffusion-m9

AltDiffusion模型地址：https://github.com/FlagAI-Open/FlagAI/tree/master/examples/AltDiffusion

Huggingface diffusers 地址：https://huggingface.co/docs/diffusers/index

Diffusers github 地址：https://github.com/huggingface/diffusers

背景

Dreambooth

论文地址：https://arxiv.org/abs/2208.12242

Dreambooth是一个可以满足不同用户特定的需要的一种文生图微调方式，不同用户说的就是正在看的你~只需要自己准备少数图片，论文中说3~5张就可以，但实际上越多越好，在我们的实践过程中，收集20~30张图片可以产生非常满意的效果。

那么Dreambooth到底是怎么做的呢？Dreambooth是谷歌的工作，所以在原论文中介绍的是在谷歌的文生图模型Imagen做的，但是这不要紧，我们清楚他们的原理之后，在stable diffusion上一样是适用的。话不多说，我们来看一下训练的原理，直接上图：

数据集的准备

假如现在的输入图像如上图做左上角所示，我们的目的是想要让模型学会画这种特殊类别的狗，那么我们就将这种狗给定一个特殊的标注为："A [V] dog"，这样新准备的图像与这个特殊的标注"A [V] dog"，就组成了一个新的小样本训练集。

除了准备一个小样本的数据集以外，还需要用使用我们的文生图模型生成一组"先验"的图像。比如现在要想要生成的是一种特殊类别的狗，它是属于狗这个超类的。所以使用"A dog"这个prompt输入给我们想要微调的文生图模型，论文中是Imagen，也可以使用Stable diffusion，或者是我们的AltDiffusion-m9生成一组先验图像，大概生成200~300张图像即可，生成的图像和"A dog"这个prompt组成一个先验的数据集。

为什么要使用先验数据集呢，是因为我们想要现在想要把新的类别的狗的知识注入到模型中，但是要保留模型原有的对狗这个超类的知识，来防止过拟合和语言飘移的现象。

训练的目标函数

Math warning!先来看一下训练的loss函数：

想要拥有一个属于自己的个性化多语言stable diffusion吗？试试在Altdiffusion上使用Dreambooth!

是不是看起来有点复杂，没事，我们只需要知道这个loss函数在做什么就可以了。首先看到这个目标函数一共有两项，第一项为：

第二项为：

是不是基本长得一模一样，实际上，这个目标函数就是diffusion model的目标函数（这里只是一个形象的写法，跟实际训练有一些出入，具体关于diffusion模型的基础知识可以参考这篇blog）。两项的唯一区别就是变成了，代表在数据集准备阶段的新的小样本数据集，代表模型生成的先验数据。也就是我们使用的训练方式跟原来训练stable diffusion的方式是一样的，只不过现在我们数据集的每一个batch当中有一半是我们新准备的数据集，一半是模型生成的先验数据集，将这两个loss相加起来训练就是所谓的大名鼎鼎的Dreambooth了。

对应到上图中，就是上半部分是使用准备的新数据集算的Reconstruction Loss,下半部分是使用先验数据集生成的Class-Specific Prior Preservation Loss。训练的使用使用一个参数将两个loss结合起来就行了。

注意事项

上面说到的[V]这个特殊token，要尽量跟词表中正常的词区分开来。例如"sks"等，或者一些罕见的单词，这样是为了跟模型原有的语言理解区分开来，避免产生干扰。

使用diffusers库来实现Dreambooth

这里以AltDiffusion-m9为例子，来演示如何使用diffusers库实现dreambooth.

环境准备

首先Git clone diffusers库：

git clone https://github.com/huggingface/diffusers.git

cd 到dreambooth代码所在的地址：

cd /diffusers/examples/dreambooth

安装所需要的环境：

pip install -U -r requirements.txt

diffusers库中有支持原版stable diffusion的训练脚本，更改为AltDiffusion-m9，需要做一些改动，在/diffusers/examples/dreambooth路径下新建一个训练脚本命名为train_dreambooth_alt.py,代码内容如具体可以见我们提交的pr：https://github.com/huggingface/diffusers/pull/1390

数据准备

准备你想要训练的一个特殊类别的图像，放入到一个文件夹中，例如鸣人数据集：

然后将图片resize到512*512大小，可以使用这个站:

https://www.onlinephotosoft.com/birme/

训练

首先执行以下命令初始化accelerate:

accelerate config

其中会让你选择一些配置，根据自己的硬件条件进行选择即可；

当不使用prior loss进行训练时，用如下训练命令：

export MODEL_NAME="BAAI/AltDiffusion-m9"
export INSTANCE_DIR="path-to-instance-images"
export OUTPUT_DIR="path-to-save-model"

accelerate launch train_dreambooth_alt.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--output_dir=$OUTPUT_DIR \
--instance_prompt="一张<鸣人>男孩的照片" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--max_train_steps=1600

其中INSTANCE_DIR是新准备的数据路径地址，需要指定instance_prompt，例如"一张<鸣人>男孩的照片"，训练结束后会将所有的模型权重保存在"path-to-save-model"。

以上训练方式是最原始的stable diffusion训练方式，那么如果要使用上文提到的dreambooth训练，则使用如下训练命令：

export MODEL_NAME="BAAI/AltDiffusion-m9"
export INSTANCE_DIR="path-to-instance-images"
export CLASS_DIR="path-to-class-images"
export OUTPUT_DIR="path-to-save-model"

accelerate launch train_dreambooth_alt.py \
--pretrained_model_name_or_path=$MODEL_NAME \
--instance_data_dir=$INSTANCE_DIR \
--class_data_dir=$CLASS_DIR \
--output_dir=$OUTPUT_DIR \
--with_prior_preservation --prior_loss_weight=1.0 \
--instance_prompt="一张<鸣人>男孩的照片"\
--class_prompt="一张男孩的照片" \
--resolution=512 \
--train_batch_size=1 \
--gradient_accumulation_steps=1 \
--learning_rate=5e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--num_class_images=200 \
--max_train_steps=1600

其中比原来的训练命令多出来的部分是，要指定一个CLASS_DIR，这个路径是模型生成的先验图片的路径。指定一个保存的文件夹位置即可，训练脚本会自动帮你生成你所想要的先验图片，指定生成多少先验图像可以使用num_class_images进行指定。除此之外还需要额外指定class_prompt，例如"一张男孩的照片"。