Update app.py

app.py

@@ -35,6 +35,9 @@ vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dt
 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)
 
+# Add this near the top after imports
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+
 def find_and_move_object_to_cpu():
     for obj in gc.get_objects():
         try:
@@ -52,52 +55,48 @@ def clear_gpu():
     gc.collect()
 
 def infer(image_path, prompt, orbit_type, progress=gr.Progress(track_tqdm=True)):
+    # Move everything to CPU initially
+    pipe.to("cpu")
+    torch.cuda.empty_cache()
     
     lora_path = "checkpoints/"
-    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"
+    weight_name = "orbit_left_lora_weights.safetensors" if orbit_type == "Left" else "orbit_up_lora_weights.safetensors"
     lora_rank = 128
-
     adapter_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
 
-    # Load LoRA weights on CPU, move to GPU afterward
+    # Load LoRA weights on CPU
     pipe.load_lora_weights(lora_path, weight_name=weight_name, adapter_name=f"adapter_{adapter_timestamp}")
     pipe.fuse_lora(lora_scale=1 / lora_rank)
     
-    # Move the pipeline to GPU for inference
-    pipe.to("cuda")
+    try:
+        # Move to GPU just before inference
+        pipe.to("cuda")
+        torch.cuda.empty_cache()
+        
+        prompt = f"{prompt}. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+        image = load_image(image_path)
+        seed = random.randint(0, 2**8 - 1)
+        
+        with torch.inference_mode():
+            video = pipe(
+                image, 
+                prompt, 
+                num_inference_steps=25,
+                guidance_scale=7.0,
+                use_dynamic_cfg=True,
+                generator=torch.Generator(device="cpu").manual_seed(seed)
+            )
+    finally:
+        # Ensure cleanup happens even if inference fails
+        pipe.to("cpu")
+        pipe.unfuse_lora()
+        pipe.unload_lora_weights()
+        torch.cuda.empty_cache()
+        gc.collect()
     
-    # Set the inference prompt
-    prompt = f"{prompt}. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
-    image = load_image(image_path)
-    seed = random.randint(0, 2**8 - 1)
-    
-    
-    video = pipe(
-        image, 
-        prompt, 
-        num_inference_steps=25,
-        guidance_scale=7.0,
-        use_dynamic_cfg=True,
-        generator=torch.Generator(device="cpu").manual_seed(seed)
-    )
-
-    torch.cuda.empty_cache()
-    pipe.unfuse_lora()
-    pipe.unload_lora_weights()
-
-    
-    # Generate and save output video
+    # Generate output video
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     export_to_video(video.frames[0], f"output_{timestamp}.mp4", fps=8)
-
-    # Move objects to CPU and clear GPU memory immediately after inference
-    find_and_move_object_to_cpu()
-    clear_gpu()
     
     return f"output_{timestamp}.mp4"