本文使用「署名 4.0 國(guó)際 (CC BY 4.0)」許可協(xié)議，歡迎轉(zhuǎn)載、或重新修改使用，但需要注明來(lái)源。 署名 4.0 國(guó)際 (CC BY 4.0)

本文作者: 蘇洋

創(chuàng)建時(shí)間: 2023年07月30日
統(tǒng)計(jì)字?jǐn)?shù): 11485字
閱讀時(shí)間: 23分鐘閱讀
本文鏈接: https://soulteary.com/2023/07/30/stable-diffusion-hardcore-survival-guide-vae-in-webui.html

Stable Diffusion 硬核生存指南：WebUI 中的 VAE

本篇文章聊聊 Stable Diffusion 生態(tài)中呼聲最高、也是最復(fù)雜的開(kāi)源模型管理圖形界面 “stable-diffusion-webui” 中和 VAE 相關(guān)的事情。

寫(xiě)在前面

Stable Diffusion 生態(tài)中有一個(gè)很重要的項(xiàng)目，它對(duì)于 SD 生態(tài)繁榮做出的貢獻(xiàn)可以說(shuō)居功至偉，自去年八月下旬推出后，至今狂攬近十萬(wàn)顆 Stars，足以說(shuō)明社區(qū)用戶對(duì)它的認(rèn)同和感激，它就是 AUTOMATIC1111/stable-diffusion-webui。

目前項(xiàng)目中已累計(jì)收到了來(lái)自全球 465 位開(kāi)源玩家的代碼提交，積累了 5300+ 提交記錄，項(xiàng)目代碼庫(kù)開(kāi)始變的非常龐大、許多功能的設(shè)計(jì)和運(yùn)行原理也變的不是那么的清晰。

在最近順帶研究 SDXL 和適配 SDXL 和 WebUI 的過(guò)程中，正好要走讀代碼，就從一個(gè)大家可能經(jīng)?？吹剑谴蟾怕什磺宄?xiàng)目中調(diào)用流程和作用原理的功能開(kāi)始吧。

在之前的內(nèi)容中，我提起過(guò)多次 Stable Diffusion 相關(guān)的內(nèi)容，如果你對(duì)它還不了解，可以自行翻閱。

VAE 相關(guān)的前置知識(shí)點(diǎn)

Stable Diffusion 和 VAE 涉及非常非常多有趣的知識(shí)點(diǎn)。

不過(guò)，這篇文章中，我不想過(guò)多展開(kāi)學(xué)術(shù)相關(guān)的內(nèi)容。所以，我會(huì)盡量簡(jiǎn)單的列舉和 VAE 相關(guān)的，我們需要知道的內(nèi)容，方便后續(xù)我們閱讀代碼中相關(guān)的部分和進(jìn)行理解。

VAE 是什么？為什么需要它？

Stable Diffusion（穩(wěn)定擴(kuò)散）基于 Latent Diffusion（潛在擴(kuò)展）提出，最出名的一篇論文：《High-Resolution Image Synthesis with Latent Diffusion Models》。使用的神經(jīng)網(wǎng)絡(luò)架構(gòu)是基于《U-Net: Convolutional Networks for Biomedical Image Segmentation》這篇研究工作中提出的 U-Net。

簡(jiǎn)單來(lái)說(shuō)是通過(guò)生成一副用戶指定大小的充滿隨機(jī)噪聲的圖片（可以想象為充滿雪花點(diǎn)的圖），通過(guò)上面的算法，將噪音中的元素（藍(lán)天、建筑物、草地）一步步的還原成我們想要的內(nèi)容，并進(jìn)行組裝，最終成為圖片（晴天廣闊的草地上有一座建筑物）。

這個(gè)圖片還原過(guò)程中，需要大量的重復(fù)、連續(xù)的計(jì)算，所以通常需要花費(fèi)大量的時(shí)間，并且需要消耗大量的內(nèi)存或顯存資源。圖片像素內(nèi)容越豐富、圖片尺寸越大，所需要的計(jì)算資源就越多、甚至非?？鋸?。譬如，如果不采取任何算法手段優(yōu)化性能，想要生成一張 512x512 的圖片，我們可能需要一臺(tái)至少有 128G 顯存的設(shè)備。

