万字长文全面理解与应用Stable Diffusion
1. Stable Diffusion简介 1.1 基本概念 1.2 主体结构 1.3 训练细节 1.4 模型评测 1.5 模型应用 1.6 模型版本 1.7 其他类型的条件生成模型 1.8 使用DreamBooth进行微调 2. 实战Stable Diffusion 2.1 环境准备 2.2 从文本生成图像 2.3 Stable Diffusion Pipeline 2.3.1 2.3.2 分词器和文本编码器 2.3.3 UNet 2.3.4 调度器 2.3.5 DIY采样循环 2.4 其他管线应用 2.4.1 Img2Img 2.4.2 Inpainting 2.4.3 Depth2Image 3. Stable Diffusion的特色应用 3.1 个性化生成 3.2 风格化finetune模型 3.3 图像编辑 3.4 可控生成 3.5 Stable-Diffusion-WebUI 参考资料Stable Diffusion是一个强大的文本条件隐式扩散模型(text-conditioned latent diffusion model),它具有根据文字描述生成精美图片的能力。它不仅是一个完全开源的模型(代码,数据,模型全部开源),而且是它的参数量只有
1B
左右,大部分人可以在普通的显卡上进行推理甚至精调模型。毫不夸张的说,Stable Diffusion的出现和开源对AIGC的火热和发展是有巨大推动作用的,因为它让更多的人能快地上手AI作画。本文将基于Hugging Face的diffusers库深入讲解Stable Diffusion的技术原理以及部分的实现细节,然后也会介绍Stable Diffusion的常用功能。
1. Stable Diffusion简介
Stable Diffusion是CompVis、Stability AI和LAION等公司研发的一个文生图模型,它的模型和代码是开源的,而且训练数据LAION-5B
也是开源的。
Stable Diffusion是一个基于Latent的扩散模型,它在UNet中引入text condition来实现基于文本生成图像。Stable Diffusion的核心来源于Latent Diffusion这个工作,常规的扩散模型是基于Pixel(像素)的生成模型,而Latent Diffusion是基于Latent的生成模型,它先采用一个autoencoder将图像压缩到Latent空间,然后用扩散模型来生成图像的Latents,最后送入autoencoder的decoder模块就可以得到生成的图像。
Stable Diffusion模型的主体结构如下图所示,主要包括三个模型:
AutoEncoder
:Encoder将图像压缩到Latent空间,而Decoder将Latent解码为图像;
CLIP text encoder
:提取输入text的text embeddings,通过cross attention方式送入扩散模型的UNet中作为condition;
UNet
:扩散模型的主体,用来实现文本引导下的Latent生成。
对于Stable Diffusion模型,其Autoencoder模型参数大小为84M,CLIP text encoder模型大小为123M,而UNet参数大小为860M,所以Stable Diffusion模型的总参数量约为1B。
基于Latent的扩散模型的优势在于计算效率更高效,因为图像的Latent空间要比图像Pixel空间要小,这也是Stable Diffusion的核心优势
。文生图模型往往参数量比较大,基于Pixel的方法往往限于算力只生成64x64大小的图像,比如OpenAI的DALL-E2和谷歌的Imagen,然后再通过超分辨模型将图像分辨率提升至256x256和1024x1024;而基于Latent的SD是在Latent空间操作的,它可以直接生成256x256和512x512甚至更高分辨率的图像。
1.1 基本概念
隐式扩散
:使用VAE(Variational Auto-Encoder)将图片映射到一个较小的隐式表征,再将其映射到原始图片,通过在隐式表征上进行扩散,可以使用更少的内存,减少UNet层数并加速图片的生成。还可以将结果输入VAE解码器中,得到高分辨率图像。Stable Diffusion中的VAE能够接收一张三通道图片作为输入,从而生成一个四通道的隐式表征,同时每一个空间维度都将减少为原来的八分之一。
以文本为生成条件
:在推理阶段,输入期望图像的文本描述,将纯噪声数据作为起点,然后模型对噪声输入进行“去噪”,生成能匹配文本描述的图像。为此,Stable Diffusion使用了一个名为CLIP的预训练Transformer
模型。
下图展示了文本条件信息(以及基于时间步的条件)是如何在不同位置被输入的。
文本条件信息通过交叉注意力被输入到各个模块,时间步信息通过时间嵌入的映射被输入到各个模块。
无分类器引导
(Classifier-Free Guidance, CFG):主要解决可能得到与文字描述根本不相关的图片,具体方法如下:
1.2 主体结构
(1)AutoEncoder
AutoEncoder是一个基于encoder-decoder
架构的图像压缩模型,对于一个大小为
H
×
W
×
3
H\times W\times 3
H×W×3的输入图像,encoder模块将其编码为一个大小为
h
×
w
×
c
h\times w\times c
h×w×c的Latent,其中
f
=
H
/
h
=
W
/
h
f=H/h=W/h
f=H/h=W/h为下采样率(downsampling factor)。
在训练AutoEncoder过程中,除了采用L1重建损失外,还增加了感知损失(Perceptual Loss,即LPIPS,具体见论文The Unreasonable Effectiveness of Deep Features as a Perceptual Metric)以及基于Patch的对抗训练。
辅助Loss主要是为了确保重建的图像局部真实性以及避免模糊,具体损失函数见Latent Diffusion的Loss部分。同时为了防止得到的latent的标准差过大,采用了两种正则化方法:
第一种是KL-reg,类似VAE增加一个latent和标准正态分布的KL loss,不过这里为了保证重建效果,采用比较小的权重(~10e-6); 第二种是VQ-reg,引入一个VQ (vector quantization)layer,此时的模型可以看成是一个VQ-GAN,不过VQ层是在decoder模块中,这里VQ的codebook采样较高的维度(8192)来降低正则化对重建效果的影响。Stable Diffusion采用基于KL-reg的autoencoder,其中下采样率 f = 8 f=8 f=8,特征维度为 c = 4 c=4 c=4,当输入图像为512x512大小时将得到64x64x4大小的Latent。AutoEncoder模型时在OpenImages数据集上基于256x256大小训练的,但是由于AutoEncoder的模型是全卷积结构的(基于ResnetBlock),所以它可以扩展应用在尺寸>256的图像上。
下面给出使用diffusers库来加载AutoEncoder模型,并使用AutoEncoder来实现图像的压缩和重建的代码:
import torch
import numpy as np
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionPipeline
from PIL import Image
# 加载模型
autoencoder = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
autoencoder.to("cuda", dtype=torch.float16)
# 读取图像并预处理
raw_image = Image.open("boy.jpg").convert("RGB").resize((256, 256))
image = np.array(raw_image).astype(np.float32) / 127.5 - 1.0
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image)
# 压缩图像为Latent并重建
with torch.inference_mode():
latent = autoencoder.encode(image.to("cuda", dtype=torch.float16)).latent_dist.sample()
rec_image = autoencoder.decode(latent).sample
rec_image = (rec_image/2 + 0.5).clamp(0, 1)
rec_image = rec_image.cpu().permute(0,2,3,1).numpy()
rec_image = (rec_image * 255).round().astype("uint8")
rec_image = Image.fromarray(rec_image[0])
rec_image
由于Stable Diffusion采用的AutoEncoder是基于KL-reg的,所以这个AutoEncoder在编码图像时其实得到的是一个高斯分布DiagonalGaussianDistribution
(分布的均值和标准差),然后通过调用sample方法来采样一个具体的Latent(调用mode方法可以得到均值)。
由于KL-reg
的权重系数非常小,实际得到Latent的标准差还是比较大的,Latent diffusion论文中提出了一种rescaling
方法:首先计算出第一个batch数据中的latent的标准差
σ
^
\hat{\sigma}
σ^,然后采用
1
/
σ
^
1/\hat{\sigma}
1/σ^的系数来rescale latent,这样就尽量保证latent的标准差接近1(防止扩散过程的SNR较高,影响生成效果),然后扩散模型也是应用在rescaling的latent上,在解码时只需要将生成的latent除以
1
/
σ
^
1/\hat{\sigma}
1/σ^,然后再送入autoencoder的decoder即可。对于SD所使用的autoencoder,这个rescaling系数为0.18215。
(2)CLIP text encoder
Stable Diffusion采用CLIP text encoder来对输入text提取text embeddings,具体的是采用目前OpenAI所开源的最大CLIP模型:clip-vit-large-patch14
,这个CLIP的text encoder是一个transformer
模型(只有encoder模块):层数为12,特征维度为768,模型参数大小是123M。
对于输入text,送入CLIP text encoder
后得到最后的hidden states
(即最后一个transformer block得到的特征),其特征维度大小为77x768(77是token的数量),这个细粒度的text embeddings将以cross attention的方式送入UNet中。在transofmers
库中,可以如下使用CLIP text encoder:
from transformers import CLIPTextModel, CLIPTokenizer
text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder").to("cuda")
# text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda")
text_tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
# tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
# 对输入的text进行tokenize,得到对应的token ids
prompt = "a photograph of an astronaut riding a horse"
text_input_ids = text_tokenizer(
prompt,
padding="max_length",
max_length=text_tokenizer.model_max_length,
truncation=True,
return_tensors="pt"
).input_ids
# 将token ids送入text model得到77x768的特征
text_embeddings = text_encoder(text_input_ids.to("cuda"))[0]
值得注意的是,这里的text_tokenizer最大长度为77(CLIP训练时所采用的设置),当输入text的tokens数量超过77后,将进行截断,如果不足则进行paddings,这样将保证无论输入任何长度的文本(甚至是空文本)都得到77x768大小的特征。而且在训练Stable Diffusion的过程中,CLIP text encoder模型是冻结的。
补充
在早期的工作中,比如OpenAI的GLIDE
和latent diffusion
中的LDM均采用一个随机初始化的tranformer
模型来提取text的特征,但是最新的工作都是采用预训练好的text model。比如谷歌的Imagen
采用纯文本模型T5 encoder
来提出文本特征,而Stable Diffusion则采用CLIP text encoder
,预训练好的模型往往已经在大规模数据集上进行了训练,它们要比直接采用一个从零训练好的模型要好。
(3)UNet
Stable Diffusion的扩散模型是一个860M的UNet
,其主要结构如下图所示(这里以输入的Latent为64x64x4维度为例),
其中encoder部分包括3个CrossAttnDownBlock2D模块和1个DownBlock2D模块,而decoder部分包括1个UpBlock2D模块和3个CrossAttnUpBlock2D模块,中间还有一个UNetMidBlock2DCrossAttn模块。encoder和decoder两个部分是完全对应的,中间存在skip connection。注意3个CrossAttnDownBlock2D模块最后均有一个2x的downsample操作,而DownBlock2D模块是不包含下采样的。
其中CrossAttnDownBlock2D
模块的主要结构如下图所示,text condition
将通过CrossAttention
模块嵌入进来,此时Attention的query是UNet的中间特征,而key和value则是text embeddings。
Stable Diffusion和DDPM一样采用预测noise的方法来训练UNet,其训练损失也和DDPM一样:
L
s
i
m
p
l
e
=
E
x
0
,
ϵ
∼
N
(
0
,
I
)
,
t
[
∣
∣
ϵ
−
ϵ
θ
(
α
ˉ
t
x
0
+
1
−
α
ˉ
t
ϵ
,
t
,
c
)
∣
∣
2
]
L^{simple}=\mathbb{E}_{x_0,\epsilon \sim \mathcal{N}(0,\mathrm{I}),t}[||\epsilon -\epsilon _{\theta }(\sqrt{\bar{\alpha }_t}\mathrm{x}_0+\sqrt{1-\bar{\alpha }_t}\epsilon,t,\mathrm{c})||^2]
Lsimple=Ex0,ϵ∼N(0,I),t[∣∣ϵ−ϵθ(αˉt
x0+1−αˉt
ϵ,t,c)∣∣2]
这里的
c
c
c为text embeddings,此时的模型是一个条件扩散模型。
在训练条件扩散模型时,往往会采用Classifier-Free Guidance(简称为CFG),所谓的CFG简单来说就是在训练条件扩散模型的同时也训练一个无条件的扩散模型,同时在采样阶段将条件控制下预测的噪音和无条件下的预测噪音组合在一起来确定最终的噪音,具体的计算公式如下所示:
这里的
w
w
w为guidance scale
,当
w
w
w越大时,condition起的作用越大,即生成的图像其更和输入文本一致。CFG的具体实现非常简单,在训练过程中,我们只需要以一定的概率(比如10%)随机drop掉text即可,这里我们可以将text置为空字符串(前面说过此时依然能够提取text embeddings)。
1.3 训练细节
Stable Diffusion的训练主要包括训练数据和训练资源,这方面也是在Stable Diffusion的Model Card上有说明。首先是训练数据,Stable Diffusion在laion2B-en数据集上训练的,它是laion-5b数据集的一个子集,更具体的说它是laion-5b中的英文(文本为英文)数据集。laion-5b数据集是从网页数据Common Crawl中筛选出来的图像-文本对数据集,它包含5.85B的图像-文本对,其中文本为英文的数据量为2.32B,这就是laion2B-en数据集。
下面是laion2B-en数据集的元信息(图片width和height,以及文本长度)统计分析:
其中图片的width和height均在256以上的样本量为1324M,在512以上的样本量为488M,而在1024以上的样本为76M;文本的平均长度为67。Laion数据集中除了图片(下载URL,图像width和height)和文本(描述文本)的元信息外,还包含以下信息:
similarity
:使用CLIP ViT-B/32计算出来的图像和文本余弦相似度;
pwatermark
:使用一个图片水印检测器检测的概率值,表示图片含有水印的概率;
punsafe
:图片是否安全,或者图片是不是NSFW,使用基于CLIP的检测器来估计;
AESTHETIC_SCORE
:图片的美学评分(1-10),这个是后来追加的,首先选择一小部分图片数据集让人对图片的美学打分,然后基于这个标注数据集来训练一个打分模型,并对所有样本计算估计的美学评分。
上面是Laion数据集的情况,下面介绍Stable Diffusion训练数据集的具体情况,Stable Diffusion的训练是多阶段的(先在256x256尺寸上预训练,然后在512x512尺寸上精调),不同的阶段产生了不同的版本:
SD v1.1:在Laion2B-en数据集上以256x256大小训练237,000步,如上所述,laion2B-en数据集中256以上的样本量共1324M;然后在Laion5B的高分辨率数据集以512x512尺寸训练194,000步,这里的高分辨率数据集是图像尺寸在1024x1024以上,共170M样本。 SD v1.2:以SD v1.1为初始权重,在improved_aesthetics_5plus数据集上以512x512尺寸训练515,000步数,这个improved_aesthetics_5plus数据集上laion2B-en数据集中美学评分在5分以上的子集(共约600M样本),注意这里过滤了含有水印的图片(pwatermark>0.5)以及图片尺寸在512x512以下的样本。 SD v1.3:以SD v1.2为初始权重,在improved_aesthetics_5plus数据集上继续以512x512尺寸训练195,000步数,不过这里采用了CFG(以10%的概率随机drop掉text)。 SD v1.4:以SD v1.2为初始权重,在improved_aesthetics_5plus数据集上采用CFG以512x512尺寸训练225,000步数。 SD v1.5:以SD v1.2为初始权重,在improved_aesthetics_5plus数据集上采用CFG以512x512尺寸训练595,000步数。可以看到SD v1.3、SD v1.4和SD v1.5其实是以SD v1.2为起点在improved_aesthetics_5plus数据集上采用CFG训练过程中的不同checkpoints,目前最常用的版本是SD v1.4和SD v1.5。SD的训练是采用了32台8卡的A100机器(32 x 8 x A100_40GB GPUs),所需要的训练硬件还是比较多的。可以简单计算一下,单卡的训练batch size为2,并采用gradient accumulation,其中gradient accumulation steps=2,那么训练的总batch size就是32x8x2x2=2048。训练优化器采用AdamW,训练采用warmup,在初始10,000步后学习速率升到0.0001,后面保持不变。至于训练时间,文档上只说了用了150,000小时,这个应该是A100卡时,如果按照256卡A100来算的话,那么大约需要训练25天左右。
1.4 模型评测
对于文生图模型,目前常采用的定量指标是FID
(Fréchet inception distance)和CLIP score
,其中FID可以衡量生成图像的逼真度(image fidelity),而CLIP score评测的是生成的图像与输入文本的一致性,其中FID越低越好,而CLIP score是越大越好。当CFG的gudiance scale参数设置不同时,FID和CLIP score会发生变化,下图为不同的gudiance scale参数下,SD模型在COCO2017验证集上的评测结果,注意这里是zero-shot
评测,即SD模型并没有在COCO训练数据集上精调。
可以看到当gudiance scale=3时,FID最低;而当gudiance scale越大时,CLIP score越大,但是FID同时也变大。在实际应用时,往往会采用较大的gudiance scale,比如SD模型默认采用7.5,此时生成的图像和文本有较好的一致性。
从不同版本的对比曲线上看,Stable Diffusion的采用CFG训练后三个版本其实差别并没有那么大,其中SD v1.5相对好一点,但是明显要未采用CFG训练的版本要好的多,这说明CFG训练是比较关键的。
但是FID有很大的局限性,它并不能很好地衡量生成图像的质量,也是因为这个原因,谷歌的Imagen引入了人工评价,先建立一个评测数据集DrawBench
(包含200个不同类型的text),然后用不同的模型来生成图像,让人去评价同一个text下不同模型生成的图像,这种评测方式比较直接,但是可能也受一些主观因素的影响。
1.5 模型应用
(1)文生图
根据文本生成图像这是文生图的最核心的功能,下图为Stable Diffusion的文生图的推理流程图:
使用diffusers
库,可以直接调用StableDiffusionPipeline
来实现文生图:
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
# 判断当前的设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')
# 组合图像,生成grid
def image_grid(imgs, rows, cols):
assert len(imgs) == rows*cols
w, h = imgs[0].size
grid = Image.new('RGB', size=(cols*w, rows*h))
grid_w, grid_h = grid.size
for i, img in enumerate(imgs):
grid.paste(img, box=(i%cols*w, i//cols*h))
return grid
# 加载文生图pipeline
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", # 或者使用 SD v1.4: "CompVis/stable-diffusion-v1-4"
torch_dtype=torch.float16
).to(device)
# 输入text,这里text又称为prompt
prompts = [
"a photograph of an astronaut riding a horse",
"A cute otter in a rainbow whirlpool holding shells, watercolor",
"An avocado armchair",
"A white dog wearing sunglasses"
]
# 定义随机seed,保证可重复性
generator = torch.Generator(device).manual_seed(42)
# 执行推理
images = pipe(
prompts,
height=512,
width=512,
num_inference_steps=50,
guidance_scale=7.5,
negative_prompt=None,
num_images_per_prompt=1,
generator=generator
).images
grid = image_grid(images, rows=1, cols=4)
grid
生成的图像效果如下所示:
可以通过指定width
和height
来决定生成图像的大小,前面提到Stable Diffusion最后是在512x512尺度上训练的,所以生成512x512尺寸效果是最好的,但是实际上Stable Diffusion可以生成任意尺寸的图片:一方面AutoEncoder支持任意尺寸的图片的编码和解码,另外一方面扩散模型UNet也是支持任意尺寸的latents生成的(UNet是卷积+attention的混合结构)。
但是注意生成低分辨率图像时,图像的质量会大幅度下降,这是因为训练是在固定尺寸上(512x512)进行的,生成其它尺寸图像还是会存在一定的问题。解决这个问题的相对比较简单,就是采用多尺度策略训练,比如NovelAI提出采用Aspect Ratio Bucketing
策略来在二次元数据集上精调模型,这样得到的模型就很大程度上避免Stable Diffusion的这个问题,目前大部分开源的基于Stable Diffusion的精调模型往往都采用类似的多尺度策略来精调。
这里的另一个比较重要的参数是num_inference_steps
,它是指推理过程中的去噪步数或者采样步数。Stable Diffusion在训练过程采用的是步数为1000的noise scheduler,但是在推理时往往采用速度更快的scheduler:只需要少量的采样步数就能生成不错的图像,比如Stable Diffusion默认采用PNDM scheduler
,它只需要采样50步就可以出图。当然我们也可以换用其它类型的scheduler,比如DDIM scheduler
和DPM-Solver scheduler
。
可以在diffusers中直接替换scheduler,比如使用DDIM:
from diffusers import DDIMScheduler
# 注意这里的clip_sample要关闭,否则生成图像存在问题,因为不能对latent进行clip
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config, clip_sample=False)
换成DDIM后,同样的采样步数生成的图像如下所示,在部分细节上和PNDM有差异:
当然采样步数越大,生成的图像质量越好,但是相应的推理时间也更久。
第三个要讨论的参数是guidance_scale
,当CFG的guidance_scale
越大时,生成的图像应该会和输入文本更一致,但是过大的guidance_scale
会出现问题,这主要是由于训练和测试的不一致,过大的guidance_scale会导致生成的样本超出范围。谷歌的Imagen
论文提出一种dynamic thresholding
策略来解决这个问题,所谓的dynamic thresholding是相对于原来的static thresholding:
static thresholding策略
是直接将生成的样本clip到[-1, 1]范围内(Imagen是基于pixel的扩散模型,这里是将图像像素值归一化到-1到1之间),但是会在过大的guidance_scale时产生很多的饱含像素点。
dynamic thresholding策略
是先计算样本在某个百分位下(比如99%)的像素绝对值
s
s
s,然后如果它超过1时就采用
s
s
s来进行clip,这样就可以大大减少过饱和的像素。
这两种策略的具体代码如下:
dynamic thresholding策略对于Imagen是比较关键的,它使得Imagen可以采用较大的guidance_scale来生成更自然的图像。
另一个比较容易忽略的参数是negative_prompt
,这个参数和CFG有关,前面说过Stable Diffusion采用了CFG来提升生成图像的质量。使用CFG,去噪过程的噪音预测不仅仅依赖条件扩散模型,也依赖无条件扩散模型:
KaTeX parse error: Expected 'EOF', got '_' at position 11: \text{pred_̲noise}=w\epsilo…
这里的negative_prompt为无条件扩散模型的text输入,上文说到训练过程中将text置为空字符串来实现无条件扩散模型,即negative_prompt = None = ""
。但是有时候可以使用不为空的negative_prompt来避免模型生成的图像包含不想要的东西,因为从上述公式可以看到这里的无条件扩散模型是想远离的部分。
合理使用negative prompt
能够帮助去除不想要的东西来提升图像生成效果。 一般情况下,输入的text或者prompt我们称之为“正向提示词”,而negative prompt称之为“反向提示词”,想要生成的好的图像,不仅要选择好的正向提示词,也需要好的反向提示词,这和文本生成模型也比较类似:都需要好的prompt。
一个对正向prompt优化的例子:
原始prompt为"A rabbit is wearing a space suit",直接生成的效果:
将prompt改为"A rabbit is wearing a space suit, digital Art, Greg rutkowski, Trending cinematographic artstation",其生成的效果就大大提升:
(2)图生图
图生图(image2image)是对文生图功能的一个扩展,其核心思路也非常简单:给定一个笔画的色块图像,可以先给它加一定的高斯噪音(执行扩散过程)得到噪音图像,然后基于扩散模型对这个噪音图像进行去噪,就可以生成新的图像,但是这个图像在结构和布局和输入图像基本一致。
对于Stable Diffusion来说,图生图的流程图如下所示,相比文生图流程来说,这里的初始latent不再是一个随机噪音,而是由初始图像经过autoencoder编码之后的latent加高斯噪音得到,这里的加噪过程就是扩散过程。要注意的是,去噪过程的步数要和加噪过程的步数一致,就是说你加了多少噪音,就应该去掉多少噪音,这样才能生成想要的无噪音图像。
在diffusers中,可以使用StableDiffusionImg2ImgPipeline
来实现文生图,具体代码如下:
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image
# 加载图生图pipeline
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float32).to(device)
# 读取初始图片
init_image = Image.open("boy.jpg").convert("RGB")
# 推理
prompt = "A fantasy landscape, trending on artstation"
generator = torch.Generator(device=device).manual_seed(2023)
image = pipe(
prompt=prompt,
image=init_image,
strength=0.8,
guidance_scale=7.5,
generator=generator
).images[0]
image
相比文生图的pipeline,图生图的pipeline还多了一个参数strength,这个参数介于0-1之间,表示对输入图片加噪音的程度,这个值越大加的噪音越多,对原始图片的破坏也就越大,当strength=1时,其实就变成了一个随机噪音,此时就相当于纯粹的文生图pipeline了。
总结来看,图生图其实核心也是依赖了文生图的能力,其中strength这个参数需要灵活调节来得到满意的图像。
(3)图像修补
图像修补,image inpainting
,它和图生图一样也是文生图功能的一个扩展。Stable Diffusion的图像inpainting不是用在图像修复上,而是主要用在图像编辑上:给定一个输入图像和想要编辑的区域mask,通过文生图来编辑mask区域的内容。Stable Diffusion的图像inpainting原理图如下所示:
和图生图一样:首先将输入图像通过autoencoder编码为latent,然后加入一定的高斯噪音生成noisy latent,再进行去噪生成图像,但是这里为了保证mask以外的区域不发生变化,在去噪过程的每一步,都将扩散模型预测的noisy latent用真实图像同level的nosiy latent替换。
在diffusers中,使用StableDiffusionInpaintPipelineLegacy
可以实现文本引导下的图像inpainting,具体代码如下所示:
import torch
from diffusers import StableDiffusionInpaintPipelineLegacy
from PIL import Image
# 加载inpainting pipeline
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionInpaintPipelineLegacy.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
# 读取输入图像和输入mask
input_image = Image.open("overture-creations-5sI6fQgYIuo.png").resize((512, 512))
input_mask = Image.open("overture-creations-5sI6fQgYIuo_mask.png").resize((512, 512))
# 执行推理
prompt = ["a mecha robot sitting on a bench", "a cat sitting on a bench"]
generator = torch.Generator("cuda").manual_seed(0)
with torch.autocast("cuda"):
images = pipe(
prompt=prompt,
image=input_image,
mask_image=input_mask,
num_inference_steps=50,
strength=0.75,
guidance_scale=7.5,
num_images_per_prompt=1,
generator=generator,
).images
下面是一个具体的生成效果,这里我们将输入图像的dog换成了mecha robot或者cat,从而实现了图像编辑。
要注意的是这里的参数guidance_scale也和图生图一样比较重要,要生成好的图像,需要选择合适的guidance_scale。如果guidance_scale=0.5时,生成的图像由于过于受到原图干扰而产生一些不协调。
无论是上面的图生图还是这里的图像inpainting,其实并没有去finetune Stable Diffusion模型,只是扩展了它的能力,但是这两样功能就需要精确调整参数才能得到满意的生成效果。 这里也给出StableDiffusionInpaintPipelineLegacy
这个pipeline内部的核心代码:
import PIL
import numpy as np
import torch
from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler
from transformers import CLIPTextModel, CLIPTokenizer
from tqdm.auto import tqdm
def preprocess_mask(mask):
mask = mask.convert("L")
w, h = mask.size
w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32
mask = mask.resize((w // 8, h // 8), resample=PIL.Image.NEAREST)
mask = np.array(mask).astype(np.float32) / 255.0
mask = np.tile(mask, (4, 1, 1))
mask = mask[None].transpose(0, 1, 2, 3) # what does this step do?
mask = 1 - mask # repaint white, keep black
mask = torch.from_numpy(mask)
return mask
def preprocess(image):
w, h = image.size
w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32
image = image.resize((w, h), resample=PIL.Image.LANCZOS)
image = np.array(image).astype(np.float32) / 255.0
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image)
return 2.0 * image - 1.0
model_id = "runwayml/stable-diffusion-v1-5"
# 1. 加载autoencoder
vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae")
# 2. 加载tokenizer和text encoder
tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder")
# 3. 加载扩散模型UNet
unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet")
# 4. 定义noise scheduler
noise_scheduler = DDIMScheduler(
num_train_timesteps=1000,
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False, # don't clip sample, the x0 in stable diffusion not in range [-1, 1]
set_alpha_to_one=False,
)
# 将模型复制到GPU上
device = "cuda"
vae.to(device, dtype=torch.float16)
text_encoder.to(device, dtype=torch.float16)
unet = unet.to(device, dtype=torch.float16)
prompt = "a mecha robot sitting on a bench"
strength = 0.75
guidance_scale = 7.5
batch_size = 1
num_inference_steps = 50
negative_prompt = ""
generator = torch.Generator(device).manual_seed(0)
with torch.no_grad():
# 获取prompt的text_embeddings
text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
# 获取unconditional text embeddings
max_length = text_input.input_ids.shape[-1]
uncond_input = tokenizer(
[negative_prompt] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
)
uncond_embeddings = text_encoder(uncond_input.input_ids.to(device))[0]
# 拼接batch
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
# 设置采样步数
noise_scheduler.set_timesteps(num_inference_steps, device=device)
# 根据strength计算timesteps
init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
t_start = max(num_inference_steps - init_timestep, 0)
timesteps = noise_scheduler.timesteps[t_start:]
# 预处理init_image
init_input = preprocess(input_image)
init_latents = vae.encode(init_input.to(device, dtype=torch.float16)).latent_dist.sample(generator)
init_latents = 0.18215 * init_latents
init_latents = torch.cat([init_latents] * batch_size, dim=0)
init_latents_orig = init_latents
# 处理mask
mask_image = preprocess_mask(input_mask)
mask_image = mask_image.to(device=device, dtype=init_latents.dtype)
mask = torch.cat([mask_image] * batch_size)
# 给init_latents加噪音
noise = torch.randn(init_latents.shape, generator=generator, device=device, dtype=init_latents.dtype)
init_latents = noise_scheduler.add_noise(init_latents, noise, timesteps[:1])
latents = init_latents # 作为初始latents
# Do denoise steps
for t in tqdm(timesteps):
# 这里latens扩展2份,是为了同时计算unconditional prediction
latent_model_input = torch.cat([latents] * 2)
latent_model_input = noise_scheduler.scale_model_input(latent_model_input, t) # for DDIM, do nothing
# 预测噪音
noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
# CFG
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# 计算上一步的noisy latents:x_t -> x_t-1
latents = noise_scheduler.step(noise_pred, t, latents).prev_sample
# 将unmask区域替换原始图像的nosiy latents
init_latents_proper = noise_scheduler.add_noise(init_latents_orig, noise, torch.tensor([t]))
latents = (init_latents_proper * mask) + (latents * (1 - mask))
# 注意要对latents进行scale
latents = 1 / 0.18215 * latents
image = vae.decode(latents).sample
另外,runwayml在发布SD 1.5版本的同时还发布了一个inpainting模型:runwayml/stable-diffusion-inpainting,与前面所讲不同的是,这是一个在SD 1.2上finetune的模型。原来Stable Diffusion的UNet的输入是64x64x4,为了实现inpainting,现在给UNet的第一个卷积层增加5个channels,分别为masked图像的latents(经过autoencoder编码,64x64x4)和mask图像(直接下采样8x,64x64x1),增加的权重填零初始化。
在diffusers中,可以使用StableDiffusionInpaintPipeline来调用这个模型,具体代码如下:
import torch
from diffusers import StableDiffusionInpaintPipeline
from PIL import Image
from tqdm.auto import tqdm
import PIL
# Load pipeline
model_id = "runwayml/stable-diffusion-inpainting/"
pipe = StableDiffusionInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
prompt = ["a mecha robot sitting on a bench", "a dog sitting on a bench", "a bench"]
generator = torch.Generator("cuda").manual_seed(2023)
input_image = Image.open("overture-creations-5sI6fQgYIuo.png").resize((512, 512))
input_mask = Image.open("overture-creations-5sI6fQgYIuo_mask.png").resize((512, 512))
images = pipe(
prompt=prompt,
image=input_image,
mask_image=input_mask,
num_inference_steps=50,
generator=generator,
).images
其生成的效果图如下所示:
经过finetune的inpainting在生成细节上可能会更好,但是有可能会丧失部分文生图的能力,而且也比较难迁移其它finetune的Stable Diffusion模型。
1.6 模型版本
(1)Stable Diffusion V2
Stability AI公司在2022年11月(stable-diffusion-v2-release)放出了SD 2.0版本,SD 2.0相比SD 1.x版本的主要变动在于模型结构和训练数据两个部分。
首先是模型结构,SD 1.x版本的text encoder采用的是OpenAI的CLIP ViT-L/14模型,其模型参数量为123.65M;而SD 2.0采用了更大的text encoder:基于OpenCLIP在laion-2b数据集上训练的CLIP ViT-H/14模型,其参数量为354.03M,相比原来的text encoder模型大了约3倍。
然后是训练数据,前面说过SD 1.x版本其实最后主要采用Laion-2B中美学评分为5以上的子集来训练,而SD 2.0版本采用评分在4.5以上的子集,相当于扩大了训练数据集,具体的训练细节见model card。 另外SD 2.0除了512x512版本的模型,还包括768x768版本的模型(https://huggingface.co/stabilityai/stable-diffusion-2),所谓的768x768模型是在512x512模型基础上用图像分辨率大于768x768的子集继续训练的,不过优化目标不再是noise_prediction,而是采用Progressive Distillation for Fast Sampling of Diffusion Models论文中所提出的 v-objective。
Stability AI在发布SD 2.0的同时,还发布了另外3个模型:stable-diffusion-x4-upscaler,stable-diffusion-2-inpainting和stable-diffusion-2-depth。
在diffusers库中,可以如下使用这个超分模型(这里的noise level是指推理时对低分辨率图像加入噪音的程度):
import requests
from PIL import Image
from io import BytesIO
from diffusers import StableDiffusionUpscalePipeline
import torch
# load model and scheduler
model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipeline = pipeline.to("cuda")
# let's download an image
url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale/low_res_cat.png"
response = requests.get(url)
low_res_img = Image.open(BytesIO(response.content)).convert("RGB")
low_res_img = low_res_img.resize((128, 128))
prompt = "a white cat"
upscaled_image = pipeline(prompt=prompt, image=low_res_img, noise_level=20).images[0]
upscaled_image.save("upsampled_cat.png")
stable-diffusion-2-depth是也是在SD 2.0的512x512版本上finetune的模型,它是额外增加了图像的深度图作为condition,这里是直接将深度图下采样8x,然后和nosiy latent拼接在一起送入UNet模型中。深度图可以作为一种结构控制,下图展示了加入深度图后生成的图像效果:
原图如下:
可以调用diffusers库中的StableDiffusionDepth2ImgPipeline
来实现基于深度图控制的文生图:
import torch
import requests
from PIL import Image
from diffusers import StableDiffusionDepth2ImgPipeline
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
).to("cuda")
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
init_image = Image.open(requests.get(url, stream=True).raw)
prompt = "two tigers"
n_propmt = "bad, deformed, ugly, bad anotomy"
image = pipe(prompt=prompt, image=init_image, negative_prompt=n_propmt, strength=0.7).images[0]
除此之外,Stability AI公司还开源了两个加强版的autoencoder:ft-EMA
和ft-MSE
(前者使用L1 loss后者使用MSE loss),它们是在LAION数据集继续finetune decoder来增强重建效果。
(2)Stable Diffusion V2.1
在SD 2.0版本发布几周后,Stability AI又发布了SD 2.1。SD 2.0在训练过程中采用NSFW检测器过滤掉了可能包含色情的图像(punsafe=0.1),但是也同时过滤了很多人像图片,这导致SD 2.0在人像生成上效果可能较差,所以SD 2.1是在SD 2.0的基础上放开了限制(punsafe=0.98)继续finetune,所以增强了人像的生成效果。
(3)Stable Diffusion unclip
Stability AI在2023年3月份,又放出了基于Stable Diffusion的另外一个模型:stable-diffusion-reimagine,它可以实现单个图像的变换,即image variations,目前该模型已经在在huggingface上开源:stable-diffusion-2-1-unclip。
这个模型是借鉴了OpenAI的DALLE2(又称unCLIP),unCLIP是基于CLIP的image encoder提取的image embeddings作为condition来实现图像的生成。
Stable Diffusion unCLIP是在原来的Stable Diffusion模型的基础上增加了CLIP的image encoder的nosiy image embeddings作为condition。具体来说,它在训练过程中是对提取的image embeddings施加一定的高斯噪音(也是通过扩散过程),然后将noise level对应的time embeddings和image embeddings拼接在一起,最后再以class labels的方式送入UNet。
在diffusers中,可以调用StableUnCLIPImg2ImgPipeline
来实现图像的变换:
import requests
import torch
from PIL import Image
from io import BytesIO
from diffusers import StableUnCLIPImg2ImgPipeline
#Start the StableUnCLIP Image variations pipeline
pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1-unclip", torch_dtype=torch.float16, variation="fp16"
)
pipe = pipe.to("cuda")
#Get image from URL
url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/tarsila_do_amaral.png"
response = requests.get(url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
#Pipe to make the variation
images = pipe(init_image).images
images[0].save("tarsila_variation.png")
其实在Stable Diffusion unCLIP之前,已经有Lambda Labs开源的sd-image-variations-diffusers,它是在SD 1.4的基础上finetune的模型,不过实现方式是直接将text embeddings替换为image embeddings,这样也同样可以实现图像的变换。
这里Stable Diffusion unCLIP有两个版本:sd21-unclip-l
和sd21-unclip-h
,两者分别是采用OpenAI CLIP-L
和OpenCLIP-H
模型的image embeddings作为condition。如果要实现文生图,还需要像DALLE2那样训练一个prior模型,它可以实现基于文本来预测对应的image embeddings,将prior模型和SD unCLIP接在一起就可以实现文生图了。
1.7 其他类型的条件生成模型
Img2Img
Img2Img是图片到图片的转换,包括多种类型,如风格转换(从照片风格转换为动漫风格)和图片超分辨率(给定一张低分辨率图片作为条件,让模型生成对应的高分辨率图片,类似于Stable Diffusion Upscaler)。
Inpainting
Inpainting又称为图片修复,它是图片的部分掩膜到图片的转换,模型会根据掩膜的区域信息和掩膜之外的全局结构信息,用掩膜区域生成连贯的图片,而掩膜区域之外则与原图保持一致。
Depth2Img
Depth2Img采用图片的深度图作为条件,模型会生成与深度图本身相似的具有全局结构的图片。
1.8 使用DreamBooth进行微调
DreamBooth是一种个性化训练一个文本到图像模型的方法,只需要提供一个主题的3~5张图像,就能教会模型有关这个主题的各种概念,从而在不同的场景和视图中生成这个主题的相关图像。
2. 实战Stable Diffusion
2.1 环境准备
!pip install -Uq diffusers fifty accelerate transformers
查看GPU:
import torch
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')
下载图像:
import torch
import requests
from PIL import Image
from io import BytesIO
from matplotlib import pyplot as plt
from diffusers import (
StableDiffusionPipeline,
StableDiffusionImg2ImgPipeline,
StableDiffusionInpaintPipeline,
StableDiffusionDepth2ImgPipeline
)
def download_image(url):
response = requests.get(url)
return Image.open(BytesIO(response.content)).convert("RGB")
img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
init_image = download_image(img_url).resize((512, 512))
mask_image = download_image(mask_url).resize((512, 512))
device = "cuda" if torch.cuda.is_available() else "cpu"
2.2 从文本生成图像
载入管线:
如果GPU显存不足,可以尝试通过如下方法减少对GPU显存的使用:
pipe = StableDiffusionPipeline.from_pretrained(model_id, revision="fp16", torch_dtype=torch.float16).to(device)
开启注意力切分功能,旨在通过降低速度来减少GPU显存的使用,代码如下:
pipe.enable_attention_slicing()
减少所生成图片的尺寸。
管线加载完毕后,通过如下代码,利用文本提示语生成图片:
# 给生成器设置一个随机种子
generator = torch.Generator(device=device).manual_seed(42)
pipe_output = pipe(
prompt="Palette knife painting of an autumn cityscape", # 提示语:哪里要生成
negative_prompt="Oversaturated, blurry, low quality", # 提示语:哪里不需要生成
height=480, width=640, # 图片大小
guidance_scale=8, # 提示文字的影响程度
num_inference_steps=35, # 推理步数
generator=generator # 设置随机种子生成器
)
pipe_output.images[0]
主要的调节参数如下:
以下代码能加大guidance_scale参数的作用:
# 对比不同的guidance_scale效果(该参数决定了无分类器引导的影响强度)
cfg_scales = [1.1, 8, 12]
prompt = "A collie with a pink hat"
fig, axs = plt.subplots(1, len(cfg_scales), figsize=(16, 5))
for i, ax in enumerate(axs):
im = pipe(prompt, height=480, width=480,
guidance_scale=cfg_scales[i], num_inference_steps=35,
generator=torch.Generator(device=device).manual_seed(42)).images[0]
ax.imshow(im); ax.set_title(f'CFG Scale {cfg_scales[i]}');
2.3 Stable Diffusion Pipeline
查看Stable Diffusion Pipeline的组成部分:
2.3.1
可变分自编码器(VAE)是一种模型,VAE可以对输入图像进行编码,得到“压缩过”的信息,之后再解码“隐式的”压缩信息,得到接近输入的输出。
UNet的输入不是完整的图片,而是缩小版的特征,这样可以极大地减少对计算资源的使用。
# 创建区间为(-1, 1)的伪数据
images = torch.rand(1, 3, 512, 512).to(device) * 2 - 1
print("Input images shape:", images.shape)
# 编码到隐空间
with torch.no_grad():
latents = 0.18215 * pipe.vae.encode(images).latent_dist.mean
print("Encoded latents shape:", latents.shape)
# 解码
with torch.no_grad():
decoded_images = pipe.vae.decode(latents / 0.18215).sample
print("Decoded images shape:", decoded_images.shape)
代码输出内容:
Input images shape: torch.Size([1, 3, 512, 512])
Encoded latents shape: torch.Size([1, 4, 64, 64])
Decoded images shape: torch.Size([1, 3, 512, 512])
2.3.2 分词器和文本编码器
文本编码器的作用是将输入的字符串(文本提示语)转换成数值表示形式,这样才能将其输入UNet作为条件。文本则被管线中的分词器(tokenizer)转换成一系列的token(分词)。文本编码器是一个Transformer模型,它被训练用于CLIP。
下面首先手动分词,并将分词结果输入文本编码器,然后使用管线的_encode_prompt方法调用这个编码过程,其间补全或截断分词串的长度;最后使得分词串的长度等于最大长度7。
接下来,获取最终的文本特征:
文本嵌入(text embedding)是指文本编码器中最后一个Transformer模块的“隐状态”(hidden state),它们将作为UNet中的forward函数的一个额外输入。
2.3.3 UNet
在扩散模型中,UNet的作用是接收“带噪”的输入并预测噪声,以实现“去噪”。此处输入模型的并非图片,而是图片的隐式表示形式,此外,除了将用于暗示“带噪”程度的时间步作为条件输入UNet之外,模型还将文本提示语和文本嵌入也作为UNet的输入。
首先使用伪输入尝试让模型进行预测。注意在此过程中各个输入输出的形状和大小:
# 创建伪输入
timestep = pipe.scheduler.timesteps[0]
latents = torch.randn(1, 4, 64, 64).to(device)
text_embeddings = torch.randn(1, 77, 1024).to(device)
# 模型预测
with torch.no_grad():
unet_output = pipe.unet(latents, timestep, text_embeddings).sample
print('UNet output shape:', unet_output.shape) # UNet output shape: torch.Size([1, 4, 64, 64])
2.3.4 调度器
调度器保存了关于如何添加噪声的信息,并管理如何基于模型的预测更新“带噪”样本。默认的调度器是PNDMScheduler。
可以通过绘制图片来观察在添加噪声的过程中噪声水平(基于参数 α ˉ \bar{\alpha} αˉ)随时间步增加的变化趋势:
plt.plot(pipe.scheduler.alphas_cumprod, label=r'$\bar{\alpha}$')
plt.xlabel('Timestep (high noise to low noise ->)')
plt.title('Noise schedule')
plt.legend();
尝试不同的调度器:
生成的图片:
2.3.5 DIY采样循环
将管线的不同组成部分组装在一起,复现整个管线的功能:
guidance_scale = 8
num_inference_steps=30
prompt = "Beautiful picture of a wave breaking"
negative_prompt = "zoomed in, blurry, oversaturated, warped"
# 对提示文字进行编码
text_embeddings = pipe._encode_prompt(prompt, device, 1, True, negative_prompt)
# 创建随机噪声作为起点
latents = torch.randn((1, 4, 64, 64), device=device, generator=generator)
latents *= pipe.scheduler.init_noise_sigma
# 设置调度器
pipe.scheduler.set_timesteps(num_inference_steps, device=device)
# 循环采样
for i, t in enumerate(pipe.scheduler.timesteps):
latent_model_input = torch.cat([latents] * 2)
latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)
with torch.no_grad():
noise_pred = pipe.unet(latent_model_input, t, text_embeddings).sample
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = pipe.scheduler.step(noise_pred, t, latents).prev_sample
# 将隐变量映射到图片
with torch.no_grad():
image = pipe.decode_latents(latents.detach())
pipe.numpy_to_pil(image)[0]
生成的结果:
2.4 其他管线应用
2.4.1 Img2Img
Img2Img首先会对一张已有的图片进行编码,得到隐变量后添加随机噪声,并以此作为起点。噪声的数量和“去噪”所需步数决定了Img2Img过程的“强度”.
Img2Img管线不需要任何特殊的模型,而只需要与文字到图像模型相同的模型ID,无需下载新文件。
首先载入Img2Img管线:
model_id = "stabilityai/stable-diffusion-2-1-base"
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, revision="fp16", torch_dtype=torch.float16).to(device)
Img2Img管线代码:
result_image = img2img_pipe(
prompt="An oil painting of a man on a bench",
image = init_image,
strength = 0.6, # 强度:0表示完全不起作用,1表示作用强度最大
).images[0]
fig, axs = plt.subplots(1, 2, figsize=(12, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(result_image);axs[1].set_title('Result');
2.4.2 Inpainting
关于Inpainting的内容上文已经详细说明了,这里回顾以下。Stable Diffusion模型接收一张掩模图片作为额外条件输入,该掩模图片与输入图片的尺寸一致,白色区域表示要替换的部分,黑色区域表示要保留的部分。
以下代码展示了如何载入StableDiffusionInpaintPipelineLegacy管线并将其应用于前面输入的示例图片和掩膜图片:
pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", revision="fp16", torch_dtype=torch.float16)
pipe = pipe.to(device)
prompt = "A small robot, high resolution, sitting on a park bench"
image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
fig, axs = plt.subplots(1, 3, figsize=(16, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(mask_image);axs[1].set_title('Mask')
axs[2].imshow(image);axs[2].set_title('Result');
Hugging Face Spaces上的一个示例Space应用就使用了一个名为CLIPSeg
的模型,旨在根据文字描述自动地通过掩膜去掉一个物体。
2.4.3 Depth2Image
Depth2Image采用深度预测模型来预测一个深度图,该深度图被输入微调过的UNet以生成图片。希望生成的图片既能保留原始图片的深度信息和总体结构,同时又能在相关部分填入全新的内容:
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-depth",revision="fp16", torch_dtype=torch.float16)
pipe = pipe.to(device)
prompt = "An oil painting of a man on a bench"
image = pipe(prompt=prompt, image=init_image).images[0]
fig, axs = plt.subplots(1, 2, figsize=(16, 5))
axs[0].imshow(init_image);axs[0].set_title('Input Image')
axs[1].imshow(image);axs[1].set_title('Result');
Depth2Img管线的生成结果:
3. Stable Diffusion的特色应用
3.1 个性化生成
个性化生成是指的生成特定的角色或者风格,比如给定自己几张肖像来利用SD来生成个性化头像。在个性化生成方面,比较重要的两个工作是英伟达的Textual Inversion
和谷歌的DreamBooth
。
Textual Inversion
这个工作的核心思路是基于用户提供的3~5张特定概念(物体或者风格)的图像来学习一个特定的text embeddings,实际上只用一个word embedding就足够了。Textual Inversion
不需要finetune UNet,而且由于text embeddings较小,存储成本很低。目前diffusers库已经支持textual_inversion的训练。
DreamBooth
原本是谷歌提出的应用在Imagen上的个性化生成,但是它实际上也可以扩展到Stable Diffusion上。DreamBooth首先为特定的概念寻找一个特定的描述词[V],这个特定的描述词只要是稀有的就可以,然后与Textual Inversion不同的是DreamBooth需要finetune UNet,这里为了防止过拟合,增加了一个class-specific prior preservation loss(基于SD生成同class图像加入batch里面训练)来进行正则化。由于finetune了UNet,DreamBooth往往比Textual Inversion要表现的要好,但是DreamBooth的存储成本较高。目前diffusers库已经支持dreambooth训练,也可以在sd-dreambooth-library中找到其他人上传的模型。
还有很多其它的研究工作,比如Adobe提出的Custom Diffusion,相比DreamBooth,它只finetune了UNet的attention模块的KV权重矩阵,同时优化一个新概念的token。
3.2 风格化finetune模型
Stable Diffusion的另外一大应用是采用特定风格的数据集进行finetune,这使得模型“过拟合”在特定的风格上。之前比较火的novelai就是基于二次元数据在Stable Diffusion上finetune的模型,虽然它失去了生成其它风格图像的能力,但是它在二次元图像的生成效果上比原来的Stable Diffusion要好很多。
目前已经有很多风格化的模型在huggingface上开源,这里也列出一些:
值得说明的一点是,目前finetune Stable Diffusion模型的方法主要有两种:一种是直接finetune了UNet,但是容易过拟合,而且存储成本;另外一种低成本的方法是基于微软的LoRA,LoRA本来是用于finetune语言模型的,但是现在已经可以用来finetune Stable Diffusion模型了
3.3 图像编辑
图像编辑也是Stable Diffusion比较火的应用方向,这里所说的图像编辑是指的是使用Stable Diffusion来实现对图片的局部编辑。这里列举两个比较好的工作:谷歌的prompt-to-prompt
和加州伯克利的instruct-pix2pix
。
3.4 可控生成
可控生成是Stable Diffusion最近比较火的应用,这主要归功于ControlNet,基于ControlNet可以实现对很多种类的可控生成,比如边缘,人体关键点,草图和深度图等等。
其实在ControlNet之前,也有一些可控生成的工作,比如stable-diffusion-2-depth也属于可控生成,但是没有太火。
与ControlNet同期的工作还有腾讯的T2I-Adapter以及阿里的composer-page:
3.5 Stable-Diffusion-WebUI
最后要介绍的一个比较火的应用stable-diffusion-webui其实是用来支持Stable Diffusion出图的一个web工具,它算是基于gradio
框架实现了Stable Diffusion的快速部署,不仅支持Stable Diffusion的最基础的文生图、图生图以及图像inpainting功能,还支持Stable Diffusion的其它拓展功能,很多基于Stable Diffusion的拓展应用可以用插件的方式安装在webui上。