2024不可不会的StableDiffusion之变分自编码器（三）_stable diffusion vae_encoder vae_decoder

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")
1
2
3
4
5
6
7
8
9
10
11
12
13

由于我们之前已经下载过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
1
2
3
4
5

上述函数，接收图像和我们的变分自编码器作为输入，通过调用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()
1
2
3
4
5

得到结果如下：

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}")
1
2

得到结果如下：

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()
1
2
3
4
5

得到可视化结果如下：

上图所示，可以看出，这种潜在的特征表示应该捕捉到许多关于原始图像的信息。

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
1
2
3
4
5
6
7
8
9
10

8. 验证解码器

让我们在上述编码后得到的特征表示1x4x64x64上使用解码辅助函数，看看我们可以得到什么。相关函数调用如下所示：

# decode
decoded_img = latent_to_pil(latent_img,vae)
plt.imshow(decoded_img[0])
plt.show()
1
2
3
4

得到结果如下：

从上图中我们可以看出，VAE解码器能够从48x压缩的潜在表示中恢复原始图像。是不是很神奇！
如果大家仔细观察解码后的图像，它与原始图像不同，请注意眼睛周围的差异。
这就是为什么VAE编码器/解码器不是无损压缩的原因。

9. VAE 在SD中的用途

稳定扩散模型（Stable Diffusion）可以在没有VAE组件的情况下进行，但是我们使用VAE的原因是减少生成高分辨率图像的计算时间。潜在扩散模型（latent diffusion models）可以在VAE编码器产生的这个潜在空间中（Latent Space）进行扩散，一旦我们通过扩散过程产生了所需的潜在输出，我们就可以使用VAE解码器将其转换回高分辨率图像。

关于VAE更多的信息，大家可以访问论文进行全面的了解。

10. 总结

本文重点介绍了SD模型中的变分自编码器VAE的相关功能和具体工作原理，并详细介绍了其编码器和解码器的操作步骤，并给出了相应的代码示例。

您学废了嘛！