LightVAE: WanとH3向けの高速な動画VAEデコード

ComfyUI Wikinews

ModelTCがLightVAEを公開。オリジナルの潜在インターフェースを維持したままWan2.1、Wan2.2、MiniMax-H3の動画VAEデコーダを圧縮し、ComfyUIで最大2.94倍高速化します。

LightX2Vチームによる<strong>LightVAE</strong>は、時間方向の低ランク畳み込みとブロックプルーニングによって事前学習済みの動画VAEデコーダを圧縮しつつ、オリジナルの潜在インターフェースを維持するため、オリジナルのエンコーダもそのまま動作します。Wan2.1、Wan2.2、MiniMax-H3に対応し、Hugging Faceで6つのチェックポイントとして配布され、ComfyUIノード、すぐに使える6つの生成ワークフロー、実測のデコードベンチマークも付属しています。
Wan2.1 VAE、Turbo-VAED、Flash-VAED、LightVAEのアーキテクチャ比較

プロジェクトリポジトリより: Wan2.1 VAEと、Turbo-VAED、Flash-VAED、LightVAEの比較。

LightVAEが行うこと

動画VAEのデコードは動画生成のコストの大きな部分を占めます。潜在クリップは毎回数百フレームに戻す必要があるためです。LightVAEはデコーダの役割をそのままに、それを実行するネットワークを小さくします。

  • LightVAE-Pro はオリジナルのデコーダと同じ深さで時間方向の低ランク畳み込みを使用するため、再構成の忠実度はオリジナルに近いままです。
  • LightVAE-Lite はさらにブロックプルーニングを追加し、わずかな忠実度と引き換えにさらなる高速化を実現します。
  • オリジナルの潜在インターフェース。 Wanモデルは独自の非圧縮エンコーダを維持するため、既存のワークフローはデコード側を差し替えるだけで済みます。MiniMax-H3にはプルーニング済みデコーダと、別途Liteエンコーダが用意されています。
表1: デコーダのスループットと再構成品質

論文の表1: デコーダのスループットと再構成PSNRの比較。

デコード速度

著者らがComfyUI上、H200で768 x 1344、バッチサイズ1、Wanは125フレーム、H3は124フレーム、ウォームアップ1回と3回の実行の中央値で測定した結果です。

バックボーンオリジナルVAELightVAE-ProLightVAE-Lite
Wan2.16.844秒 / 18.26 FPS3.754秒 / 33.30 FPS (1.82倍)2.330秒 / 53.64 FPS (2.94倍)
Wan2.26.886秒 / 18.15 FPS4.420秒 / 28.28 FPS (1.56倍)2.999秒 / 41.68 FPS (2.30倍)
MiniMax-H36.303秒 / 19.67 FPS提供なし3.574秒 / 34.69 FPS (1.76倍)

この測定はGPU同期されたVAE.decodeのみを対象としているため、生成全体ではなくデコードの高速化です。LightVAEは拡散モデルやテキストエンコーダのメモリコストを削減するものではありません。

チェックポイントの選択

生成モデルLightVAEデコーダエンコーダ
Wan2.1 / Wan2.2 14Blightvae-{pro,lite}-wan21-decoder-comfy.safetensorsオリジナルのWan2.1 VAE
Wan2.2 5B TI2Vlightvae-{pro,lite}-wan22-decoder-comfy.safetensorsオリジナルのWan2.2 VAE
MiniMax-H3lightvae-lite-h3-decoder-comfy.safetensorslightvae-lite-h3-encoder-comfy.safetensors またはオリジナルのH3動画VAE

Wan2.2 14BはWan2.1 VAE(16潜在チャンネル)を使用し、5BモデルはWan2.2 VAE(48チャンネル)を使用するため、自由に選べるわけではありません。WanのエンコーダとH3の音声VAEは、どの構成でも変更されません。

ComfyUIでの使用

python comfyui/install.py --comfy-root /path/to/ComfyUIでアダプタがComfyUI/custom_nodes/ComfyUI-LightVAE/にコピーされ、*-comfy.safetensorsの重みはComfyUI/models/vae/LightVAE/に配置します。ノードはLoad LightVAE (TLR)で、重要なのは、パッケージ化されたエクスポートがスタンダードなVAELoaderの経路ではなく、通常のLoad VAEを置き換えるという点です。

  1. オリジナルのLoad VAEノードを残し、LightVAEのbase_vaeに接続します。
  2. 使用するバックボーンに応じてProまたはLiteデコーダを選択し、encoder = noneのままにします(Wan)。またはH3 Liteの両方のファイルを選択します。
  3. LightVAEのVAE出力を、ワークフロー内のエンコードおよびデコードノードに接続します。

Wanではデコード側のみが変更され、参照条件付けはオリジナルのエンコーダを使い続けます。H3ではオプションのLiteエンコーダが参照エンコード経路とデコードの両方を置き換え、その構成ではオリジナルの動画VAEは不要です。両方の選択をnoneに設定すると、オリジナルのVAEが正確に復元されます。

リポジトリにはこのノード用の6つのフラットなワークフローが同梱されており、ComfyUI本体以外に必要なものはありません。

ワークフロータスクLightVAEコンポーネント
wan21_t2v_lightvae.jsonWan2.1 1.3B テキストから動画、832x480、81フレームwan21 Pro / Liteデコーダ
wan22_5b_t2v_lightvae.jsonWan2.2 5B テキストから動画、1280x704、121フレームwan22 Pro / Liteデコーダ
wan22_5b_i2v_lightvae.jsonWan2.2 5B 画像から動画wan22 Pro / Liteデコーダ
wan22_14b_i2v_lightvae.jsonWan2.2 14B 画像から動画、2エキスパートwan21 Pro / Liteデコーダ
h3_ref_image_lightvae.jsonH3 参照画像から動画と音声、1344x768H3 Liteエンコーダ + デコーダ
h3_ref_video_lightvae.jsonH3 参照動画とサウンドトラックから動画と音声H3 Liteエンコーダ + デコーダ

6つすべてはworkflowsディレクトリにもあります。

ハードウェアサポート

LightVAEはHopperに限定されません。NVIDIA BF16でTritonが利用可能な場合は融合カーネルを使用します。動作するTritonコンパイラがない場合やその他のPyTorchデバイスでは、TLR畳み込みは因数分解されたままPyTorchフォールバックで実行され、デバッグ用のCPU FP32経路も用意されています。H3の経路はWanのTritonカーネルではなくPyTorchの畳み込みとアテンションを使用しており、著者らは未検証のGPUファミリーでは速度を保証しないと述べています。このノードはPython 3.11のComfyUI e8665caで検証され、H3の完全な生成が1344x768、362フレーム、24 FPSで確認されています。

入手方法

重み: Hugging Faceのlightx2v/LightVAE。コード、ベンチマーク、ワークフロー: GitHubのModelTC/LightVAE。

コメント

GitHubでサインインしてディスカッションに参加しましょう。

コメントを読み込み中…
LightVAE: WanとH3向けの高速な動画VAEデコード | ComfyUI Wiki