Update app.py

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
 import os
 import torch
-from diffusers import CogVideoXImageToVideoPipeline
+from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel
 from diffusers.utils import export_to_video, load_image
+from transformers import T5EncoderModel, T5Tokenizer
+
 from datetime import datetime
 
 from huggingface_hub import hf_hub_download
@@ -22,15 +24,25 @@ hf_hub_download(
     local_dir="checkpoints"
 )
 
-pipe = CogVideoXImageToVideoPipeline.from_pretrained("THUDM/CogVideoX-5b-I2V", torch_dtype=torch.bfloat16)
+
+model_id = "THUDM/CogVideoX-5b-I2V"
+transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16)
+text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16)
+vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16)
+tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer")
+pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16)
 lora_path = "your lora path"
 lora_rank = 256
+
 def infer(prompt, image_path, orbit_type, progress=gr.Progress(track_tqdm=True)):
     lora_path = "checkpoints/"
+    adapter_name = None
     if orbit_type == "Left":
         weight_name = "orbit_left_lora_weights.safetensors"
+        adapter_name = "orbit_left_lora_weights"
     elif orbit_type == "Up":
         weight_name = "orbit_up_lora_weights.safetensors"
+        adapter_name = "orbit_up_lora_weights"
     lora_rank = 256
     pipe.load_lora_weights(lora_path, weight_name=weight_name, adapter_name="test_1")
     pipe.fuse_lora(lora_scale=1 / lora_rank)