Latent Diffusion 有一個(gè)“多快好省”的妙用：可以通過(guò)在低維空間上進(jìn)行高效的計(jì)算，來(lái)替代在真實(shí)空間對(duì)實(shí)際圖片像素進(jìn)行計(jì)算，減少大量的計(jì)算資源和內(nèi)存資源的消耗。通俗理解，它可以對(duì)我們要生成的圖像內(nèi)容進(jìn)行壓縮和解壓縮，壓縮后的數(shù)據(jù)，還可以參與計(jì)算。譬如，當(dāng)我們的“壓縮工具”的壓縮倍數(shù)是 8 的話，上面例子中，我們需要的顯存資源就能夠從 128G 減少到 16G 啦。

VAE 就是這個(gè)神奇的“壓縮/解壓縮工具”。

當(dāng)我們想使? Stable Diffusion ?成?張圖?時(shí)（推理過(guò)程），VAE 起到的是解壓縮過(guò)程，? 般稱之為“解碼器”；當(dāng)我們想將各種圖?訓(xùn)練為圖??成模型時(shí)（訓(xùn)練過(guò)程），VAE 充當(dāng)?shù)???就是編碼器了，將圖?或圖?中的元素轉(zhuǎn)換為低維空間向量表?，并傳球給上?中提到的 U-Net ?絡(luò)使?。

常見(jiàn) VAE 模型類型

一般情況下，我們只需要重點(diǎn)關(guān)注 Stability AI 推出的 EMA （Exponential Moving Average）和 MSE （Mean Square Error ）兩個(gè)類型的 VAE 模型即可。

• stabilityai/sd-vae-ft-ema
• stabilityai/sd-vae-ft-mse

在上面的鏈接中有這兩個(gè)模型在輔助生成圖片時(shí)的效果對(duì)比。就使用經(jīng)驗(yàn)而言，EMA 會(huì)更銳利、MSE 會(huì)更平滑。

除此之外，還有兩個(gè)比較知名的 VAE 模型，主要用在動(dòng)漫風(fēng)格的圖片生成中：

• WarriorMama777/OrangeMixs
• hakurei/waifu-diffusion-v1-4

除了上面的幾種 VAE 模型之外，有一些模型會(huì)自帶自己的 VAE 模型，比如最近發(fā)布的 SDXL 模型，在項(xiàng)目中，我們能夠看到模型自己的 VAE 模型。

• stabilityai/stable-diffusion-xl-refiner-1.0/vae
• stabilityai/stable-diffusion-xl-base-1.0/vae

當(dāng)然，Stability AI 也單獨(dú)發(fā)布了一個(gè)名為 stabilityai/sdxl-vae 的項(xiàng)目，雖然項(xiàng)目更新時(shí)間比兩個(gè) SDXL 繪圖模型晚一天，但是其中的 VAE模型的版本，卻比繪圖模型中內(nèi)置的 VAE 模型要老一個(gè)版本，推測(cè)這里應(yīng)該是一個(gè)烏龍。

此外，在一些陳舊的文章中，可能會(huì)指引你使用：

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

但其實(shí)，這兩個(gè)項(xiàng)目早已被新版本模型取代，我們只需要使用前面提到的 EMA 和 MSE 倉(cāng)庫(kù)中的模型即可。

VAE 可以組合工作嗎？

結(jié)合前文中，我們不難理解 UNET 這種接近于原始數(shù)據(jù)的內(nèi)容，是比較好組合（合并）在一起進(jìn)行計(jì)算的，而將 VAE 進(jìn)行組合（合并）就不一定了，比如數(shù)據(jù)分別被壓縮成了 rar 和 7z 格式，在解壓縮的時(shí)候，我們需要分別進(jìn)行這兩種格式的數(shù)據(jù)處理。

VAE 的工作方式類似，如果我們需要套用多種模型，那么可選的操作是針對(duì)每一個(gè)模型應(yīng)用和它“兼容”的 VAE 模型，而不能夠?qū)⒛Ｐ秃喜⒊梢粋€(gè)，或者使用一個(gè) VAE 模型“一條路走到黑”。

為什么有 VAE 可以“美顏” 的說(shuō)法

在 Stable Diffusion 的世界，修復(fù)人臉主要依賴的是下面兩個(gè)項(xiàng)目的能力：

• TencentARC/GFPGAN
• sczhou/CodeFormer

不過(guò)，前文提到的 Stability AI 推出的常用的 VAE 模型，是基于 LAION-Aesthetics和 LAION-Humans，對(duì) CompVis/latent-diffusion 項(xiàng)目進(jìn)行了模型微調(diào)而來(lái)的模型。而這兩個(gè)數(shù)據(jù)集特別針對(duì)人對(duì)于圖片的喜愛(ài)程度進(jìn)行了整理，其中后者包含大量的人臉。

