使用推理端點進行解碼的遠端 VAE 🤗

釋出於 2025 年 2 月 24 日
在 GitHub 上更新

當使用潛在空間擴散模型進行高解析度影像和影片合成時,VAE 解碼器可能會消耗相當多的記憶體。這使得使用者很難在消費級 GPU 上執行這些模型,而無需犧牲延遲等。

例如,在解除安裝時,會產生裝置傳輸開銷,從而導致整體推理延遲。平鋪是另一種解決方案,它允許我們對所謂的“輸入塊”進行操作。然而,它可能對最終影像的質量產生負面影響。

因此,我們希望與社群共同試行一個想法——將解碼過程委託給遠端端點。

不儲存或跟蹤任何資料,並且程式碼是開源的。我們對 huggingface-inference-toolkit 進行了一些更改,並使用 自定義處理程式

此實驗功能由 Diffusers 🧨 開發

目錄:

開始使用

下面,我們涵蓋了我們認為遠端 VAE 推理會受益的三個用例。

程式碼

首先,我們建立了一個用於與遠端 VAE 互動的輔助方法。

從 `main` 安裝 `diffusers` 以執行程式碼。 `pip install git+https://github.com/huggingface/diffusers@main`

程式碼

from diffusers.utils.remote_utils import remote_decode

基本示例

這裡,我們展示瞭如何在隨機張量上使用遠端 VAE。

程式碼

image = remote_decode(
    endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
    tensor=torch.randn([1, 4, 64, 64], dtype=torch.float16),
    scaling_factor=0.18215,
)

Flux 的用法略有不同。Flux 潛在向量是打包的,因此我們需要傳送 `height` 和 `width`。

程式碼

image = remote_decode(
    endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/",
    tensor=torch.randn([1, 4096, 64], dtype=torch.float16),
    height=1024,
    width=1024,
    scaling_factor=0.3611,
    shift_factor=0.1159,
)

最後,是 HunyuanVideo 的一個示例。

程式碼

video = remote_decode(
    endpoint="https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/",
    tensor=torch.randn([1, 16, 3, 40, 64], dtype=torch.float16),
    output_type="mp4",
)
with open("video.mp4", "wb") as f:
    f.write(video)

生成

但我們希望在實際管道上使用 VAE 來獲得實際影像,而不是隨機噪聲。下面的示例展示瞭如何使用 SD v1.5 來實現。

程式碼

from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    variant="fp16",
    vae=None,
).to("cuda")

prompt = "Strawberry ice cream, in a stylish modern glass, coconut, splashing milk cream and honey, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious"

latent = pipe(
    prompt=prompt,
    output_type="latent",
).images
image = remote_decode(
    endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
    tensor=latent,
    scaling_factor=0.18215,
)
image.save("test.jpg")

這是 Flux 的另一個例子。

程式碼

from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
    vae=None,
).to("cuda")

prompt = "Strawberry ice cream, in a stylish modern glass, coconut, splashing milk cream and honey, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious"

latent = pipe(
    prompt=prompt,
    guidance_scale=0.0,
    num_inference_steps=4,
    output_type="latent",
).images
image = remote_decode(
    endpoint="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/",
    tensor=latent,
    height=1024,
    width=1024,
    scaling_factor=0.3611,
    shift_factor=0.1159,
)
image.save("test.jpg")

這是 HunyuanVideo 的一個例子。

程式碼

from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel

model_id = "hunyuanvideo-community/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id, subfolder="transformer", torch_dtype=torch.bfloat16
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id, transformer=transformer, vae=None, torch_dtype=torch.float16
).to("cuda")

latent = pipe(
    prompt="A cat walks on the grass, realistic",
    height=320,
    width=512,
    num_frames=61,
    num_inference_steps=30,
    output_type="latent",
).frames

video = remote_decode(
    endpoint="https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/",
    tensor=latent,
    output_type="mp4",
)

if isinstance(video, bytes):
    with open("video.mp4", "wb") as f:
        f.write(video)

排隊

