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

新闻 1970-01-01
435

目录

简介

近日,智源研究院发布了中英双语图像生成模型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函数在做什么就可以了。首先看到这个目标函数一共有两项,第一项为:

第二项为:

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

是不是基本长得一模一样,实际上,这个目标函数就是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

数据准备

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

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

然后将图片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,例如"一张男孩的照片"。

测试

训练时间根据使用的硬件不同而不同,也取决于你所准备的数据集数目。但大约都在一个小时内可以完成训练,只要你使用的显卡是可以把模型跑起来的。

测试的代码如下:

import torch
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

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

Prompt: 一张<鸣人>男孩的照片,背景是富士山,masterpieces

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

Prompt: 一张<鸣人>男孩的照片,背景是海洋,masterpieces

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

Prompt: 一张<鸣人>男孩的照片,背景是草原,masterpieces

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

也可以转换不同风格的鸣人:

铅笔素描风格:

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

油画梵高风格:

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

齐天大圣

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

加藤惠

Prompt: 一张<sks>女孩的照片

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

总结

Ok,that is all!感谢看到这里的小伙伴,我们可以看到AIGC令人惊叹的能力,智源研究院也将持续优化AltDiffusion,并且欢迎关注AltDiffusion背后的核心技术AltCLIP。

有任何问题都欢迎联系我们~ <open.platform@baai.ac.cn>

如果有兴趣加入我们,欢迎联系:jjli@baai.ac.cn

相关文章