Update main.py

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, UploadFile, File, Form
+from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks
 from fastapi.responses import FileResponse
 import tempfile
 import shutil
@@ -77,6 +77,7 @@ async def merge_endpoint(
     files: list[UploadFile] = File(...),
     orig_vol: float = Form(1.0),
     music_vol: float = Form(0.5),
+    background_tasks: BackgroundTasks,
 ):
     temp_dir = tempfile.mkdtemp()
     try:
@@ -93,31 +94,28 @@ async def merge_endpoint(
         # Separate video and audio files
         video_files = [f for f in saved_files if f.lower().endswith(".mp4")]
         audio_files = [
-            f
-            for f in saved_files
-            if f.lower().endswith((".mp3", ".wav", ".aac", ".m4a", ".ogg"))
+            f for f in saved_files if f.lower().endswith((".mp3", ".wav", ".aac", ".m4a", ".ogg"))
         ]
 
         if len(saved_files) < 2:
+            shutil.rmtree(temp_dir, ignore_errors=True)
             return {"error": "Please upload at least 2 files (videos or audios)."}
 
-        # Prepare output path
         output_path = os.path.join(temp_dir, "merged_output.mp4")
 
-        # Merge videos and audios
         merge_videos_and_audios(video_files, audio_files, output_path, temp_dir)
 
-        # Return merged file
+        # Schedule directory removal after response is sent
+        background_tasks.add_task(shutil.rmtree, temp_dir, ignore_errors=True)
+
         return FileResponse(output_path, media_type="video/mp4", filename="merged_output.mp4")
 
     except Exception as e:
+        shutil.rmtree(temp_dir, ignore_errors=True)
         error_msg = f"Error: {str(e)}\n\n{traceback.format_exc()}"
         logger.error(error_msg)
         return {"error": error_msg}
 
-    finally:
-        shutil.rmtree(temp_dir, ignore_errors=True)
-
 # Optional: log public URL
 def log_api_url():
     url = os.getenv("SPACE_PUBLIC_URL")