使用遠端 VAE 的一大優點是我們可以將多個生成請求排隊。噹噹前的潛在向量正在進行解碼處理時,我們已經可以排隊另一個請求。這有助於提高併發性。

程式碼

import queue
import threading
from IPython.display import display
from diffusers import StableDiffusionPipeline

def decode_worker(q: queue.Queue):
    while True:
        item = q.get()
        if item is None:
            break
        image = remote_decode(
            endpoint="https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
            tensor=item,
            scaling_factor=0.18215,
        )
        display(image)
        q.task_done()

q = queue.Queue()
thread = threading.Thread(target=decode_worker, args=(q,), daemon=True)
thread.start()

def decode(latent: torch.Tensor):
    q.put(latent)

prompts = [
    "Blueberry ice cream, in a stylish modern glass , ice cubes, nuts, mint leaves, splashing milk cream, in a gradient purple background, fluid motion, dynamic movement, cinematic lighting, Mysterious",
    "Lemonade in a glass, mint leaves, in an aqua and white background, flowers, ice cubes, halo, fluid motion, dynamic movement, soft lighting, digital painting, rule of thirds composition, Art by Greg rutkowski, Coby whitmore",
    "Comic book art, beautiful, vintage, pastel neon colors, extremely detailed pupils, delicate features, light on face, slight smile, Artgerm, Mary Blair, Edmund Dulac, long dark locks, bangs, glowing, fashionable style, fairytale ambience, hot pink.",
    "Masterpiece, vanilla cone ice cream garnished with chocolate syrup, crushed nuts, choco flakes, in a brown background, gold, cinematic lighting, Art by WLOP",
    "A bowl of milk, falling cornflakes, berries, blueberries, in a white background, soft lighting, intricate details, rule of thirds, octane render, volumetric lighting",
    "Cold Coffee with cream, crushed almonds, in a glass, choco flakes, ice cubes, wet, in a wooden background, cinematic lighting, hyper realistic painting, art by Carne Griffiths, octane render, volumetric lighting, fluid motion, dynamic movement, muted colors,",
]

pipe = StableDiffusionPipeline.from_pretrained(
    "Lykon/dreamshaper-8",
    torch_dtype=torch.float16,
    vae=None,
).to("cuda")

pipe.unet = pipe.unet.to(memory_format=torch.channels_last)
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

_ = pipe(
    prompt=prompts[0],
    output_type="latent",
)

for prompt in prompts:
    latent = pipe(
        prompt=prompt,
        output_type="latent",
    ).images
    decode(latent)

q.put(None)
thread.join()

可用 VAE

使用遠端 VAE 的優勢

這些表格展示了不同 GPU 的 VRAM 要求。記憶體使用百分比決定了某些 GPU 的使用者是否需要解除安裝。解除安裝時間因 CPU、RAM 和 HDD/NVMe 而異。平鋪解碼會增加推理時間。

SD v1.5

GPU 解析度 時間(秒) 記憶體 (%) 平鋪時間(秒) 平鋪記憶體 (%)
英偉達 GeForce RTX 4090 512x512 0.031 5.60% 0.031 (0%) 5.60%
英偉達 GeForce RTX 4090 1024x1024 0.148 20.00% 0.301 (+103%) 5.60%
英偉達 GeForce RTX 4080 512x512 0.05 8.40% 0.050 (0%) 8.40%
英偉達 GeForce RTX 4080 1024x1024 0.224 30.00% 0.356 (+59%) 8.40%
英偉達 GeForce RTX 4070 Ti 512x512 0.066 11.30% 0.066 (0%) 11.30%
英偉達 GeForce RTX 4070 Ti 1024x1024 0.284 40.50% 0.454 (+60%) 11.40%
英偉達 GeForce RTX 3090 512x512 0.062 5.20% 0.062 (0%) 5.20%
英偉達 GeForce RTX 3090 1024x1024 0.253 18.50% 0.464 (+83%) 5.20%
英偉達 GeForce RTX 3080 512x512 0.07 12.80% 0.070 (0%) 12.80%
英偉達 GeForce RTX 3080 1024x1024 0.286 45.30% 0.466 (+63%) 12.90%
英偉達 GeForce RTX 3070 512x512 0.102 15.90% 0.102 (0%) 15.90%
英偉達 GeForce RTX 3070 1024x1024 0.421 56.30% 0.746 (+77%) 16.00%

