当前位置:AIGC资讯 > AIGC > 正文

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)

更新时间 2024-01-22