所以，在經(jīng)過(guò)高質(zhì)量的圖片、大量人臉數(shù)據(jù)的訓(xùn)練后，VAE 模型對(duì)于改善圖片色調(diào)，以及輕微修正圖片中的人臉，也具備了一些能力。

好了，前置知識(shí)到此為止，已經(jīng)足夠我們閱讀和理解代碼啦。

Stable Diffusion WebUI 代碼走讀

VAE 的引用和定義散落在項(xiàng)目的許多地方，我們就主要場(chǎng)景進(jìn)行展開(kāi)。

VAE 模型和算法調(diào)用方式

WebUI 項(xiàng)目中涉及 VAE 定義主要有三個(gè)文件：

sd_vae.py獲取存在的 VAE 模型文件列表、管理 VAE 模型的加載，文件位于：modules/sd_vae.py。

其中最重要的兩個(gè)函數(shù)方法是 resolve_vae 和 reload_vae_weights。前者會(huì)將程序啟動(dòng)或用戶手動(dòng)觸發(fā) refresh_vae_list 函數(shù)后，尋找匹配繪圖模型的 VAE 模型。后者，則會(huì)將匹配的模型在合適的時(shí)間點(diǎn)，一起裝載到內(nèi)存或顯存中，用于圖片生成計(jì)算。

sd_vae_taesd.py 使用了項(xiàng)目 madebyollin/taesd/ 作為 VAE 實(shí)現(xiàn)，核心目標(biāo)是快速呈現(xiàn)圖像，所以會(huì)損失一些質(zhì)量細(xì)節(jié)，出現(xiàn)“捏造”數(shù)據(jù)的問(wèn)題：modules/sd_vae_taesd.py。

sd_vae_approx.py 定義了核心類 VAE Approx，使用 8 個(gè)卷積層來(lái)實(shí)現(xiàn)圖片的降噪和還原，以及 cheap_approximation，用來(lái)將潛在向量映射到 RGB 空間，這個(gè)算法參考了帖子中的做法，效率得到了巨大的提升：modules/sd_vae_approx.py。

如果我們使用的是 SDXL 模型，則會(huì)從 releases/v1.0.0-pre/ 發(fā)布頁(yè)面中下載預(yù)構(gòu)建模型 vaeapprox-sdxl.pt，反之則使用項(xiàng)目中的 model.pt 模型。

WebUI 啟動(dòng)，如果缺少上述模型，則會(huì)報(bào)一些因?yàn)樽x取不到文件出現(xiàn)的奇奇怪怪的問(wèn)題，所以建議自行下載，提前放在項(xiàng)目代碼要讀取的位置。

VAE 模型文件查找邏輯

網(wǎng)上經(jīng)常有人在項(xiàng)目開(kāi)源社區(qū)或各種教程帖子中詢問(wèn)正確的 VAE 模型加載路徑。其實(shí)這個(gè)問(wèn)題很好回答，源碼中查找模型路徑的定義在這里：

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)
...

程序啟動(dòng)后，會(huì)首先尋找 models/具體模型目錄 中的 *.vae.ckpt、*.vae.pt、*.vae.safetensors 三種后綴的 VAE 模型。接著，如果我們指定了應(yīng)用啟動(dòng)的模型目錄，則會(huì)追加查找 指定模型目錄 下的三種后綴的 VAE 模型。最后，會(huì)搜索 models/VAE 目錄下，三種后綴的 VAE 模型。

上述所有的操作，都支持子目錄查找。

我們只需要將模型文件防止在上述任一位置即可，程序都能正常加載使用。

額外需要注意的是，有的帖子會(huì)讓用戶隨意修改 VAE 模型后綴類型，比如將 pt 修改為 safetensors，這樣做是錯(cuò)的。

VAE 模塊的調(diào)用邏輯

在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)

調(diào)用上面函數(shù)的主要方法有兩個(gè)，目標(biāo)都是將采樣數(shù)據(jù)轉(zhuǎn)換為一張或者多張圖片：

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])

在模型處理圖片的過(guò)程中的各種功能，會(huì)根據(jù)具體情況設(shè)置 approximation ，來(lái)調(diào)用不同的模型和算法，來(lái)生成圖片。

全局參數(shù)定義

WebUI 項(xiàng)目中，和 VAE 相關(guān)最重要的參數(shù)最重要的只有一個(gè)：dtype_vae。

這個(gè)變量最初被定義在modules/devices.py：

dtype_vae = torch.float16

