Update app.py

app.py

@@ -48,6 +48,7 @@ import torch
 import numpy as np
 from PIL import Image
 import cv2
+import gc 
 import tempfile
 import moviepy.editor as mp
 from facexlib.utils.face_restoration_helper import FaceRestoreHelper
@@ -64,6 +65,11 @@ from diffusers.models import AutoencoderKLCogVideoX
 from transformers import SiglipImageProcessor, SiglipVisionModel
 from diffposetalk.diffposetalk import DiffPoseTalk
 
+def cleanup_resources():
+    """Clear CUDA cache and garbage collect"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
 
 # Helper functions from the original script
 def parse_video(driving_frames, max_frame_num, fps=25):
@@ -284,6 +290,7 @@ def process_image_audio(image_path, audio_path, guidance_scale=3.0, steps=10, pr
     comparison_with_audio = save_video_with_audio(comparison_path, audio_path, comparison_with_audio)
     
     progress(1.0, desc="Done!")
+    cleanup_resources
     return result_path, comparison_with_audio
 
 # Create Gradio interface