Update app.py

app.py

@@ -15,15 +15,22 @@ from accelerate.utils import set_seed
 from latentsync.whisper.audio2feature import Audio2Feature
 from openai import OpenAI
 from elevenlabs import set_api_key, generate, play, clone, Voice, VoiceSettings
+from torch.cuda.amp import autocast
  
 # Initialize the Flask app
 app = Flask(__name__)
 TEMP_DIR = None
 VIDEO_DIRECTORY = None
 
+def clear_cuda_memory():
+    torch.cuda.empty_cache()
+    gc.collect(
+
 def run_inference(video_path, audio_path, video_out_path,
                   inference_ckpt_path, unet_config_path="configs/unet/second_stage.yaml",
                   inference_steps=20, guidance_scale=1.0, seed=1247):
+    clear_cuda_memory()
+    
     # Load configuration
     config = OmegaConf.load(unet_config_path)
  
@@ -77,19 +84,22 @@ def run_inference(video_path, audio_path, video_out_path,
     else:
         torch.seed()
  
-    # Run the pipeline
-    pipeline(
-        video_path=video_path,
-        audio_path=audio_path,
-        video_out_path=video_out_path,
-        video_mask_path=video_out_path.replace(".mp4", "_mask.mp4"),
-        num_frames=config.data.num_frames,
-        num_inference_steps=inference_steps,
-        guidance_scale=guidance_scale,
-        weight_dtype=dtype,
-        width=config.data.resolution,
-        height=config.data.resolution,
-    )
+    with autocast():
+        try:
+            pipeline(
+                video_path=video_path,
+                audio_path=audio_path,
+                video_out_path=video_out_path,
+                video_mask_path=video_out_path.replace(".mp4", "_mask.mp4"),
+                num_frames=config.data.num_frames,
+                num_inference_steps=inference_steps,
+                guidance_scale=guidance_scale,
+                weight_dtype=dtype,
+                width=config.data.resolution,
+                height=config.data.resolution,
+            )
+        finally:
+            clear_cuda_memory()
 
 def create_temp_dir():
     return tempfile.TemporaryDirectory()