LightVAE: ComfyUI에서 Wan 및 H3 비디오 VAE 디코딩 속도 향상

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: 재구성 PSNR에 따른 디코더 처리량.

디코드 속도

저자들이 H200에서 768 x 1344, 배치 크기 1, Wan은 125프레임, H3는 124프레임으로 ComfyUI에서 측정했으며, 워밍업 1회와 3회 실행의 중앙값을 사용했습니다.

백본	원본 VAE	LightVAE-Pro	LightVAE-Lite
Wan2.1	6.844 s / 18.26 FPS	3.754 s / 33.30 FPS (1.82배)	2.330 s / 53.64 FPS (2.94배)
Wan2.2	6.886 s / 18.15 FPS	4.420 s / 28.28 FPS (1.56배)	2.999 s / 41.68 FPS (2.30배)
MiniMax-H3	6.303 s / 19.67 FPS	제공되지 않음	3.574 s / 34.69 FPS (1.76배)

측정은 GPU 동기화된 VAE.decode만을 대상으로 하므로 전체 생성 속도 향상이 아니라 디코드 속도 향상입니다. LightVAE는 diffusion 모델이나 텍스트 인코더의 메모리 비용을 줄이지 않습니다.

체크포인트 선택

생성 모델	LightVAE 디코더	인코더
Wan2.1 / Wan2.2 14B	`lightvae-{pro,lite}-wan21-decoder-comfy.safetensors`	원본 Wan2.1 VAE
Wan2.2 5B TI2V	`lightvae-{pro,lite}-wan22-decoder-comfy.safetensors`	원본 Wan2.2 VAE
MiniMax-H3	`lightvae-lite-h3-decoder-comfy.safetensors`	`lightvae-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를 대체한다는 것입니다.

원본 Load VAE 노드를 유지하고 LightVAE의 base_vae에 연결합니다.
사용 중인 백본에 맞는 Pro 또는 Lite 디코더를 선택하고 encoder = none(Wan)으로 두거나, H3 Lite 파일 두 개를 모두 선택합니다.
LightVAE의 VAE 출력을 워크플로의 인코드 및 디코드 노드에 연결합니다.

Wan에서는 디코드 쪽만 변경되며 참조 조건화는 계속 원본 인코더를 사용합니다. H3에서는 선택적 Lite 인코더가 참조 인코딩 경로와 디코딩을 모두 대체하며, 이 구성에서는 원본 비디오 VAE가 필요하지 않습니다. 두 선택을 모두 none으로 설정하면 원본 VAE가 정확히 복원됩니다.

저장소에는 이 노드를 위한 단일 워크플로 6개가 포함되어 있으며, ComfyUI 코어 외에는 아무것도 필요하지 않습니다.

워크플로	작업	LightVAE 구성 요소
`wan21_t2v_lightvae.json`	Wan2.1 1.3B 텍스트 투 비디오, 832x480, 81프레임	wan21 Pro / Lite 디코더
`wan22_5b_t2v_lightvae.json`	Wan2.2 5B 텍스트 투 비디오, 1280x704, 121프레임	wan22 Pro / Lite 디코더
`wan22_5b_i2v_lightvae.json`	Wan2.2 5B 이미지 투 비디오	wan22 Pro / Lite 디코더
`wan22_14b_i2v_lightvae.json`	Wan2.2 14B 이미지 투 비디오, 전문가 2개	wan21 Pro / Lite 디코더
`h3_ref_image_lightvae.json`	H3 참조 이미지에서 비디오 및 오디오 생성, 1344x768	H3 Lite 인코더 + 디코더
`h3_ref_video_lightvae.json`	H3 참조 비디오와 사운드트랙에서 비디오 및 오디오 생성	H3 Lite 인코더 + 디코더

여섯 개 모두 워크플로 디렉터리에도 있습니다.

하드웨어 지원

LightVAE는 Hopper에 국한되지 않습니다. Triton을 사용하는 NVIDIA BF16에서는 융합 커널을 사용하고, 작동하는 Triton 컴파일러가 없거나 다른 PyTorch 장치에서는 TLR 컨볼루션이 분해된 형태로 유지되어 PyTorch 폴백을 통해 실행되며, 디버깅용 CPU FP32 경로도 있습니다. H3 경로는 Wan Triton 커널 대신 PyTorch 컨볼루션과 어텐션을 사용하며, 저자들은 테스트되지 않은 GPU 계열에 대해서는 속도를 보장하지 않는다고 밝혔습니다. 이 노드는 Python 3.11의 ComfyUI e8665ca에서 검증되었고, 1344x768, 362프레임, 24 FPS의 완전한 H3 생성이 확인되었습니다.

사용 가능 여부

가중치: Hugging Face의 lightx2v/LightVAE. 코드, 벤치마크, 워크플로: GitHub의 ModelTC/LightVAE.