本篇文章聊聊 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。

VAE 工作原理(图片来自 HF 技术博客)

简单来说是通过生成一副用户指定大小的充满随机噪声的图片(可以想象为充满雪花点的图),通过上面的算法,将噪音中的元素(蓝天、建筑物、草地)一步步的还原成我们想要的内容,并进行组装,最终成为图片(晴天广阔的草地上有一座建筑物)。

这个图片还原过程中,需要大量的重复、连续的计算,所以通常需要花费大量的时间,并且需要消耗大量的内存或显存资源。图片像素内容越丰富、图片尺寸越大,所需要的计算资源就越多、甚至非常夸张。譬如,如果不采取任何算法手段优化性能,想要生成一张 512x512 的图片,我们可能需要一台至少有 128G 显存的设备。

Latent Diffusion 有一个“多快好省”的妙用:可以通过在低维空间上进行高效的计算,来替代在真实空间对实际图片像素进行计算,减少大量的计算资源和内存资源的消耗。通俗理解,它可以对我们要生成的图像内容进行压缩和解压缩,压缩后的数据,还可以参与计算。譬如,当我们的“压缩工具”的压缩倍数是 8 的话,上面例子中,我们需要的显存资源就能够从 128G 减少到 16G 啦。

VAE 就是这个神奇的“压缩/解压缩工具”。

图片推理过程中,VAE 的“位置” (图片来自 HF 技术博客)

当我们想使⽤ Stable Diffusion ⽣成⼀张图⽚时(推理过程),VAE 起到的是解压缩过程,⼀ 般称之为“解码器”;当我们想将各种图⽚训练为图⽚⽣成模型时(训练过程),VAE 充当的 ⾓⾊就是编码器了,将图⽚或图⽚中的元素转换为低维空间向量表⽰,并传球给上⽂中提到的 U-Net ⽹络使⽤。

常见 VAE 模型类型

一般情况下,我们只需要重点关注 Stability AI 推出的 EMA (Exponential Moving Average)和 MSE (Mean Square Error )两个类型的 VAE 模型即可。

在上面的链接中有这两个模型在辅助生成图片时的效果对比。就使用经验而言,EMA 会更锐利、MSE 会更平滑。

除此之外,还有两个比较知名的 VAE 模型,主要用在动漫风格的图片生成中:

除了上面的几种 VAE 模型之外,有一些模型会自带自己的 VAE 模型,比如最近发布的 SDXL 模型,在项目中,我们能够看到模型自己的 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 进行组合(合并)就不一定了,比如数据分别被压缩成了 rar7z 格式,在解压缩的时候,我们需要分别进行这两种格式的数据处理。

VAE 的工作方式类似,如果我们需要套用多种模型,那么可选的操作是针对每一个模型应用和它“兼容”的 VAE 模型,而不能够将模型合并成一个,或者使用一个 VAE 模型“一条路走到黑”。

为什么有 VAE 可以“美颜” 的说法

在 Stable Diffusion 的世界,修复人脸主要依赖的是下面两个项目的能力:

不过,前文提到的 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_vaereload_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_vaeno_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 项目。

适配 SDXL 1.0 模型的 WebUI

虽然能够正常生成图片,但是感觉执行效率相比之前文章中的方法,有 20% 左右的性能下降,为了清清楚楚的使用项目,于是就有了这篇内容。

接下来有空的时候,我会继续阅读这个项目的代码,改进模型适配,期待在这个项目中,图片生成效率能够变的更高。

–EOF