简介
Stable Diffusion在逐步去噪的过程中,按照图片的原尺寸进行运算,由于图片尺寸的太大,需要很大的计算量。sd通常是在latent space(潜在空间)运算的,相当于在压缩的图片上运算的,然后再恢复到原尺寸。VAE的encode和decode就是sd连接latent space原真实图的通道。
sd生图的时候,有无VAE的区别很大,效果图如下。
VAE,作为生成式模型,有许多应用。VAE在Stable Diffusion中是重要的一环,我们从sd的角度解析一下VAE。
VAE全称Variational Autoencoder(变分自编码器),主要有encode和decode两个过程,encode将图片转到latent space,是一个压缩的过程,压缩时只保留了关键信息,可以再通过decode复原图片。压缩再复原的过程并不是无损的。encode过程主要是一些卷积运算和下采样,和ResNet很类似,就不赘述了。
encode
首先展示encode过程,原图如下,尺寸512 × 512:
latent space的图片如下,尺寸64 × 64,这个图片是我放大后截屏的,放64 × 64的图太小了,平台的水印太大:
这个图的颜色不太正,有可能是我相关系数没调整好,但是可以看到相应的结构信息都保存下来了,我使用的是sd提供的vae模型,代码如下:
import cv2
import torch
import numpy as np
from PIL import Image
from diffusers.models import AutoencoderKL
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema")
image = Image.open("1.png")
image = image.convert('RGB')
tensor = torch.from_numpy(np.array(image)).float().div(255.0)
tensor = tensor.unsqueeze(0)
tensor = tensor.permute(0, 3, 1, 2)
latents = vae.encode(tensor,return_dict=False)
latents = latents[0].mean
def quantize(latents):
quantized_latents = (latents / (255 * 0.18215) + 0.5).clamp(0,1)
quantized = quantized_latents.cpu().permute(0, 2, 3, 1).detach().numpy()[0]
quantized = (quantized * 255.0 + 0.5).astype(np.uint8)
return quantized
zx = quantize(latents)
imageOut = cv2.cvtColor(zx, cv2.COLOR_RGBA2BGRA)
cv2.imwrite('2.png', imageOut)
其中quantize函数是用的deephub的,https://colab.research.google.com/drive/1Ci1VYHuFJK5eOX9TB0Mq4NsqkeDrMaaH?usp=sharing#scrollTo=K29uYxMZTD82
vae模型是从diffusers导入的,运行代码自动下载,也就是和sd用的模型是一样的,具体的模型信息在下载的时候有,这里也可以改成你想用的模型。
decode
该过程输入是latent space的结果,再解码为尺寸512 × 512的图片,如下:
可以看到基本还原了图片,通过多组对比,我们可以看到复原的图片和原图仍有差别,但是复原的图片质量是很高的。代码如下,我们只需要接着上面的写就行,输入是latents:
decodeOut = vae.decode(latents,return_dict=False)
decodeOut = decodeOut[0].permute(2, 3, 1, 0).squeeze()
numpy_img = decodeOut.detach().numpy()
numpy_img = numpy_img*255
numpy_img = cv2.cvtColor(numpy_img, cv2.COLOR_RGB2BGR)
cv2.imwrite('3.png', numpy_img)