Update main.py

main.py

@@ -1,149 +1,96 @@
-from fastapi import FastAPI, UploadFile, File, Form
-from fastapi.responses import FileResponse
-import tempfile
-import shutil
 import os
-import logging
-import traceback
-from moviepy.editor import (
-    VideoFileClip,
-    concatenate_videoclips,
-    AudioFileClip,
-    CompositeAudioClip,
-    concatenate_audioclips,
-)
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-def safe_audio_clip(path):
-    try:
-        clip = AudioFileClip(path)
-        if clip.reader is None or clip.duration is None or clip.duration <= 0:
-            clip.close()
-            raise ValueError(f"Invalid or empty audio file: {os.path.basename(path)}")
-        return clip
-    except Exception as e:
-        raise ValueError(f"Failed to load audio file '{os.path.basename(path)}': {str(e)}")
-
-def merge_videos_and_audios(
-    video_files, audio_files, output_path, temp_dir,
-    orig_vol=1.0, music_vol=0.5,
-):
-    try:
-        # Load video clips
-        video_clips = [VideoFileClip(v) for v in video_files]
-
-        # Load audio clips safely
-        audio_clips = []
-        for a in audio_files:
-            try:
-                clip = safe_audio_clip(a)
-                audio_clips.append(clip)
-            except Exception as e:
-                logger.warning(f"Skipping invalid audio file: {a} → {str(e)}")
-
-        # Prepare final audio composite if audio clips exist
-        if audio_clips:
-            final_audio = concatenate_audioclips(audio_clips).volumex(music_vol)
-        else:
-            final_audio = None
-
-        # Concatenate video clips with method 'compose' for safety
-        final_video = concatenate_videoclips(video_clips, method="compose")
-
-        # Combine video audio with provided original volume
-        original_audio = final_video.audio.volumex(orig_vol) if final_video.audio else None
-
-        # Combine audios if both exist
-        if original_audio and final_audio:
-            composite_audio = CompositeAudioClip([original_audio, final_audio])
-            final_video = final_video.set_audio(composite_audio)
-        elif final_audio:
-            final_video = final_video.set_audio(final_audio)
-        elif original_audio:
-            final_video = final_video.set_audio(original_audio)
-        else:
-            final_video = final_video.set_audio(None)
-
-        # Explicit writable temp audio file path inside temp_dir
-        temp_audiofile_path = os.path.join(temp_dir, "temp_audiofile.m4a")
-
-        # Write output video with explicit temp audio file to avoid permission errors
-        final_video.write_videofile(
-            output_path,
-            codec="libx264",
-            audio_codec="aac",
-            temp_audiofile=temp_audiofile_path,
-            remove_temp=True,
-            verbose=False,
-            logger=None,
-        )
-
-        # Cleanup
-        final_video.close()
-        for vc in video_clips:
-            vc.close()
-        for ac in audio_clips:
-            ac.close()
-
-        return output_path
-
-    except Exception as e:
-        error_msg = f"Error during merge: {str(e)}\n{traceback.format_exc()}"
-        logger.error(error_msg)
-        return error_msg
+import shutil
+import tempfile
+from fastapi import FastAPI, File, UploadFile, Form, BackgroundTasks
+from fastapi.responses import FileResponse
+from moviepy.editor import VideoFileClip, AudioFileClip, CompositeAudioClip, concatenate_videoclips
 
 app = FastAPI()
 
-@app.post("/merge")
-async def merge_endpoint(
-    files: list[UploadFile] = File(...),
-    orig_vol: float = Form(1.0),
-    music_vol: float = Form(0.5),
-):
-    temp_dir = tempfile.mkdtemp()
-    try:
-        saved_files = []
-        for uploaded_file in files:
-            file_path = os.path.join(temp_dir, uploaded_file.filename)
-            with open(file_path, "wb") as out_file:
-                content = await uploaded_file.read()
-                out_file.write(content)
-            saved_files.append(file_path)
+def merge_videos_and_audios(video_paths, audio_paths, output_path):
+    # Load video clips
+    video_clips = [VideoFileClip(vp) for vp in video_paths]
 
-        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"))
-        ]
+    # Load audio clips (skip invalid files)
+    audio_clips = []
+    for ap in audio_paths:
+        try:
+            audio_clip = AudioFileClip(ap)
+            audio_clips.append(audio_clip)
+        except Exception as e:
+            print(f"Warning: Could not load audio file {ap}: {e}")
 
-        if len(saved_files) < 2:
-            return {"error": "Please upload at least 2 files (videos or audios)."}
+    # Combine audio clips into one composite audio clip
+    final_audio = CompositeAudioClip(audio_clips) if audio_clips else None
 
-        output_file = os.path.join(temp_dir, "merged_output.mp4")
+    # Concatenate video clips
+    final_video = concatenate_videoclips(video_clips)
 
-        result_path = merge_videos_and_audios(
-            video_files, audio_files, output_file, temp_dir, orig_vol, music_vol
-        )
+    # Set audio if exists
+    if final_audio:
+        final_video = final_video.set_audio(final_audio)
 
-        if isinstance(result_path, str) and result_path.startswith("Error"):
-            return {"error": result_path}
+    # Write output video file
+    final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")
 
-        media_type = "video/mp4" if result_path.lower().endswith(".mp4") else "audio/mpeg"
-        filename = os.path.basename(result_path)
+    # Close clips
+    final_video.close()
+    for vc in video_clips:
+        vc.close()
+    for ac in audio_clips:
+        ac.close()
 
-        return FileResponse(result_path, media_type=media_type, filename=filename)
 
-    finally:
-        shutil.rmtree(temp_dir)
+@app.post("/merge")
+async def merge_endpoint(
+    video_files: list[UploadFile] = File(...),
+    audio_files: list[UploadFile] = File(None),
+    background_tasks: BackgroundTasks = None
+):
+    # Create temporary directory for uploaded files
+    temp_dir = tempfile.mkdtemp()
 
-# Optional: log space public URL
-def log_api_url():
-    url = os.getenv("SPACE_PUBLIC_URL")
-    if url:
-        logger.info(f"API is available at: {url}/merge")
-    else:
-        logger.info("SPACE_PUBLIC_URL environment variable not found")
+    try:
+        # Save uploaded video files to temp_dir
+        video_paths = []
+        for vf in video_files:
+            video_path = os.path.join(temp_dir, vf.filename)
+            with open(video_path, "wb") as f:
+                content = await vf.read()
+                f.write(content)
+            video_paths.append(video_path)
+
+        # Save uploaded audio files to temp_dir (if any)
+        audio_paths = []
+        if audio_files:
+            for af in audio_files:
+                audio_path = os.path.join(temp_dir, af.filename)
+                with open(audio_path, "wb") as f:
+                    content = await af.read()
+                    f.write(content)
+                audio_paths.append(audio_path)
+
+        # Prepare permanent output directory
+        output_dir = os.path.abspath("output")
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Define output video path
+        output_file = os.path.join(output_dir, "merged_output.mp4")
+
+        # Merge videos and audios
+        merge_videos_and_audios(video_paths, audio_paths, output_file)
+
+        # Add background task to delete temp_dir after response sent
+        if background_tasks:
+            background_tasks.add_task(shutil.rmtree, temp_dir)
+        else:
+            # If no background_tasks support, you can delete later manually or omit
+            pass
 
-log_api_url()
+        # Return merged video file
+        return FileResponse(output_file, media_type="video/mp4", filename="merged_output.mp4")
+
+    except Exception as e:
+        # Cleanup temp_dir if exception occurs
+        shutil.rmtree(temp_dir)
+        return {"error": str(e)}