Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import os
 import torch
-from diffusers import CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel
+from diffusers import CogVideoXPipeline, CogVideoXDPMScheduler, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel
 from diffusers.utils import export_to_video, load_image
 from datetime import datetime
 
@@ -22,13 +22,20 @@ hf_hub_download(
     local_dir="checkpoints"
 )
 
-pipe = CogVideoXImageToVideoPipeline.from_pretrained(
-    "THUDM/CogVideoX-5b-I2V", 
+pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16).to(device)
+pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+
+pipe_image = 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)
-
+    vae=pipe.vae,
+    scheduler=pipe.scheduler,
+    tokenizer=pipe.tokenizer,
+    text_encoder=pipe.text_encoder,
+    torch_dtype=torch.bfloat16,
+)
 def infer(prompt, image_path, orbit_type, progress=gr.Progress(track_tqdm=True)):
     lora_path = "checkpoints/"
     if orbit_type == "Left":
@@ -45,7 +52,7 @@ def infer(prompt, image_path, orbit_type, progress=gr.Progress(track_tqdm=True))
     image = load_image(image_path)
     seed = random.randint(0, 2**8 - 1)
     
-    video = pipe(
+    video = pipe_image(
         image, 
         prompt, 
         num_inference_steps=50,  # NOT Changed