Update main.py

main.py

@@ -1,96 +1,152 @@
-import os
-import shutil
-import tempfile
-from fastapi import FastAPI, File, UploadFile, Form, BackgroundTasks
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import FileResponse
-from moviepy.editor import VideoFileClip, AudioFileClip, CompositeAudioClip, concatenate_videoclips
-
-app = FastAPI()
+import tempfile
+import shutil
+import os
+import logging
+import traceback
+from moviepy.editor import (
+    VideoFileClip,
+    concatenate_videoclips,
+    AudioFileClip,
+    CompositeAudioClip,
+    concatenate_audioclips,
+)
+
+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_paths, audio_paths, output_path):
-    # Load video clips
-    video_clips = [VideoFileClip(vp) for vp in video_paths]
+def merge_videos_and_audios(
+    video_files, audio_files, orig_vol=1.0, music_vol=0.5, temp_dir=None
+):
+    try:
+        # Use temp_dir for output path
+        output_path = os.path.join(temp_dir, "merged_output.mp4" if video_files else "merged_output.mp3")
+
+        if not video_files and audio_files:
+            audio_clips = []
+            for a in audio_files:
+                try:
+                    audio_clip = safe_audio_clip(a)
+                    audio_clips.append(audio_clip)
+                except Exception as e:
+                    logger.warning(f"Skipping invalid audio file: {a} → {str(e)}")
+
+            if not audio_clips:
+                raise ValueError("No valid audio files to merge.")
+
+            final_audio = concatenate_audioclips(audio_clips)
+            final_audio.write_audiofile(output_path)
+            for clip in audio_clips:
+                clip.close()
+            final_audio.close()
+            return output_path
+
+        video_clips = [VideoFileClip(v) for v in video_files]
+        final_video = concatenate_videoclips(video_clips, method="compose")
 
-    # 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 audio_files:
+            audio_clips = []
+            for a in audio_files:
+                try:
+                    audio_clip = safe_audio_clip(a)
+                    audio_clips.append(audio_clip)
+                except Exception as e:
+                    logger.warning(f"Skipping invalid audio file: {a} → {str(e)}")
+
+            if audio_clips:
+                final_audio = concatenate_audioclips(audio_clips).volumex(music_vol)
+                original_audio = final_video.audio.volumex(orig_vol) if final_video.audio else None
+
+                if original_audio:
+                    composite_audio = CompositeAudioClip([original_audio, final_audio])
+                else:
+                    composite_audio = final_audio
+
+                final_video = final_video.set_audio(composite_audio)
+
+                for clip in audio_clips:
+                    clip.close()
+            else:
+                logger.warning("No valid audio files found. Using only video audio (if present).")
+                if final_video.audio:
+                    final_video = final_video.volumex(orig_vol)
+        else:
+            if final_video.audio:
+                final_video = final_video.volumex(orig_vol)
 
-    # Combine audio clips into one composite audio clip
-    final_audio = CompositeAudioClip(audio_clips) if audio_clips else None
+        final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")
 
-    # Concatenate video clips
-    final_video = concatenate_videoclips(video_clips)
+        for clip in video_clips:
+            clip.close()
+        final_video.close()
 
-    # Set audio if exists
-    if final_audio:
-        final_video = final_video.set_audio(final_audio)
+        return output_path
 
-    # Write output video file
-    final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")
+    except Exception as e:
+        error_msg = f"Error during merge: {str(e)}\n{traceback.format_exc()}"
+        logger.error(error_msg)
+        raise RuntimeError(error_msg)
 
-    # Close clips
-    final_video.close()
-    for vc in video_clips:
-        vc.close()
-    for ac in audio_clips:
-        ac.close()
 
+app = FastAPI()
 
 @app.post("/merge")
 async def merge_endpoint(
-    video_files: list[UploadFile] = File(...),
-    audio_files: list[UploadFile] = File(None),
-    background_tasks: BackgroundTasks = None
+    files: list[UploadFile] = File(...),
+    orig_vol: float = Form(1.0),
+    music_vol: float = Form(0.5),
 ):
-    # Create temporary directory for uploaded files
     temp_dir = tempfile.mkdtemp()
-
     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
+        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)
 
-        # Return merged video file
-        return FileResponse(output_file, media_type="video/mp4", filename="merged_output.mp4")
+        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"))]
 
-    except Exception as e:
-        # Cleanup temp_dir if exception occurs
-        shutil.rmtree(temp_dir)
-        return {"error": str(e)}
+        if len(saved_files) < 2:
+            raise HTTPException(status_code=400, detail="Please upload at least 2 files (videos or audios).")
+
+        result_path = merge_videos_and_audios(video_files, audio_files, orig_vol, music_vol, temp_dir)
+
+        media_type = "video/mp4" if result_path.lower().endswith(".mp4") else "audio/mpeg"
+        filename = os.path.basename(result_path)
+
+        return FileResponse(result_path, media_type=media_type, filename=filename)
+
+    except RuntimeError as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    finally:
+        # Clean up temp directory safely
+        try:
+            shutil.rmtree(temp_dir)
+        except Exception as e:
+            logger.warning(f"Failed to remove temp directory {temp_dir}: {e}")
+
+
+# Optional logging
+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")
+
+log_api_url()