Update app.py

app.py

@@ -4,18 +4,18 @@ import shlex
 import tempfile
 import uuid
 from typing import List
+import shutil # Import shutil for rmtree
+import logging
 
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.responses import FileResponse, JSONResponse
-import logging
+from starlette.background import BackgroundTask # Import BackgroundTask
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# --- Define the core ffmpeg function (adapted for API context) ---
-# Note: We pass paths directly now, as they are managed within the API call.
-# We also return stderr for debugging purposes in case of failure.
+# --- Define the core ffmpeg function (remains the same) ---
 def run_ffmpeg_concatenation(input_files: List[str], output_file: str, ffmpeg_executable: str = "ffmpeg"):
     """
     Runs the FFmpeg concatenation process.
@@ -60,8 +60,6 @@ def run_ffmpeg_concatenation(input_files: List[str], output_file: str, ffmpeg_ex
     logger.info("\n--- FFmpeg Execution Start ---")
     stderr_output = ""
     try:
-        # Using Popen for potentially better handling of large outputs if needed,
-        # but run is simpler here. Ensure encoding/errors are handled.
         process = subprocess.run(
             command,
             check=True, # Raises CalledProcessError on non-zero exit
@@ -72,8 +70,6 @@ def run_ffmpeg_concatenation(input_files: List[str], output_file: str, ffmpeg_ex
         )
         stderr_output = process.stderr
         logger.info("--- FFmpeg Execution End ---")
-        # logger.info("\nSTDOUT:") # Often empty for concat
-        # logger.info(process.stdout)
         logger.info("\nSTDERR (Progress/Info):\n" + stderr_output)
         msg = f"Successfully concatenated videos into {os.path.basename(output_file)}"
         logger.info(msg)
@@ -93,6 +89,16 @@ def run_ffmpeg_concatenation(input_files: List[str], output_file: str, ffmpeg_ex
         logger.exception(err_msg) # Log full traceback
         return False, err_msg, stderr_output # Include any captured stderr
 
+# --- Cleanup Function ---
+def cleanup_temp_dir(temp_dir: str):
+    """Removes a temporary directory."""
+    try:
+        logger.info(f"Cleaning up temporary directory: {temp_dir}")
+        shutil.rmtree(temp_dir)
+        logger.info(f"Successfully cleaned up temporary directory: {temp_dir}")
+    except Exception as e:
+        logger.error(f"Error cleaning up temporary directory {temp_dir}: {e}")
+
 # --- FastAPI App Definition ---
 app = FastAPI()
 
@@ -103,6 +109,7 @@ async def concatenate_videos_api(
 ):
     """
     API endpoint to concatenate uploaded video files using FFmpeg re-encoding.
