本篇文章聊聊 Stable Diffusion 生态中呼声最高、也是最复杂的开源模型管理图形界面 “stable-diffusion-webui” 中和 VAE 相关的事情。
写在前面
Stable Diffusion 生态中有一个很重要的项目，它对于 SD 生态繁荣做出的贡献可以说居功至伟，自去年八月下旬推出后，至今狂揽近十万颗 Stars，足以说明社区用户对它的认同和感激，它就是 AUTOMATIC1111/stable-diffusion-webui。
目前项目中已累计收到了来自全球 465 位开源玩家的代码提交，积累了 5300+ 提交记录，项目代码库开始变的非常庞大、许多功能的设计和运行原理也变的不是那么的清晰。
在最近顺带研究 SDXL 和适配 SDXL 和 WebUI 的过程中，正好要走读代码，就从一个大家可能经常看到，但是大概率不清楚项目中调用流程和作用原理的功能开始吧。
在之前的内容中，我提起过多次 Stable Diffusion 相关的内容，如果你对它还不了解，可以自行翻阅。
VAE 相关的前置知识点
Stable Diffusion 和 VAE 涉及非常非常多有趣的知识点。
不过，这篇文章中，我不想过多展开学术相关的内容。所以，我会尽量简单的列举和 VAE 相关的，我们需要知道的内容，方便后续我们阅读代码中相关的部分和进行理解。
VAE 是什么？为什么需要它？
Stable Diffusion（稳定扩散）基于 Latent Diffusion（潜在扩展）提出，最出名的一篇论文：《High-Resolution Image Synthesis with Latent Diffusion Models》。使用的神经网络架构是基于《U-Net: Convolutional Networks for Biomedical Image Segmentation》这篇研究工作中提出的 U-Net。
简单来说是通过生成一副用户指定大小的充满随机噪声的图片（可以想象为充满雪花点的图），通过上面的算法，将噪音中的元素（蓝天、建筑物、草地）一步步的还原成我们想要的内容，并进行组装，最终成为图片（晴天广阔的草地上有一座建筑物）。
这个图片还原过程中，需要大量的重复、连续的计算，所以通常需要花费大量的时间，并且需要消耗大量的内存或显存资源。图片像素内容越丰富、图片尺寸越大，所需要的计算资源就越多、甚至非常夸张。譬如，如果不采取任何算法手段优化性能，想要生成一张 512×512 的图片，我们可能需要一台至少有 128G 显存的设备。
Latent Diffusion 有一个“多快好省”的妙用：可以通过在低维空间上进行高效的计算，来替代在真实空间对实际图片像素进行计算，减少大量的计算资源和内存资源的消耗。通俗理解，它可以对我们要生成的图像内容进行压缩和解压缩，压缩后的数据，还可以参与计算。譬如，当我们的“压缩工具”的压缩倍数是 8 的话，上面例子中，我们需要的显存资源就能够从 128G 减少到 16G 啦。
VAE 就是这个神奇的“压缩/解压缩工具”。
当我们想使⽤ Stable Diffusion ⽣成⼀张图⽚时（推理过程），VAE 起到的是解压缩过程，⼀ 般称之为“解码器”；当我们想将各种图⽚训练为图⽚⽣成模型时（训练过程），VAE 充当的 ⾓⾊就是编码器了，将图⽚或图⽚中的元素转换为低维空间向量表⽰，并传球给上⽂中提到的 U-Net ⽹络使⽤。
常见 VAE 模型类型
一般情况下，我们只需要重点关注 Stability AI 推出的 EMA （Exponential Moving Average）和 MSE （Mean Square Error ）两个类型的 VAE 模型即可。
- stabilityai/sd-vae-ft-ema
- stabilityai/sd-vae-ft-mse
在上面的链接中有这两个模型在辅助生成图片时的效果对比。就使用经验而言，EMA 会更锐利、MSE 会更平滑。
除此之外，还有两个比较知名的 VAE 模型，主要用在动漫风格的图片生成中：
- WarriorMama777/OrangeMixs
- hakurei/waifu-diffusion-v1-4
除了上面的几种 VAE 模型之外，有一些模型会自带自己的 VAE 模型，比如最近发布的 SDXL 模型，在项目中，我们能够看到模型自己的 VAE 模型。
- stabilityai/stable-diffusion-xl-refiner-1.0/vae
- stabilityai/stable-diffusion-xl-base-1.0/vae
当然，Stability AI 也单独发布了一个名为 stabilityai/sdxl-vae 的项目，虽然项目更新时间比两个 SDXL 绘图模型晚一天，但是其中的 VAE模型的版本，却比绘图模型中内置的 VAE 模型要老一个版本，推测这里应该是一个乌龙。
此外，在一些陈旧的文章中，可能会指引你使用：
https://huggingface.co/stabilityai/sd-vae-ft-ema-original/resolve/main/vae-ft-ema-560000-ema-pruned.ckpt
https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt
但其实，这两个项目早已被新版本模型取代，我们只需要使用前面提到的 EMA 和 MSE 仓库中的模型即可。
VAE 可以组合工作吗？
结合前文中，我们不难理解 UNET 这种接近于原始数据的内容，是比较好组合（合并）在一起进行计算的，而将 VAE 进行组合（合并）就不一定了，比如数据分别被压缩成了 rar 和 7z 格式，在解压缩的时候，我们需要分别进行这两种格式的数据处理。
VAE 的工作方式类似，如果我们需要套用多种模型，那么可选的操作是针对每一个模型应用和它“兼容”的 VAE 模型，而不能够将模型合并成一个，或者使用一个 VAE 模型“一条路走到黑”。
为什么有 VAE 可以“美颜” 的说法
在 Stable Diffusion 的世界，修复人脸主要依赖的是下面两个项目的能力：
- TencentARC/GFPGAN
- sczhou/CodeFormer
不过，前文提到的 Stability AI 推出的常用的 VAE 模型，是基于 LAION-Aesthetics和 LAION-Humans，对 CompVis/latent-diffusion 项目进行了模型微调而来的模型。而这两个数据集特别针对人对于图片的喜爱程度进行了整理，其中后者包含大量的人脸。
所以，在经过高质量的图片、大量人脸数据的训练后，VAE 模型对于改善图片色调，以及轻微修正图片中的人脸，也具备了一些能力。
好了，前置知识到此为止，已经足够我们阅读和理解代码啦。
Stable Diffusion WebUI 代码走读
VAE 的引用和定义散落在项目的许多地方，我们就主要场景进行展开。
VAE 模型和算法调用方式
WebUI 项目中涉及 VAE 定义主要有三个文件：
sd_vae.py获取存在的 VAE 模型文件列表、管理 VAE 模型的加载，文件位于：modules/sd_vae.py。
其中最重要的两个函数方法是 resolve_vae 和 reload_vae_weights。前者会将程序启动或用户手动触发 refresh_vae_list 函数后，寻找匹配绘图模型的 VAE 模型。后者，则会将匹配的模型在合适的时间点，一起装载到内存或显存中，用于图片生成计算。
sd_vae_taesd.py 使用了项目 madebyollin/taesd/ 作为 VAE 实现，核心目标是快速呈现图像，所以会损失一些质量细节，出现“捏造”数据的问题：modules/sd_vae_taesd.py。
sd_vae_approx.py 定义了核心类 VAE Approx，使用 8 个卷积层来实现图片的降噪和还原，以及 cheap_approximation，用来将潜在向量映射到 RGB 空间，这个算法参考了帖子中的做法，效率得到了巨大的提升：modules/sd_vae_approx.py。
如果我们使用的是 SDXL 模型，则会从 releases/v1.0.0-pre/ 发布页面中下载预构建模型 vaeapprox-sdxl.pt，反之则使用项目中的 model.pt 模型。
WebUI 启动，如果缺少上述模型，则会报一些因为读取不到文件出现的奇奇怪怪的问题，所以建议自行下载，提前放在项目代码要读取的位置。
VAE 模型文件查找逻辑
网上经常有人在项目开源社区或各种教程帖子中询问正确的 VAE 模型加载路径。其实这个问题很好回答，源码中查找模型路径的定义在这里：
def refresh_vae_list():
...
paths = [
os.path.join(sd_models.model_path, '**/*.vae.ckpt'),
os.path.join(sd_models.model_path, '**/*.vae.pt'),
os.path.join(sd_models.model_path, '**/*.vae.safetensors'),
os.path.join(vae_path, '**/*.ckpt'),
os.path.join(vae_path, '**/*.pt'),
os.path.join(vae_path, '**/*.safetensors'),
]
if shared.cmd_opts.ckpt_dir is not None and os.path.isdir(shared.cmd_opts.ckpt_dir):
paths += [
os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.ckpt'),
os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.pt'),
os.path.join(shared.cmd_opts.ckpt_dir, '**/*.vae.safetensors'),
]
if shared.cmd_opts.vae_dir is not None and os.path.isdir(shared.cmd_opts.vae_dir):
paths += [
os.path.join(shared.cmd_opts.vae_dir, '**/*.ckpt'),
os.path.join(shared.cmd_opts.vae_dir, '**/*.pt'),
os.path.join(shared.cmd_opts.vae_dir, '**/*.safetensors'),
]
...
for path in paths:
candidates += glob.iglob(path, recursive=True)
...
程序启动后，会首先寻找 models/具体模型目录 中的 *.vae.ckpt、*.vae.pt、*.vae.safetensors 三种后缀的 VAE 模型。接着，如果我们指定了应用启动的模型目录，则会追加查找 指定模型目录 下的三种后缀的 VAE 模型。最后，会搜索 models/VAE 目录下，三种后缀的 VAE 模型。
上述所有的操作，都支持子目录查找。
我们只需要将模型文件防止在上述任一位置即可，程序都能正常加载使用。
额外需要注意的是，有的帖子会让用户随意修改 VAE 模型后缀类型，比如将 pt 修改为 safetensors，这样做是错的。
VAE 模块的调用逻辑
在modules/sd_samplers_common.py程序中，定义了四种模型的加载方式：
approximation_indexes = {"Full": 0, "Approx NN": 1, "Approx cheap": 2, "TAESD": 3}
def single_sample_to_image(sample, approximation=None):
if approximation is None:
approximation = approximation_indexes.get(opts.show_progress_type, 0)
if approximation == 2:
x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5
elif approximation == 1:
x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach() * 0.5 + 0.5
elif approximation == 3:
x_sample = sample * 1.5
x_sample = sd_vae_taesd.model()(x_sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
else:
x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0] * 0.5 + 0.5
x_sample = torch.clamp(x_sample, min=0.0, max=1.0)
x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
x_sample = x_sample.astype(np.uint8)
return Image.fromarray(x_sample)
调用上面函数的主要方法有两个，目标都是将采样数据转换为一张或者多张图片：
def sample_to_image(samples, index=0, approximation=None):
return single_sample_to_image(samples[index], approximation)
def samples_to_image_grid(samples, approximation=None):
return images.image_grid([single_sample_to_image(sample, approximation) for sample in samples])
在模型处理图片的过程中的各种功能，会根据具体情况设置 approximation ，来调用不同的模型和算法，来生成图片。
全局参数定义
WebUI 项目中，和 VAE 相关最重要的参数最重要的只有一个：dtype_vae。
这个变量最初被定义在modules/devices.py：
dtype_vae = torch.float16
当程序代码跑起来之后，modules/shared.py 中会根据当前显卡支持的数据类型，自动调整数值：
devices.dtype = torch.float32 if cmd_opts.no_half else torch.float16
devices.dtype_vae = torch.float32 if cmd_opts.no_half or cmd_opts.no_half_vae else torch.float16
这个数值的改变，正常情况只能够通过命令行参数 no_half_vae 和 no_half 来改变。但也有意外，在 modules/processing.py 中：
def decode_latent_batch(model, batch, target_device=None, check_for_nans=False):
...
for i in range(batch.shape[0]):
sample = decode_first_stage(model, batch[i:i + 1])[0]
if check_for_nans:
try:
devices.test_for_nans(sample, "vae")
except devices.NansException as e:
if devices.dtype_vae == torch.float32 or not shared.opts.auto_vae_precision:
raise e
errors.print_error_explanation(
"A tensor with all NaNs was produced in VAE.\n"
"Web UI will now convert VAE into 32-bit float and retry.\n"
"To disable this behavior, disable the 'Automaticlly revert VAE to 32-bit floats' setting.\n"
"To always start with 32-bit VAE, use --no-half-vae commandline flag."
)
devices.dtype_vae = torch.float32
model.first_stage_model.to(devices.dtype_vae)
batch = batch.to(devices.dtype_vae)
...
return samples
在批量处理潜在向量、检查张量，用来“还原”图片时，假如我们指定了 VAE 模型的精度不为 auto_vae_precision，且 dtype_vae 为全精度时，则强制指定 devices.dtype_vae = torch.float32 并输出错误提醒告知用户。
通过 Model 合并对 VAE 进行 “Bake”
在项目的modules/extras.py程序中，当选择将 VAE Bake 到模型内时，将会把绘图模型和 VAE 打包成一个文件：
...
bake_in_vae_filename = sd_vae.vae_dict.get(bake_in_vae, None)
if bake_in_vae_filename is not None:
print(f"Baking in VAE from {bake_in_vae_filename}")
shared.state.textinfo = 'Baking in VAE'
vae_dict = sd_vae.load_vae_dict(bake_in_vae_filename, map_location='cpu')
for key in vae_dict.keys():
theta_0_key = 'first_stage_model.' + key
if theta_0_key in theta_0:
theta_0[theta_0_key] = to_half(vae_dict[key], save_as_half)
...
_, extension = os.path.splitext(output_modelname)
if extension.lower() == ".safetensors":
safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata)
else:
torch.save(theta_0, output_modelname)
性能优化，量化计算相关
项目的内置“超分插件 LDSR”中，包含了 VAE 量化相关的功能，extensions-builtin/LDSR/vqvae_quantize.py。这段代码改写自“官方项目”CompVis/taming-transformers/taming/modules/vqvae/quantize.py，但是做了大量的精简。
原始程序中，定义了四个量化计算的“工具人”，包括：
- VectorQuantizer类：对原始向量进行离散化处理，解决计算瓶颈。
- GumbelQuantize类：使用Gumbel Softmax 重参数化技巧处理向量，将原始向量转换为概率分布。
- VectorQuantizer2类：更高效的VectorQuantizer 的实现，避免使用昂贵的矩阵乘法操作。
- EmbeddingEMA类：使用 EMA（Exponential Moving Average）的方式对数据进行平滑。
新的实现中，实现了一个更精简的模块，可以在保留输入向量梯度信息的前提下，对原始向量进行高效的离散处理。
macOS 和小显存设备相关
WebUI 默认没有全局启用半精度，但是考虑到 macOS 设备资源有限，所以在默认提供的启动脚本中，尤其是针对 macOS 设备，默认添加了 --no-half-vae 启动参数。
export COMMANDLINE_ARGS="--skip-torch-cuda-test --upcast-sampling --no-half-vae --use-cpu interrogate"
所以，如果你的设备显存比较少，也可以考虑默认添加这个启动参数，来让资源有限的设备能够将模型程序跑起来。
其他
感谢在 WebUI 项目发布的同期，HuggingFace 团队发布了一篇详尽的技术博客：《Stable Diffusion with 🧨 Diffusers》。
感谢 @史业民 博士，中文 Llama2 7B 项目发起人，针对本文学术内容部分进行勘误。
最后
折腾完前两篇文章《使用 Docker 快速上手 Stability AI 的 SDXL 1.0 正式版》、《基于 Docker 的深度学习环境：Windows 篇》后，昨天尝试将 SDXL 1.0 适配到了本文中提到的 WebUI 项目。
虽然能够正常生成图片，但是感觉执行效率相比之前文章中的方法，有 20% 左右的性能下降，为了清清楚楚的使用项目，于是就有了这篇内容。
接下来有空的时候，我会继续阅读这个项目的代码，改进模型适配，期待在这个项目中，图片生成效率能够变的更高。
–EOF
本文使用「署名 4.0 国际 (CC BY 4.0)」许可协议，欢迎转载、或重新修改使用，但需要注明来源。 署名 4.0 国际 (CC BY 4.0)
本文作者: 苏洋
创建时间: 2023年07月30日
统计字数: 11485字
阅读时间: 23分钟阅读
本文链接: https://soulteary.com/2023/07/30/stable-diffusion-hardcore-survival-guide-vae-in-webui.html
文章出处登录后可见！