SDXL

GPU 解析度 時間(秒) 記憶體消耗 (%) 平鋪時間(秒) 平鋪記憶體 (%)
英偉達 GeForce RTX 4090 512x512 0.057 10.00% 0.057 (0%) 10.00%
英偉達 GeForce RTX 4090 1024x1024 0.256 35.50% 0.257 (+0.4%) 35.50%
英偉達 GeForce RTX 4080 512x512 0.092 15.00% 0.092 (0%) 15.00%
英偉達 GeForce RTX 4080 1024x1024 0.406 53.30% 0.406 (0%) 53.30%
英偉達 GeForce RTX 4070 Ti 512x512 0.121 20.20% 0.120 (-0.8%) 20.20%
英偉達 GeForce RTX 4070 Ti 1024x1024 0.519 72.00% 0.519 (0%) 72.00%
英偉達 GeForce RTX 3090 512x512 0.107 10.50% 0.107 (0%) 10.50%
英偉達 GeForce RTX 3090 1024x1024 0.459 38.00% 0.460 (+0.2%) 38.00%
英偉達 GeForce RTX 3080 512x512 0.121 25.60% 0.121 (0%) 25.60%
英偉達 GeForce RTX 3080 1024x1024 0.524 93.00% 0.524 (0%) 93.00%
英偉達 GeForce RTX 3070 512x512 0.183 31.80% 0.183 (0%) 31.80%
英偉達 GeForce RTX 3070 1024x1024 0.794 96.40% 0.794 (0%) 96.40%

提供反饋

如果您喜歡這個想法和功能,請幫助我們提供反饋,說明我們如何做得更好,以及您是否有興趣將此類功能更原生整合到 Hugging Face 生態系統中。如果這個試點進展順利,我們計劃為更多模型建立最佳化的 VAE 端點,包括那些可以生成高解析度影片的模型!

步驟:

  1. 透過 此連結 在 Diffusers 上提出問題。
  2. 回答問題並提供您想要的任何額外資訊。
  3. 點選提交!

社群

這篇文章有些地方不夠清晰

  • “時間”是指“推理時間”還是“解除安裝時間”?
  • 平鋪記憶體/平鋪時間是什麼意思?
  • VAE 平鋪總是發生嗎?
·
文章作者

我不確定這有什麼不清楚的,但為了完整性,特此說明。

VAE 平鋪總是發生嗎?

我們從未提及這會發生。

平鋪記憶體/平鋪時間是什麼意思?

平鋪記憶體/時間意味著我們正在應用平鋪。

“時間”是指“推理時間”還是“解除安裝時間”?

總往返時間。如果它意味著其他任何東西,它就會像其他一樣被指定。

這需要專業賬戶嗎?

文章作者

不。

你好。
非常棒的實現,執行完美,解碼時間為 2-4 秒,具體取決於影像大小。

一個簡單的問題是,是否有可能在另一臺機器(非 HF)上建立本地端點並將其用作 vae-decode 機器?

例如,一個舒適的使用者介面實現。

文章作者

舒適實現:https://github.com/kijai/ComfyUI-HFRemoteVae

如果輸入和輸出方案匹配,您可以在本地機器上託管端點並按部落格文章中所示使用它。

·

好的,我再讀一遍。
謝謝。

📻 🎙️ 嘿,我為這篇博文生成了一個 AI 播客,快來聽聽看吧!

此播客透過 ngxson/kokoro-podcast-generator 生成,使用 DeepSeek-R1Kokoro-TTS

·
文章作者

無音訊。

非常棒的功能。

我們還能獲得 Wan Video 2.1 VAE 用於遠端解碼嗎?

註冊登入 發表評論

© . This site is unofficial and not affiliated with Hugging Face, Inc.