Update gradio_app.py

gradio_app.py

@@ -11,6 +11,10 @@ from attn_ctrl.attention_control import (AttentionStore,
                                          register_temporal_self_attention_flip_control,
 )
 from torch.cuda.amp import autocast
+import gc
+
+# Set PYTORCH_CUDA_ALLOC_CONF
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 
 # Set up device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -64,6 +68,7 @@ register_temporal_self_attention_flip_control(pipe.unet, controller, controller_
 def cuda_memory_cleanup():
     torch.cuda.empty_cache()
     torch.cuda.ipc_collect()
+    gc.collect()
 
 def check_outputs_folder(folder_path):
     if os.path.exists(folder_path) and os.path.isdir(folder_path):
@@ -82,7 +87,7 @@ def check_outputs_folder(folder_path):
 @torch.no_grad()
 def infer(frame1_path, frame2_path):
     seed = 42
-    num_inference_steps = 10
+    num_inference_steps = 5  # Reduced from 10
     noise_injection_steps = 0
     noise_injection_ratio = 0.5
     weighted_average = False
@@ -92,30 +97,50 @@ def infer(frame1_path, frame2_path):
         generator = generator.manual_seed(seed)
 
     frame1 = load_image(frame1_path)
-    frame1 = frame1.resize((512, 288))
+    frame1 = frame1.resize((256, 144))  # Reduced from (512, 288)
 
     frame2 = load_image(frame2_path)
-    frame2 = frame2.resize((512, 288))
+    frame2 = frame2.resize((256, 144))  # Reduced from (512, 288)
 
+    # Clear CUDA cache
     cuda_memory_cleanup()
 
-    with autocast():
-        frames = pipe(image1=frame1, image2=frame2, 
-            num_inference_steps=num_inference_steps,
-            generator=generator,
-            weighted_average=weighted_average,
-            noise_injection_steps=noise_injection_steps,
-            noise_injection_ratio=noise_injection_ratio,
-        ).frames[0]
-
-    frames = [frame.cpu() for frame in frames]
-
-    out_dir = "result"
-    check_outputs_folder(out_dir)
-    os.makedirs(out_dir, exist_ok=True)
-    out_path = "result/video_result.gif"
+    # Move model to CPU and clear CUDA cache
+    pipe.to("cpu")
+    cuda_memory_cleanup()
 
-    return "done"
+    # Move model back to GPU
+    pipe.to(device)
+
+    try:
+        with autocast(device_type='cuda', dtype=torch.float16):
+            frames = pipe(
+                image1=frame1, 
+                image2=frame2, 
+                num_inference_steps=num_inference_steps,
+                generator=generator,
+                weighted_average=weighted_average,
+                noise_injection_steps=noise_injection_steps,
+                noise_injection_ratio=noise_injection_ratio,
+            ).frames[0]
+
+        frames = [frame.cpu() for frame in frames]
+
+        out_dir = "result"
+        check_outputs_folder(out_dir)
+        os.makedirs(out_dir, exist_ok=True)
+        out_path = "result/video_result.gif"
+
+        return "done"
+    except RuntimeError as e:
+        if "CUDA out of memory" in str(e):
+            return "Error: CUDA out of memory. Try reducing the image size or using fewer inference steps."
+        else:
+            return f"An error occurred: {str(e)}"
+    finally:
+        # Move model back to CPU and clear CUDA cache
+        pipe.to("cpu")
+        cuda_memory_cleanup()
 
 @torch.no_grad()
 def load_model():