當(dāng)程序代碼跑起來(lái)之后，modules/shared.py 中會(huì)根據(jù)當(dāng)前顯卡支持的數(shù)據(jù)類型，自動(dòng)調(diào)整數(shù)值：

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

這個(gè)數(shù)值的改變，正常情況只能夠通過(guò)命令行參數(shù) no_half_vae 和 no_half 來(lái)改變。但也有意外，在 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

在批量處理潛在向量、檢查張量，用來(lái)“還原”圖片時(shí)，假如我們指定了 VAE 模型的精度不為 auto_vae_precision，且 dtype_vae 為全精度時(shí)，則強(qiáng)制指定 devices.dtype_vae = torch.float32 并輸出錯(cuò)誤提醒告知用戶。

通過(guò) Model 合并對(duì) VAE 進(jìn)行 “Bake”

在項(xiàng)目的modules/extras.py程序中，當(dāng)選擇將 VAE Bake 到模型內(nèi)時(shí)，將會(huì)把繪圖模型和 VAE 打包成一個(gè)文件：

...
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)

性能優(yōu)化，量化計(jì)算相關(guān)

項(xiàng)目的內(nèi)置“超分插件 LDSR”中，包含了 VAE 量化相關(guān)的功能，extensions-builtin/LDSR/vqvae_quantize.py。這段代碼改寫(xiě)自“官方項(xiàng)目”CompVis/taming-transformers/taming/modules/vqvae/quantize.py，但是做了大量的精簡(jiǎn)。

原始程序中，定義了四個(gè)量化計(jì)算的“工具人”，包括：

• VectorQuantizer類：對(duì)原始向量進(jìn)行離散化處理，解決計(jì)算瓶頸。
• GumbelQuantize類：使用Gumbel Softmax 重參數(shù)化技巧處理向量，將原始向量轉(zhuǎn)換為概率分布。
• VectorQuantizer2類：更高效的VectorQuantizer 的實(shí)現(xiàn)，避免使用昂貴的矩陣乘法操作。
• EmbeddingEMA類：使用 EMA（Exponential Moving Average）的方式對(duì)數(shù)據(jù)進(jìn)行平滑。

新的實(shí)現(xiàn)中，實(shí)現(xiàn)了一個(gè)更精簡(jiǎn)的模塊，可以在保留輸入向量梯度信息的前提下，對(duì)原始向量進(jìn)行高效的離散處理。

macOS 和小顯存設(shè)備相關(guān)

WebUI 默認(rèn)沒(méi)有全局啟用半精度，但是考慮到 macOS 設(shè)備資源有限，所以在默認(rèn)提供的啟動(dòng)腳本中，尤其是針對(duì) macOS 設(shè)備，默認(rèn)添加了 --no-half-vae 啟動(dòng)參數(shù)。

export COMMANDLINE_ARGS="--skip-torch-cuda-test --upcast-sampling --no-half-vae --use-cpu interrogate"

所以，如果你的設(shè)備顯存比較少，也可以考慮默認(rèn)添加這個(gè)啟動(dòng)參數(shù)，來(lái)讓資源有限的設(shè)備能夠?qū)⒛Ｐ统绦蚺芷饋?lái)。

其他

感謝在 WebUI 項(xiàng)目發(fā)布的同期，HuggingFace 團(tuán)隊(duì)發(fā)布了一篇詳盡的技術(shù)博客：《Stable Diffusion with ?? Diffusers》。

感謝 @史業(yè)民 博士，中文 Llama2 7B 項(xiàng)目發(fā)起人，針對(duì)本文學(xué)術(shù)內(nèi)容部分進(jìn)行勘誤。

最后

折騰完前兩篇文章《使用 Docker 快速上手 Stability AI 的 SDXL 1.0 正式版》、《基于 Docker 的深度學(xué)習(xí)環(huán)境：Windows 篇》后，昨天嘗試將 SDXL 1.0 適配到了本文中提到的 WebUI 項(xiàng)目。

雖然能夠正常生成圖片，但是感覺(jué)執(zhí)行效率相比之前文章中的方法，有 20% 左右的性能下降，為了清清楚楚的使用項(xiàng)目，于是就有了這篇內(nèi)容。

接下來(lái)有空的時(shí)候，我會(huì)繼續(xù)閱讀這個(gè)項(xiàng)目的代碼，改進(jìn)模型適配，期待在這個(gè)項(xiàng)目中，圖片生成效率能夠變的更高。

–EOF文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-625591.html

到了這里，關(guān)于Stable Diffusion 硬核生存指南：WebUI 中的 VAE的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！