+    Cleans up temporary files after response is sent.
     """
     if not files:
         raise HTTPException(status_code=400, detail="No files were uploaded.")
@@ -111,61 +118,65 @@ async def concatenate_videos_api(
 
     logger.info(f"Received {len(files)} files for concatenation. Output name: {output_filename}")
 
-    # Create a temporary directory to store uploaded files and the output
-    with tempfile.TemporaryDirectory() as temp_dir:
-        input_file_paths = []
-        original_filenames = []
-
-        try:
-            # Save uploaded files to the temporary directory
-            for uploaded_file in files:
-                original_filenames.append(uploaded_file.filename)
-                # Create a unique-ish temp filename to avoid clashes, keep extension
-                _, ext = os.path.splitext(uploaded_file.filename)
-                temp_input_path = os.path.join(temp_dir, f"{uuid.uuid4()}{ext}")
-                logger.info(f"Saving uploaded file '{uploaded_file.filename}' to '{temp_input_path}'")
+    # Create a temporary directory manually
+    temp_dir = tempfile.mkdtemp()
+    logger.info(f"Created temporary directory: {temp_dir}")
+
+    input_file_paths = []
+    original_filenames = []
+
+    try:
+        # Save uploaded files to the temporary directory
+        for uploaded_file in files:
+            original_filenames.append(uploaded_file.filename)
+            _, ext = os.path.splitext(uploaded_file.filename)
+            # Ensure unique name within the temp dir
+            temp_input_path = os.path.join(temp_dir, f"{uuid.uuid4()}{ext}")
+            logger.info(f"Saving uploaded file '{uploaded_file.filename}' to '{temp_input_path}'")
+            try:
                 with open(temp_input_path, "wb") as buffer:
                     buffer.write(await uploaded_file.read())
                 input_file_paths.append(temp_input_path)
-                await uploaded_file.close() # Close the file handle
-
-            logger.info(f"Saved files: {original_filenames}")
-
-            # Define the output path within the temporary directory
-            temp_output_path = os.path.join(temp_dir, output_filename)
-
-            # Run the FFmpeg concatenation logic
-            success, message, ffmpeg_stderr = run_ffmpeg_concatenation(input_file_paths, temp_output_path)
-
-            if success:
-                logger.info(f"Concatenation successful. Preparing file response for: {temp_output_path}")
-                # Return the concatenated file
-                return FileResponse(
-                    path=temp_output_path,
-                    filename=output_filename, # Send back with the desired name
-                    media_type='video/mp4' # Adjust if needed based on output_filename extension
-                )
-            else:
-                logger.error(f"Concatenation failed: {message}")
-                # Return a JSON error response including ffmpeg's stderr if available
-                return JSONResponse(
-                    status_code=500,
-                    content={
-                        "detail": f"Video concatenation failed: {message}",
-                        "ffmpeg_stderr": ffmpeg_stderr,
-                        "input_files": original_filenames
-                    }
-                )
-
-        except Exception as e:
-            logger.exception("An unexpected error occurred in the API endpoint.")
-            raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
+            finally:
+                await uploaded_file.close() # Ensure file handle is closed
+
+        logger.info(f"Saved files: {original_filenames}")
+
+        # Define the output path within the temporary directory
+        temp_output_path = os.path.join(temp_dir, output_filename)
+
+        # Run the FFmpeg concatenation logic
+        success, message, ffmpeg_stderr = run_ffmpeg_concatenation(input_file_paths, temp_output_path)
+
+        if success:
+            logger.info(f"Concatenation successful. Preparing file response for: {temp_output_path}")
+            # Return the concatenated file with a background task for cleanup
+            return FileResponse(
+                path=temp_output_path,
+                filename=output_filename, # Send back with the desired name
+                media_type='video/mp4', # Adjust if needed
+                background=BackgroundTask(cleanup_temp_dir, temp_dir) # Cleanup AFTER sending
+            )
+        else:
+            logger.error(f"Concatenation failed: {message}")
+            # Explicitly clean up now since we are not returning a FileResponse
+            cleanup_temp_dir(temp_dir)
+            # Return a JSON error response
+            return JSONResponse(
+                status_code=500,
+                content={
+                    "detail": f"Video concatenation failed: {message}",
+                    "ffmpeg_stderr": ffmpeg_stderr,
+                    "input_files": original_filenames
+                }
+            )
+
+    except Exception as e:
+        logger.exception("An unexpected error occurred in the API endpoint.")
+        # Explicitly clean up in case of other errors before returning FileResponse
+        cleanup_temp_dir(temp_dir)
+        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
 
 @app.get("/")
 async def read_root():
-    return {"message": "Video Concatenation API is running. Use the /concatenate/ endpoint (POST) to process videos."}
-
-# Example of how to run locally (Hugging Face Spaces handles this automatically)
-# if __name__ == "__main__":
-#     import uvicorn
-#     uvicorn.run(app, host="0.0.0.0", port=8000)
+    return {"message": "Video Concatenation API is running. Use the /concatenate/ endpoint (POST) to process videos."}