Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import os
 import torch
-from diffusers import CogVideoXImageToVideoPipeline
+from diffusers import CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel
 from diffusers.utils import export_to_video, load_image
 from datetime import datetime
 
@@ -22,9 +22,14 @@ hf_hub_download(
     local_dir="checkpoints"
 )
 
-pipe = CogVideoXImageToVideoPipeline.from_pretrained("THUDM/CogVideoX-5b-I2V", torch_dtype=torch.bfloat16)
+pipe = CogVideoXImageToVideoPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b-I2V", 
+    transformer=CogVideoXTransformer3DModel.from_pretrained(
+        "THUDM/CogVideoX-5b-I2V", subfolder="transformer", torch_dtype=torch.bfloat16
+    ),
+    torch_dtype=torch.bfloat16)
 
-def infer(prompt, image_path, orbit_type):
+def infer(prompt, image_path, orbit_type, progress=gr.Progress(track_tqdm=True)):
     lora_path = None
     if orbit_type == "Left":
         lora_path = "checkpoints/orbit_left_lora_weights.safetensors"