Update app.py

app.py

@@ -1,24 +1,23 @@
 import os
-import uuid
 import shutil
 import subprocess
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException
+import tempfile
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException, BackgroundTasks
 from fastapi.responses import FileResponse, RedirectResponse
 from fastapi.middleware.cors import CORSMiddleware
-import tempfile
 
 app = FastAPI()
 
-# Add CORS middleware to allow requests from your frontend
+# Add CORS middleware (adjust allow_origins for production)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # In production, set this to your frontend URL
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Create a directory to store temporary files
+# Ensure the base temporary directory exists.
 TEMP_DIR = "temp_files"
 os.makedirs(TEMP_DIR, exist_ok=True)
 
@@ -28,38 +27,38 @@ def read_root():
 
 @app.post("/burn-subtitles")
 async def burn_subtitles(
+    background_tasks: BackgroundTasks,
     video: UploadFile = File(...),
     subtitle: UploadFile = File(...),
     font_size: str = Form("20"),
     font_color: str = Form("white"),
     output_format: str = Form("mp4")
 ):
-    # Create unique ID for this job
-    job_id = str(uuid.uuid4())
-    job_dir = os.path.join(TEMP_DIR, job_id)
-    os.makedirs(job_dir, exist_ok=True)
+    # Create a unique temporary directory for this job.
+    job_dir = tempfile.mkdtemp(prefix='job_', dir=TEMP_DIR)
     
     try:
-        # Save uploaded files
+        # Build file paths.
         video_path = os.path.join(job_dir, f"input{os.path.splitext(video.filename)[1]}")
         subtitle_path = os.path.join(job_dir, "subtitles.srt")
         output_path = os.path.join(job_dir, f"output.{output_format}")
         
-        # Write the files to disk
+        # Save uploaded video file.
         with open(video_path, "wb") as video_file:
             shutil.copyfileobj(video.file, video_file)
         
+        # Save uploaded subtitle file.
         with open(subtitle_path, "wb") as subtitle_file:
             shutil.copyfileobj(subtitle.file, subtitle_file)
         
-        # Ensure font color is in the correct format for FFmpeg
+        # Format the font color for FFmpeg if provided as hex.
         if font_color.startswith('#'):
             font_color = font_color.lstrip('#')
             if len(font_color) == 6:
                 r, g, b = font_color[0:2], font_color[2:4], font_color[4:6]
                 font_color = f"&H{b}{g}{r}&"
         
-        # Run FFmpeg command to burn subtitles
+        # Build the subtitle style string.
         subtitle_style = f"FontSize={font_size},PrimaryColour={font_color}"
         command = [
             "ffmpeg",
@@ -69,12 +68,14 @@ async def burn_subtitles(
             output_path
         ]
         
+        # Run the FFmpeg command.
         process = subprocess.run(command, capture_output=True, text=True)
-        
         if process.returncode != 0:
             raise HTTPException(status_code=500, detail=f"FFmpeg Error: {process.stderr}")
         
-        # Return the processed video
+        # Schedule cleanup of the job directory after the response is sent.
+        background_tasks.add_task(shutil.rmtree, job_dir)
+        
         return FileResponse(
             output_path, 
             media_type=f"video/{output_format}",
@@ -82,29 +83,23 @@ async def burn_subtitles(
         )
     
     except Exception as e:
-        if os.path.exists(job_dir):
-            shutil.rmtree(job_dir)
+        # Clean up immediately on error.
+        shutil.rmtree(job_dir, ignore_errors=True)
         raise HTTPException(status_code=500, detail=str(e))
-    
-    finally:
-        if os.path.exists(job_dir):
-            shutil.rmtree(job_dir)
 
 @app.post("/convert-srt")
 async def convert_srt(
+    background_tasks: BackgroundTasks,
     subtitle: UploadFile = File(...),
 ):
-    # Create unique ID for this job
-    job_id = str(uuid.uuid4())
-    job_dir = os.path.join(TEMP_DIR, job_id)
-    os.makedirs(job_dir, exist_ok=True)
+    job_dir = tempfile.mkdtemp(prefix='job_', dir=TEMP_DIR)
     
     try:
         subtitle_path = os.path.join(job_dir, "subtitles.srt")
-        
         with open(subtitle_path, "wb") as subtitle_file:
             shutil.copyfileobj(subtitle.file, subtitle_file)
         
+        background_tasks.add_task(shutil.rmtree, job_dir)
         return FileResponse(
             subtitle_path, 
             media_type="text/plain",
@@ -112,11 +107,8 @@ async def convert_srt(
         )
     
     except Exception as e:
+        shutil.rmtree(job_dir, ignore_errors=True)
         raise HTTPException(status_code=500, detail=str(e))
-    
-    finally:
-        if os.path.exists(job_dir):
-            shutil.rmtree(job_dir)
 
 @app.get("/health")
 def health_check():