Update app.py

app.py

@@ -2,32 +2,24 @@ import spaces
 import gradio as gr
 import torch
 import numpy as np
-from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
+from diffusers import WanImageToVideoPipeline
 from diffusers.utils import export_to_video, load_image
 from transformers import CLIPVisionModel
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-## Loading Encoder
-model_id = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-print(f"Using video Model: {model_id}")
 
 
-print("###### Loading image encoder ######")
+## Loading Encoder
 
-image_encoder = CLIPVisionModel.from_pretrained(
-    model_id, subfolder="image_encoder", torch_dtype=torch.float32
-)
 
 
-print("###### Loading VAE encoder ######")
-vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+model_id = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
 
-print("Loading Pipeline...")
+print(f"Using video Model: {model_id}")
+dtype = torch.bfloat16
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-pipe = WanImageToVideoPipeline.from_pretrained(
-    model_id, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16
-)
+pipe = WanImageToVideoPipeline.from_pretrained(model_id, torch_dtype=dtype)
+pipe.to(device)
 
 try:
     pipe.to(device)