1. 引言
这是我关于StableDiffusion
学习系列的第三篇文章,如果之前的文章你还没有阅读,强烈推荐大家翻看前篇内容。在本文中,我们将学习构成StableDiffusion
的第二个基础组件变分自编码器VAE
,并针该组件的功能进行详细的阐述。
闲话少说,我们直接开始吧!
2. 概览
通常来说一个自编码器autoencoder
包含两部分:
-
Encoder
: 将图像作为输入,并将其转换为潜在特征空间的低维度表示 -
Decoder
: 将低纬度特征表示作为输入,并将其解码为图像进行输出
整体过程如下所示:
正如我们在上图看到的,编码器就像一个压缩器,将图像压缩到较低的维度,解码器从压缩表示中重新创建原始图像。需要注意的是编码器解码器压缩解压缩并不是无损的。让我们开始通过代码来研究VAE
。
我们将从导入所需的库和定义一些辅助函数开始。
3. 导入所需的库
首先让我们导入我们所需要的Python
基础库,并加载我们的VAE
模型,代码如下:
## Imaging library
from PIL import Image
from torchvision import transforms as tfms
## Basic libraries
import numpy as np
import matplotlib.pyplot as plt
## Loading a VAE model
from diffusers import AutoencoderKL
sd_path = r'/media/stable_diffusion/stable-diffusion-v1-4'
vae = AutoencoderKL.from_pretrained(sd_path,subfolder="vae",
local_files_only=True,
torch_dtype=torch.float16).to("cuda")
由于我们之前已经下载过stable-diffusion-v1-4
相关文件,在其子目录下存在vae
目录,即为本节需要测试验证的变分自编码器,此时需要将变量local_files_only
设置为True
,表示从本地读取相关权重文件。
4. 定义编码辅助函数
接着我们来实现用VAE
对图像进行编码操作的辅助函数,其相关定义如下:
def pil_to_latents(image,vae):
init_image = tfms.ToTensor()(image).unsqueeze(0) * 2.0 - 1.0
init_image = init_image.to(device="cuda",dtype=torch.float16)
init_latent_dist = vae.encode(init_image).latent_dist.sample() * 0.18215
return init_latent_dist
上述函数,接收图像和我们的变分自编码器作为输入,通过调用vae
中的encode
函数来实现将输入的image
转化为潜在低纬度特征向量。
5. 读取测试图像
接着,我们就可以读取测试图像来进行编码器功能验证了。首先,我们来进行图像读取工作:
img_path = r'/home/VAEs/Scarlet-Macaw-2.jpg'
img = Image.open(img_path).convert("RGB").resize((512,512))
print(f"Dimension of this image : {np.array(img).shape}")
plt.imshow(img)
plt.show()
得到结果如下:
Dimension of this image: (512, 512, 3)
6. 验证编码器
现在,让我们使用VAE编码器来压缩此图像,我们将使用pil_to_latents
辅助函数对齐进行编码操作,代码如下:
latent_img = pil_to_latents(img,vae)
print(f"Dimension of this latent representation: {latent_img.shape}")
得到结果如下:
Dimension of this latent representation: torch.Size([1, 4, 64, 64])
正如我们所看到的,VAE将3 x 512 x 512
维图像压缩为4 x 64 x 64
维图像的。这是48倍的压缩比!接着让我们可视化一下这四个通道的潜在特征。
# visual
fig,axs = plt.subplots(1,4,figsize=(16,4))
for c in range(4):
axs[c].imshow(latent_img[0][c].detach().cpu(),cmap='Greys')
plt.show()
得到可视化结果如下:
上图所示,可以看出,这种潜在的特征表示应该捕服务器托管网捉到许多关于原始图像的信息。
7. 定义解码辅助函数
接着,与编码操作类似,我们来定义解码的辅助函数,其过程为编码操作的逆过程,相关代码示例如下:
def latent_to_pil(latents,vae):
latents = (1/0.18215) * latents
with torch.no_grad():
服务器托管网 image = vae.decode(latents).sample
image = (image / 2 + 0.5).clamp(0,1)
image = image.detach().cpu().permute(0,2,3,1).numpy()
images = (image * 255).round().astype("uint8")
pil_images = [ Image.fromarray(image) for image in images ]
return pil_images
8. 验证解码器
让我们在上述编码后得到的特征表示1x4x64x64
上使用解码辅助函数,看看我们可以得到什么。相关函数调用如下所示:
# decode
decoded_img = latent_to_pil(latent_img,vae)
plt.imshow(decoded_img[0])
plt.show()
得到结果如下:
从上图中我们可以看出,VAE解码器能够从48x压缩的潜在表示中恢复原始图像。是不是很神奇!
如果大家仔细观察解码后的图像,它与原始图像不同,请注意眼睛周围的差异。
这就是为什么VAE编码器/解码器不是无损压缩的原因。
9. VAE 在SD中的用途
稳定扩散模型(Stable Diffusion)
可以在没有VAE
组件的情况下进行,但是我们使用VAE
的原因是减少生成高分辨率图像的计算时间。潜在扩散模型(latent diffusion models)
可以在VAE
编码器产生的这个潜在空间中(Latent Space)
进行扩散,一旦我们通过扩散过程产生了所需的潜在输出,我们就可以使用VAE
解码器将其转换回高分辨率图像。
关于VAE更多的信息,大家可以访问论文进行全面的了解。
10. 总结
本文重点介绍了SD
模型中的变分自编码器VAE
的相关功能和具体工作原理,并详细介绍了其编码器和解码器的操作步骤,并给出了相应的代码示例。
您学废了嘛!
服务器托管,北京服务器托管,服务器租用 http://www.fwqtg.net
机房租用,北京机房租用,IDC机房托管, http://www.fwqtg.net