近两年 AI 发展非常迅速,其中的 AI 绘画也越来越火爆,AI 绘画在很多应用领域有巨大的潜力,AI 甚至能模仿各种著名艺术家的风格进行绘画。
目前比较有名商业化的 AI 绘画软件有 Midjourney、DALL·E2、以及百度出品的文心一格:https://yige.baidu.com/creation
但是他们都有一个共同点,那就是要钱。为了解决这个问题,我们可以自己做一款 AI 绘图软件。
本次分享主要涉及的内容:
现在主流的两个图像生成核心模型是 GAN 和 Diffusion Models。
生成对抗网络(Generative Adversarial Nets,GAN)于 2014 年提出,是一种基于对抗学习的深度生成模型,它由两个主要组件组成:生成器和判别器。生成器通过学习输入数据的分布,生成新的数据样本;判别器则尝试区分生成器生成的数据和真实数据。通过不断迭代训练,生成器和判别器相互对抗,最终生成器能够生成越来越逼真的数据。
扩散模型(Diffusion Models,DM) 于 2015 年被提出,在提出后的好多年中并没有掀起什么波澜,直到 2020 到 2022 年期间,基于该模型提出了其他改良模型如 DDPM、DDIM 等,扩散模型开始引起大量关注,2022 年 8 月基于扩散模型设计的 Stable Diffusion 出现后,扩散模型直接爆火。
扩散模型,像分子运动一样,一点点改变。对于图像而言,就是图像上的像素点一点点改变,直到最后改变成了有意义的图像。
不管是 GAN 模型还是 DM 模型,他们本质上都是给定输出 y 和输入 x,然后通过神经网络和深度学习建立两者的模式。
对于 AI 绘画,我们一般需要给出一个提示,让 AI 返回与提示匹配的图像,对于数学来说,我们则需要找到一个函数,让它能根据我们输入,转化成我们想要的输出。这个函数背后象征着一种模式,函数则依据这个模式来将我们的输入转化为输出。
那么怎么找到这个模式呢?
这就需要引入神经网络和深度学习了。
所谓的神经网络,其实都是由许多神经单元构成,而简单的神经单元,用数学公式表示的话,最基础,最简单的就是这样一个线性函数公式:
y = a x + b y = ax + b y=ax+b
有了这个简单的神经单元我们可以拟合一些数据的表现,比如下面这个图。
就像这个图,我们有一组数据(蓝色的点),然后我们用红线对应的函数表达了这组数据,虽然红线上的值和这组数据的分布点有差距,但差距不大,所以我们可以认为这个函数表达了这组数据的模式(学会了这组数据的模式)。
扩散模型的训练过程需要遵循监督学习的模式,它分为两个过程,分别是前向过程和后向过程。前向过程可以认为是生成输出 y 的过程,后向过程则是根据输入 x 输出对应 y 的训练过程。
扩散模型的原论文链接:https://arxiv.org/pdf/2006.11239.pdf
前向过程,这个过程目的是生成一系列噪声,用于之后的后向过程训练,前向过程是不需要学习的。
我们先观察一下噪声分布,
大家觉得每个时刻加的噪音是一样的吗,一开始的时候加的多,还是后面加的多呢?
某个时刻的噪声跟哪个时刻最有关系呢?很明显当前时刻的噪声跟前一个时刻的噪声关系最密切,因为当前时刻的噪声可以用前一个时刻的噪声来求出,其实当前时刻的噪声和前一个时刻的噪声也只差了一个噪声而已。所以我们看第一个公式:
x t = a t x t − 1 + 1 − a t z 1 x_t = \sqrt{{a}_t}x_{t-1} + \sqrt{1-{{a}_t}}z_1 xt=atxt−1+1−atz1
a t = 1 − β t a_t = 1 - β_t at=1−βt
这个公式里面,β 是一个常量,它的值从 0.0001 到 0.002。
现在我们已经能够求出各个时刻需要加的噪声是多少了,但是还有一个问题,就是我们每次计算当前时刻噪声的时候,都需要从 T0 时刻开始一直计算到 T 时刻,这个过程对前向过程可能没问题,也许只需要浪费一点内存保存前一个过程的噪声就行了,
但是扩散模型不止有前向过程,还有一个后向过程,后向过程其实就是根据当前噪声倒推前一个噪声是什么,所以后向过程只关心当前时刻的噪声,其他时刻的噪声并不关心,所以根据论文给出的原始公式,我们可以推导出一个公式:
x t = a ‾ t x 0 + 1 − a ‾ t z x_t = \sqrt{\overline{a}_t}x_0 + \sqrt{1-{\overline{a}_t}}z xt=atx0+1−atz
这个公式成立的条件是 z 必须符合高斯分布,现在我们想要知道任意时刻 T 的噪声,只需要代入这条公式就可以了。
后向过程其实就是由当前时刻噪声预测出前一个时刻的噪声,然后用这个预测的噪声和我们前向过程中对应时刻的噪声做对比,最终得到一个差距不大的噪声,就算预测成功。
这个后向过程也是有公式的,而且论文中已经直接给出,不需要我们自己推导,如下:
μ t = 1 a t ( x t − β 1 − a ‾ t ϵ ) μ_t = \dfrac{1}{\sqrt{a_t}}(x_t - \dfrac{β}{\sqrt{1-\overline{a}_t}}ϵ) μt=at1(xt−1−atβϵ)
x t − 1 = μ t + σ t z x_{t-1} = μ_t + σ_tz xt−1=μt+σtz
观察一下公式,发现这里面除了 ϵ 参数,其他参数都是已知的,那么这个参数我们从哪里得到呢?
这就是从训练中得到的,在训练的过程中,神经网络需要不断调整 ϵ 的值,直到通过 ϵ 计算出来的预测噪声和前向过程对应时刻的噪声对比差距不大的时候,这个值就算定下来了,以后给定任何带噪声的图像,我们就可以用上面的公式,不断往前推到 T0 时刻,最终得到生成的图像。
这个时候又有人会问了,能不能直接通过当前时刻,一步到位直接计算出 T0 时刻的噪声呢,就像前向过程那样,能减少大量计算。
很遗憾,现在还做不到。
现在我们已经了解了 Diffusion Models 中的前向过程和后向过程了,接下来我们看看它是怎么训练的。
1、首先我们拿到要学习的图片
2、给这个图片添加一个噪声N,并把这个噪声保存下来
3、把加噪后的图片扔给神经网络,神经网络根据加噪后的图片输出预测的噪声PN
4、比较 PN 和 N 在数学尺度上的"差距" ,这个差距我们记为 D
5、把这个差距 D 扔给迭代器,迭代器会告诉神经网络应该怎么调整神经参数来缩小 N 和 PN 的差距
6、最后重复不断这个过程,直到 D 的值足够小,我们就认为神经网络学会我们期望的运动方式了
训练完毕后,我们就可以使用这个训练好的神经网络了来生成图像了,再来看看神经网络是如何一步步去掉噪声,生成最终图像的。
下面我们来演示 Demo,这个 Demo 展示了一个比较完整的 Diffusion Models 的训练和生成图片的过程,我们重点看里面的:
首先我们需要定义 β 、a 和 T 这些参数, β 随着时间不断线性变大,初始值是 0.0001,结束值是 0.002,T 为 200,也就是前向过程一共有 200 步。
# T 步骤长度
timesteps = 200
# 使用线性函数来获取 β 值序列, 初始值是 0.002,结束值是 0.0001
betas = linear_beta_schedule(timesteps=timesteps)
# 定义 a = 1 - β
alphas = 1. - betas
# 所有 a 累乘
alphas_cumprod = torch.cumprod(alphas, axis=0)
# 对 1/a 进行开方
sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
# 对 a 累乘进行开方
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
然后我们看前向过程的函数,这个函数就是完全根据公式给出的来写的,通过这个函数可以获得任意 T 时刻的噪声强度。
def q_sample(x_start, t, noise=None):
if noise is None:
noise = torch.randn_like(x_start)
sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
sqrt_one_minus_alphas_cumprod_t = extract(
sqrt_one_minus_alphas_cumprod, t, x_start.shape
)
# 前向过程计算 T 时刻公式
return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
然后我们看后向过程的函数,这个函数就是完全根据公式给出的来写的,通过这个函数可以当前时刻推理出前一时刻的噪声强度。
def p_sample(model, x, t, t_index):
betas_t = extract(betas, t, x.shape)
sqrt_one_minus_alphas_cumprod_t = extract(
sqrt_one_minus_alphas_cumprod, t, x.shape
)
sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
model_mean = sqrt_recip_alphas_t * (
x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t
)
if t_index == 0:
return model_mean
else:
posterior_variance_t = extract(posterior_variance, t, x.shape)
noise = torch.randn_like(x)
# 后向过程计算 T-1 时刻公式
return model_mean + torch.sqrt(posterior_variance_t) * noise
再看训练过程,训练过程首先是加载数据集,然后开始对数据集进行训练,
训练的过程中会随机取一个时间 T,根据这个时间 T 计算出当前时刻的噪声分布,然后根据模型计算出预测的噪声分布,两个根据这两个噪声分布计算损耗,并将损耗给到迭代器进行对参数进行优化
device = "cuda" if torch.cuda.is_available() else "cpu"
# 使用 Unet 模型初始化模型
model = Unet(
dim=image_size,
channels=channels,
dim_mults=(1, 2, 4,)
)
model.to(device)
# 定义模型优化器,用于优化和调整模型
optimizer = Adam(model.parameters(), lr=1e-3)
model_path = "model_params/model_params.pth"
# 判断文件是否存在
if os.path.exists(model_path):
print("使用本地模型")
model.load_state_dict(torch.load('model_params/model_params.pth'))
else:
print("开始训练模型")
epochs = 5
for epoch in range(epochs):
for step, batch in enumerate(dataloader):
optimizer.zero_grad()
batch_size = batch["pixel_values"].shape[0]
batch = batch["pixel_values"].to(device)
t = torch.randint(0, timesteps, (batch_size,), device=device).long()
loss = p_losses(model, batch, t, loss_type="huber")
# 计算损失的梯度
loss.backward()
# 根据优化器来更新模型的权重
optimizer.step()
# 保存模型参数
torch.save(model.state_dict(), 'model_params/model_params.pth')
再看生成过程,首先我们需要准备一张还有噪声分布的图片,然后通过后向过程公式,不断对这个图片进行去噪,最终生成一个 T 等于 0 时刻的图像。
def p_sample_loop(model, shape):
device = next(model.parameters()).device
# 生成随机噪音图
img = torch.randn(shape, device=device)
# imgs 用于存储采样结果
imgs = []
# 逆序循环遍历所有时刻,对于每个时刻,使用后向过程公式进行去噪,并将去噪后的图片添加到 imgs 列表中
for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps):
img = p_sample(model, img, torch.full((shape[0],), i, device=device, dtype=torch.long), i)
imgs.append(img.cpu().numpy())
return imgs
上面这个 demo 中,我们已经学会了如何训练模型和使用模型了,但是局限性还是很大,只能训练分辨率和质量非常低的模型,生成的图片质量也很差。
想要训练一个能够生成高分辨率,高质量图片的大模型,可能需要使用上百个 NVIDIA A100 GPU 训练上万个 GPU 时,这个训练过程是相当耗时。
为了解决上面这个问题,我们就需要使用别人训练好的模型,这些模型大小,大概在 1-10G 不等。
另外还有一个关键问题:我们现在是输入图片,输出图片(图生图),但我们希望的是输入文字,输出图片(文生图)。
实际上现在所有的文生图本质上都是图生图,其实就是将不同文本转化为不同噪声分布的图像,并依据模型一步步去噪,最终得到目标图像。
将不同文本转化为不同噪声分布的图像,这里又需要进入另一个模型,即 CLIP 模型,CLIP 全称 Contrastive Language-Image Pre-Training(对比性语言-图像预训练模型),是 OpenAI 在 2021 年初开源的一个用文本作为监督信号来做预训练的模型,这个模型中有一个 Text Encoder,也就是文本编码器。现在主流的文生图工具都会用到这个编码器来协助生成图片。具体原理我们就不探讨了,知道有这个东西就可以。
这是基于扩散模型设计的一款开源的 AI 绘画软件,提供文生图、图生图、图片修复等多种功能。
Stable Diffusion 可以直接部署到本地,部署完毕后我们就可以使用它提供的 webui 来进行 AI 绘画功能,前提条件是我们的电脑 GPU 配置足够高,毕竟生成图片是要消耗大量 GPU 资源的,如果电脑配置一般的话,可能生成的图片质量会很差,同时分辨率很低。
现在市场上大部分的 AI 绘图软件基本都是基于 Stable Diffusion 来部署的。
GitHub地址:https://github.com/AUTOMATIC1111/stable-diffusion-webui
基于 ubuntu 系统部署步骤可以参考我写的另一片文章:Ubuntu 20.04 安装 Stable Diffusionn
其他系统部署步骤都大差不差。
当我们部署完毕后,在浏览器中输入:http://127.0.0.1:7860/,就可以访问我们本地的 Stable Diffusion WebUI 了,界面如下:
它里面提供了非常多的功能,也支持很多自定义的扩展功能,基本能解决我们对 AI 绘画的要求。
它的局限性也很明显:
现在我的目标就是将这些痛点解决掉,所以我们可能需要做:
当我们解决了这些东西,我们就能做出一款自己随时可以使用的 AI 绘图软件了。
对于服务器接口,我们可以基于 Stable Diffusion 提供的 API 进行进一步的封装,这一点非常关键,Stable Diffusion 为我们提供了丰富的 api 接口。
只需要在启动 Stable Diffusion 的时候传入 --api 参数,我们就可以在浏览器上输入:http://127.0.0.1:7860/docs 访问到所有 api 信息:
文生图接口 /sdapi/v1/txt2img 的入参如下:
{
"enable_hr": false,
"denoising_strength": 0,
"firstphase_width": 0,
"firstphase_height": 0,
"hr_scale": 2,
"hr_upscaler": "string",
"hr_second_pass_steps": 0,
"hr_resize_x": 0,
"hr_resize_y": 0,
"hr_sampler_name": "string",
"hr_prompt": "",
"hr_negative_prompt": "",
"prompt": "",
"styles": [
"string"
],
"seed": -1,
"subseed": -1,
"subseed_strength": 0,
"seed_resize_from_h": -1,
"seed_resize_from_w": -1,
"sampler_name": "string",
"batch_size": 1,
"n_iter": 1,
"steps": 50,
"cfg_scale": 7,
"width": 512,
"height": 512,
"restore_faces": false,
"tiling": false,
"do_not_save_samples": false,
"do_not_save_grid": false,
"negative_prompt": "string",
"eta": 0,
"s_min_uncond": 0,
"s_churn": 0,
"s_tmax": 0,
"s_tmin": 0,
"s_noise": 1,
"override_settings": {},
"override_settings_restore_afterwards": true,
"script_args": [],
"sampler_index": "Euler",
"script_name": "string",
"send_images": true,
"save_images": false,
"alwayson_scripts": {}
}
可以看到,区区一个文生图的接口的入参就非常的多,由此可以看出,官方的接口使用成本是很高的。
这些接口其他缺点也很明显:
除了希望解决以上问题,我们可能还希望:
为了解决以上问题,我们必须对官方提供的 api 进行进一步封装,具体步骤不演示,直接写业务代码就行了。
下面给大家提供一个我们已封装好的简单的文生图接口:
http://aiycx.cn/wallpaper/api/txt2img
请求参数
参数名称 | 是否必须 | 示例 | 备注 |
---|---|---|---|
prompt | 是 | 1fish | |
negativePrompt | 是 | lowres, bad anatomy | |
width | 是 | 512 | |
height | 是 | 512 |
返回数据
名称 | 类型 | 示例 | 备注 |
---|---|---|---|
code | string | 1001 | 1001-成功; 1002-失败 |
msg | string | success | |
data | object |
返回示例
{
"code":"1001",
"data":{
"images":[
"http://aiycx.cn/WallpaperImages/725134e5d2104ac25403269effa7a64b.png"
]
},
"msg":"success"
}
相比于官方一大推让人摸不着头脑的参数,这个接口显得非常简单。
其实有了这个接口,我们已经可以作出一个简单的可演示的 AI 绘画 demo 了。
但我们还想做得更完备,希望可以媲美市面上主流的 AI 软件,我们还需要做的其他工作:
Stable Diffusion 是基于英文标签来进行训练的,所以它的提示词只支持英文,如果我们想要提示词支持中文,首先很容易想到的方案就是在输入前先对提示词进行翻译,也就是中文转英文,实际上现在常规的做法也是如此。
为了实现这个功能,我们需要引入一款翻译插件,这款翻译插件可以将英文转化成多种语言,一款非常好用的翻译插件推荐给大家:
GitHub 地址如下:https://github.com/studyzy/sd-prompt-translator
为了实现这个功能,我们首先需要准备一份敏感词黑名单,敏感词往往有很多,可能多达 1.5W 个,也就是说每次接口请求的时候,我们都需要遍历一遍名单中的上万个敏感词做出匹配,每当匹配到敏感词,则不允许进行请求。所以,我们需要一款高性能的敏感词识别工具,GitHub 上就有(据说是行业最快的一款敏感词识别工具),地址如下:https://github.com/toolgood/ToolGood.Words
往往做了提示词识别敏感词,已经能防止 90% 以上的情况生成不良图像了。
但如果你想更严谨一点,则可以再做一层敏感图的过滤。
往往我们软件会涉及到多个用户,每个用户都有自己的创作内容,也可以自己管理自己的创作内容,所以就需要一个开发账号系统。
这里涉及到的业务有用户登录(一键登录、验证码登录、密码登录等)、用户数据库设计,数据表增删查改。
可以简单看下我的表设计结构:
CREATE TABLE IF NOT EXISTS `ugc_table`
(
`id` BIGINT NOT NULL AUTO_INCREMENT COMMENT '主键',
`user_mobile` VARCHAR(191) NOT NULL COMMENT '创作者手机号',
`url` VARCHAR(255) NOT NULL DEFAULT '' COMMENT '图片链接',
`width` INT NOT NULL DEFAULT 0 COMMENT '图片宽',
`height` INT NOT NULL DEFAULT 0 COMMENT '图片高',
`prompt` TEXT NULL DEFAULT NULL COMMENT '提示词',
`negative_prompt` TEXT NULL DEFAULT NULL COMMENT '反向提示词',
`theme` VARCHAR(255) NULL DEFAULT NULL COMMENT '模型主题',
`style` TEXT NULL DEFAULT NULL COMMENT '风格选择',
`sampling_mode` VARCHAR(255) NULL DEFAULT NULL COMMENT '采样模式',
`seed` BIGINT NULL DEFAULT NULL COMMENT '随机种子',
`prompt_relevance` VARCHAR(255) NULL DEFAULT NULL COMMENT '提示词相关性',
`clip_skip` VARCHAR(255) NULL DEFAULT NULL COMMENT 'clip skip',
`tag` VARCHAR(255) NULL DEFAULT NULL COMMENT '标签',
`start_time` VARCHAR(255) NOT NULL DEFAULT '' COMMENT '开始创作时间',
`complete_time` VARCHAR(255) NOT NULL DEFAULT '' COMMENT '完成创建时间',
`md5` TEXT NULL DEFAULT NULL COMMENT '图片的md5',
`download_count` BIGINT NULL DEFAULT NULL COMMENT '下载次数',
`generate_type` VARCHAR(255) NULL DEFAULT NULL COMMENT '生成类型(txt2img:文生图/img2img:图生图)',
`img_status` INT NOT NULL DEFAULT 0 COMMENT '图片状态(0:未完成/1:已完成/2:错误)',
`task_id` VARCHAR(255) NOT NULL DEFAULT '' COMMENT '创建任务id',
`refer_img` VARCHAR(255) NULL DEFAULT NULL COMMENT '参考图',
PRIMARY KEY (`id`)
)
ENGINE = InnoDB
AUTO_INCREMENT = 0
CHARACTER SET = utf8mb4 COMMENT ='用户创作表';
其他都是一些业务代码,略。
在 AI 绘画中,主题和风格可以表现为
在哪可以下载这些模型呢?
推荐两个地方:https://civitai.com/、https://huggingface.co/models
常规情况下,SD 支持的分辨范围是 64-2048 分辨率,分辨率越高,需要的显存就越高,生成图片的时间就越慢。
据个人使用体验,在 NVIDIA 3070 8G 显卡上,
想要提高生图效率,一方面需要提升显卡性能,另一方面需要降低图片分辨率,但是图片分辨率太低(低于 512 )又可能导致图片质量不行。
这时有三条方案可以解决:
这里只介绍第一个,因为速度最快,从 512 * 512 分辨率 —> 2048 * 2048 只需要 3 秒就能完成。
具体做法是,在生成图片后,将图片的 base64 编码等参数传给 /sdapi/v1/extra-single-image 接口,该接口的入参如下:
{
"resize_mode": 0,
"show_extras_results": true,
"gfpgan_visibility": 0,
"codeformer_visibility": 0,
"codeformer_weight": 0,
"upscaling_resize": 2,
"upscaling_resize_w": 512,
"upscaling_resize_h": 512,
"upscaling_crop": true,
"upscaler_1": "None",
"upscaler_2": "None",
"extras_upscaler_2_visibility": 0,
"upscale_first": false,
"image": "base64"
}
如果我们只有一台 GPU 服务器,而有多个用户同时请求绘图接口,某些用户可能需要等待大量的时间,因为一个服务器同一时间段只能进行一次绘图。
想要解决并发访问问题,除了需要做好线程同步问题,还需要使用至少两台服务器来支持。
其中一台服务器作为主服务器,不需要太高的 GPU 性能,只负责能分发请求和存储用户数据;其他服务器就是 GPU 服务器,拥有较高的显卡性能,负责生成图片并返回给主服务器。
做完以上这些,我们的接口就相对来说比较完善了,后续的工作只需将接口应用到我们前端、客户端上,我们的软件就是一个完整的 AI 绘图软件了。