Update gradio_app.py

gradio_app.py

@@ -77,6 +77,12 @@ def check_outputs_folder(folder_path):
     else:
         print(f'The folder {folder_path} does not exist.')
 
+# Custom CUDA memory management function
+def cuda_memory_cleanup():
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+    gc.collect()
+    
 def infer(frame1_path, frame2_path):
 
     seed = 42
@@ -96,7 +102,7 @@ def infer(frame1_path, frame2_path):
     frame2 = load_image(frame2_path)
     frame2 = frame2.resize((512, 288))
 
-    torch.cuda.empty_cache()
+    cuda_memory_cleanup()
 
     frames = pipe(image1=frame1, image2=frame2, 
         num_inference_steps=num_inference_steps, # 50
@@ -104,9 +110,11 @@ def infer(frame1_path, frame2_path):
         weighted_average=weighted_average, # True
         noise_injection_steps=noise_injection_steps, # 0
         noise_injection_ratio= noise_injection_ratio, # 0.5
-         decode_chunk_size=4
+         decode_chunk_size=6
     ).frames[0]
 
+    cuda_memory_cleanup()
+
     print(f"FRAMES: {frames}")
     
     out_dir = "result"