简介
近日，智源研究院发布了中英双语图像生成模型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函数：
是不是看起来有点复杂，没事，我们只需要知道这个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库：
cd 到dreambooth代码所在的地址：
安装所需要的环境：
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:
其中会让你选择一些配置，根据自己的硬件条件进行选择即可；
当不使用prior loss进行训练时，用如下训练命令：
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 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，例如”一张男孩的照片”。
测试
训练时间根据使用的硬件不同而不同，也取决于你所准备的数据集数目。但大约都在一个小时内可以完成训练，只要你使用的显卡是可以把模型跑起来的。
测试的代码如下：
from diffusers import AltDiffusionPipeline
output_dir = “path-to-save-model”
pipe = AltDiffusionPipeline.from_pretrained(
output_dir,
torch_dtype=torch.float16
).to(“cuda:0”)
prompt = 一张鸣人男孩的照片,背景是沙漠,masterpieces
image = pipe(prompt, guidance_scale=7.5, height=512, width=512,num_inference_steps=50).images[0]
image.save(“alt_diffusion.png”)
训练效果
好了，经过一番努力过后我们拿到了我们想要的个性化模型，下面让我们来欣赏一下模型带给我们的视觉震撼吧！
鸣人
这里我们使用如上展示的鸣人数据集进行微调，可以为鸣人换不同的背景：
Prompt: 一张<鸣人>男孩的照片，背景是沙漠，masterpieces
Prompt: 一张<鸣人>男孩的照片，背景是富士山，masterpieces
Prompt: 一张<鸣人>男孩的照片，背景是海洋，masterpieces
Prompt: 一张<鸣人>男孩的照片，背景是草原，masterpieces
也可以转换不同风格的鸣人：
铅笔素描风格：
油画梵高风格：
齐天大圣
加藤惠
Prompt: 一张<sks>女孩的照片
总结
Ok,that is all!感谢看到这里的小伙伴，我们可以看到AIGC令人惊叹的能力，智源研究院也将持续优化AltDiffusion，并且欢迎关注AltDiffusion背后的核心技术AltCLIP。
有任何问题都欢迎联系我们～ <open.platform@baai.ac.cn>
如果有兴趣加入我们，欢迎联系：jjli@baai.ac.cn
相关文章
