add AutoencoderKLWan

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 import spaces
 
 import torch
-from diffusers import WanPipeline
+from diffusers import WanPipeline, AutoencoderKLWan
 from diffusers.utils import export_to_video, load_video
 from vibt.wan import load_vibt_weight, encode_video
 from vibt.scheduler import ViBTScheduler
@@ -23,12 +23,11 @@ def get_fps(path):
 
 
 base_model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-pipe = WanPipeline.from_pretrained(
-    base_model_id,
-    torch_dtype=torch.bfloat16,
-    low_cpu_mem_usage=True,
-    keep_in_fp32_modules=False,  # <-- 或者你也可以强制关掉
-).to("cuda")
+vae = AutoencoderKLWan.from_pretrained(
+    base_model_id, subfolder="vae", torch_dtype=torch.float32
+)
+pipe = WanPipeline.from_pretrained(base_model_id, vae=vae, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
 load_vibt_weight(
     pipe.transformer,
     "Yuanshi/ViBT",