前言
近两年 AI 发展非常迅速,其中的 AI 绘画也越来越火爆,AI 绘画在很多应用领域有巨大的潜力,AI 甚至能模仿各种著名艺术家的风格进行绘画。
目前比较有名商业化的 AI 绘画软件有 Midjourney、DALL·E2、以及百度出品的文心一格:https://yige.baidu.com/creation
但是他们都有一个共同点,那就是要钱。为了解决这个问题,我们可以自己做一款 AI 绘图软件。
本次分享主要涉及的内容:
- 扩散模型(Diffusion Models)的原理
- 扩散模型(Diffusion Models)的实践
- Stable Diffusion 简单使用
- Stable Diffusion 远端部署
- 基于 Stable Diffusion 做一款属于自己的高质量 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,这个 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_meanelse: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 = 5for 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
Stable Diffusion
上面这个 demo 中,我们已经学会了如何训练模型和使用模型了,但是局限性还是很大,只能训练分辨率和质量非常低的模型,生成的图片质量也很差。
想要训练一个能够生成高分辨率,高质量图片的大模型,可能需要使用上百个 NVIDIA A100 GPU 训练上万个 GPU 时,这个训练过程是相当耗时。
为了解决上面这个问题,我们就需要使用别人训练好的模型,这些模型大小,大概在 1-10G 不等。
另外还有一个关键问题:我们现在是输入图片,输出图片(图生图),但我们希望的是输入文字,输出图片(文生图)。
实际上现在所有的文生图本质上都是图生图,其实就是将不同文本转化为不同噪声分布的图像,并依据模型一步步去噪,最终得到目标图像。
将不同文本转化为不同噪声分布的图像,这里又需要进入另一个模型,即 CLIP 模型,CLIP 全称 Contrastive Language-Image Pre-Training(对比性语言-图像预训练模型),是 OpenAI 在 2021 年初开源的一个用文本作为监督信号来做预训练的模型,这个模型中有一个 Text Encoder,也就是文本编码器。现在主流的文生图工具都会用到这个编码器来协助生成图片。具体原理我们就不探讨了,知道有这个东西就可以。
Stable Diffusion 的部署
这是基于扩散模型设计的一款开源的 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 绘画的要求。
它的局限性也很明显:
- 对电脑显卡要求很高,自己电脑带不动
- 只能本地访问
- 灵活性低,只能通过它提供的 webui 访问
现在我的目标就是将这些痛点解决掉,所以我们可能需要做:
- 租一台 GPU 云服务器,并在上面部署 Stable Diffusion
- 写服务器接口,对外提供 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": {}
}
可以看到,区区一个文生图的接口的入参就非常的多,由此可以看出,官方的接口使用成本是很高的。
这些接口其他缺点也很明显:
- 只能由本机访问
- 输入输出的参数都太繁杂,不容易理解
- 图片以 base64 方式传输,使得请求和响应非常冗长
除了希望解决以上问题,我们可能还希望:
- 接口能够识别敏感信息
- 接口能验证用户权限
- 简化请求和响应的参数
为了解决以上问题,我们必须对官方提供的 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 软件,我们还需要做的其他工作:
- 提示词支持中文
- 提示词过滤敏感词语、图像过滤敏感图
- 账号系统,管理不同用户的数据
- 丰富 AI 绘画中可使用的主题风格
- 提高做图质量和分辨率的前提下,再提高做图效率
- GPU 服务器并发访问问题
提示词支持中文
Stable Diffusion 是基于英文标签来进行训练的,所以它的提示词只支持英文,如果我们想要提示词支持中文,首先很容易想到的方案就是在输入前先对提示词进行翻译,也就是中文转英文,实际上现在常规的做法也是如此。
为了实现这个功能,我们需要引入一款翻译插件,这款翻译插件可以将英文转化成多种语言,一款非常好用的翻译插件推荐给大家:
GitHub 地址如下:https://github.com/studyzy/sd-prompt-translator
提示词过滤敏感词语、图像过滤敏感图
为了实现这个功能,我们首先需要准备一份敏感词黑名单,敏感词往往有很多,可能多达 1.5W 个,也就是说每次接口请求的时候,我们都需要遍历一遍名单中的上万个敏感词做出匹配,每当匹配到敏感词,则不允许进行请求。所以,我们需要一款高性能的敏感词识别工具,GitHub 上就有(据说是行业最快的一款敏感词识别工具),地址如下:https://github.com/toolgood/ToolGood.Words
往往做了提示词识别敏感词,已经能防止 90% 以上的情况生成不良图像了。
但如果你想更严谨一点,则可以再做一层敏感图的过滤。
账号系统
往往我们软件会涉及到多个用户,每个用户都有自己的创作内容,也可以自己管理自己的创作内容,所以就需要一个开发账号系统。
这里涉及到的业务有用户登录(一键登录、验证码登录、密码登录等)、用户数据库设计,数据表增删查改。
可以简单看下我的表设计结构:
<update id="createUGCTable">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 = InnoDBAUTO_INCREMENT = 0CHARACTER SET = utf8mb4 COMMENT ='用户创作表';</update>
其他都是一些业务代码,略。
丰富 AI 绘画中可使用的主题风格
在 AI 绘画中,主题和风格可以表现为
- 基础模型(大模型):是 SD 能够绘图的基础模型,必须搭配基础模型才能使用,不同的基础模型画风和擅长的领域会有侧重
- 微调模型(LoRA 模型):其作用就是在大模型的生成图片的基础上对图片按照一定风格进行微调
在哪可以下载这些模型呢?
推荐两个地方:https://civitai.com/、https://huggingface.co/models
提高做图质量和分辨率的前提下,再提高做图效率
常规情况下,SD 支持的分辨范围是 64-2048 分辨率,分辨率越高,需要的显存就越高,生成图片的时间就越慢。
据个人使用体验,在 NVIDIA 3070 8G 显卡上,
- 生成 512*512 分辨率图片,采样步数 50 的情况下,耗时约 15 秒
- 生成 1024*1024 分辨率图片,采样步数 50 的情况下,耗时约 60 秒
想要提高生图效率,一方面需要提升显卡性能,另一方面需要降低图片分辨率,但是图片分辨率太低(低于 512 )又可能导致图片质量不行。
这时有三条方案可以解决:
- 使用 SD 自带的 Extras 模块进行超分辨率修复图片
- 使用 SD 自带的 img2img 模块进行超分辨率修复图片
- 使用 ControlNet 模型的 tile 进行超分辨率修复图片
这里只介绍第一个,因为速度最快,从 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 性能,只负责能分发请求和存储用户数据;其他服务器就是 GPU 服务器,拥有较高的显卡性能,负责生成图片并返回给主服务器。
使用接口,完成软件交互
做完以上这些,我们的接口就相对来说比较完善了,后续的工作只需将接口应用到我们前端、客户端上,我们的软件就是一个完整的 AI 绘图软件了。