Update main.py

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.responses import FileResponse
 import tempfile
 import shutil
@@ -12,136 +12,99 @@ from moviepy.editor import (
     CompositeAudioClip,
     concatenate_audioclips,
 )
+import zipfile
 
+# Configure logging
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("main")
 
-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)}")
+# Merge function
+def merge_videos_and_audios(video_paths, audio_paths, output_path):
+    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")
-
-        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)
-
-        final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")
-
-        for clip in video_clips:
-            clip.close()
-        final_video.close()
-
-        return output_path
+    audio_clips = []
+    for ap in audio_paths:
+        try:
+            audio_clip = AudioFileClip(ap)
+            audio_clips.append(audio_clip)
+        except Exception as e:
+            logger.warning(f"Skipping invalid audio file: {ap}: {str(e)}")
 
-    except Exception as e:
-        error_msg = f"Error during merge: {str(e)}\n{traceback.format_exc()}"
-        logger.error(error_msg)
-        raise RuntimeError(error_msg)
+    final_audio = CompositeAudioClip(audio_clips) if audio_clips else None
+    final_video = concatenate_videoclips(video_clips, method="compose")
+
+    if final_audio is not None:
+        final_video = final_video.set_audio(final_audio)
+
+    final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")
 
+    # Cleanup
+    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(
-    files: list[UploadFile] = File(...),
+    zip_file: 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)
+        # Save uploaded ZIP
+        zip_path = os.path.join(temp_dir, zip_file.filename)
+        with open(zip_path, "wb") as f:
+            f.write(await zip_file.read())
+
+        # Extract ZIP
+        extract_dir = os.path.join(temp_dir, "extracted")
+        os.makedirs(extract_dir, exist_ok=True)
+
+        with zipfile.ZipFile(zip_path, "r") as zip_ref:
+            zip_ref.extractall(extract_dir)
+
+        # Collect files
+        files = os.listdir(extract_dir)
+        video_files = [
+            os.path.join(extract_dir, f)
+            for f in files
+            if f.lower().endswith(".mp4")
+        ]
+        audio_files = [
+            os.path.join(extract_dir, f)
+            for f in files
+            if f.lower().endswith((".mp3", ".wav", ".aac", ".m4a", ".ogg"))
+        ]
+
+        if len(video_files) == 0 and len(audio_files) == 0:
+            return {"error": "No valid video or audio files found in ZIP."}
+
+        # Output path
+        output_path = os.path.join(temp_dir, "merged_output.mp4")
+
+        # Merge
+        merge_videos_and_audios(video_files, audio_files, output_path)
+
+        # Return result
+        return FileResponse(
+            output_path,
+            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"))]
-
-        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))
+    except Exception as e:
+        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+        logger.error(error_msg)
+        return {"error": error_msg}
 
     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}")
-
+        shutil.rmtree(temp_dir)
 
-# Optional logging
+# Optional: log public URL
 def log_api_url():
     url = os.getenv("SPACE_PUBLIC_URL")
     if url: