Stable Diffusion VAE

简介

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)