Update app.py

app.py

@@ -27,7 +27,7 @@ hf_hub_download(
 )
 
 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")
@@ -73,7 +73,7 @@ def clear_gpu():
 
 
 def infer(image_path, prompt, orbit_type, progress=gr.Progress(track_tqdm=True)):
-    transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16)
+    
     pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16)
     lora_path = "checkpoints/"
     adapter_name = None
@@ -105,6 +105,7 @@ def infer(image_path, prompt, orbit_type, progress=gr.Progress(track_tqdm=True))
     )
 
     find_and_move_object_to_cpu()
+    pipe.to("cpu")
     clear_gpu()
     
     # Generate a timestamp for the output filename