Update main.py

main.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, UploadFile, File, Form
-from fastapi.responses import FileResponse, JSONResponse
+from fastapi.responses import FileResponse
 import tempfile
 import shutil
 import os
@@ -18,19 +18,16 @@ import numpy as np
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Safe audio clip loader
 def safe_audio_clip(path):
-    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
-
-def safe_video_clip(path):
-    clip = VideoFileClip(path)
-    if clip.reader is None or clip.duration is None or clip.duration <= 0:
-        clip.close()
-        raise ValueError(f"Invalid or empty video file: {os.path.basename(path)}")
-    return clip
+    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, orig_vol=1.0, music_vol=0.5, temp_dir=None
@@ -39,53 +36,74 @@ def merge_videos_and_audios(
         output_path = os.path.join(
             temp_dir, "merged_output.mp4" if video_files else "merged_output.mp3"
         )
-        
+
         # If no videos, just concatenate audios
         if not video_files and audio_files:
-            logger.info(f"Merging {len(audio_files)} audio files...")
-            audio_clips = [safe_audio_clip(a) for a in 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()
-            logger.info(f"Audio merge complete: {output_path}")
             return output_path
 
         # If videos are present:
-        logger.info(f"Merging {len(video_files)} video files...")
-        video_clips = [safe_video_clip(v) for v in video_files]
+        video_clips = [VideoFileClip(v) for v in video_files]
         final_video = concatenate_videoclips(video_clips, method="compose")
-        
+
         if audio_files:
-            logger.info(f"Merging {len(audio_files)} audio files into video...")
-            audio_clips = [safe_audio_clip(a) for a in audio_files]
-            final_audio = concatenate_audioclips(audio_clips).volumex(music_vol)
-            
-            # Handle video original audio
-            original_audio = (
-                final_video.audio.volumex(orig_vol) if final_video.audio is not None else None
-            )
-            if original_audio:
-                composite_audio = CompositeAudioClip([original_audio, final_audio])
+            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 is not None
+                    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:
-                composite_audio = final_audio
+                logger.warning("No valid audio files found. Using only video audio (if present).")
+                if final_video.audio is not None:
+                    final_video = final_video.volumex(orig_vol)
 
-            final_video = final_video.set_audio(composite_audio)
-            for clip in audio_clips:
-                clip.close()
         else:
-            # No audio uploaded → optionally adjust video audio volume
+            # No audio files provided
             if final_video.audio is not None:
                 final_video = final_video.volumex(orig_vol)
-            # else → no audio at all
 
-        # Final video output
         final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")
+
         for clip in video_clips:
             clip.close()
         final_video.close()
-        logger.info(f"Video merge complete: {output_path}")
+
         return output_path
 
     except Exception as e:
@@ -93,7 +111,6 @@ def merge_videos_and_audios(
         logger.error(error_msg)
         return error_msg
 
-# FastAPI app
 app = FastAPI()
 
 @app.post("/merge")
@@ -104,7 +121,6 @@ async def merge_endpoint(
 ):
     temp_dir = tempfile.mkdtemp()
     try:
-        # Save all uploaded files
         saved_files = []
         for uploaded_file in files:
             file_path = os.path.join(temp_dir, uploaded_file.filename)
@@ -112,42 +128,35 @@ async def merge_endpoint(
                 content = await uploaded_file.read()
                 out_file.write(content)
             saved_files.append(file_path)
-            logger.info(f"Saved file: {uploaded_file.filename}")
 
-        # Separate video and audio files
         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"))]
+        audio_files = [
+            f
+            for f in saved_files
+            if f.lower().endswith((".mp3", ".wav", ".aac", ".m4a", ".ogg"))
+        ]
 
-        # Require at least 2 files
         if len(saved_files) < 2:
-            return JSONResponse(
-                status_code=400,
-                content={"error": "Please upload at least 2 files (videos or audios)."},
-            )
+            return {"error": "Please upload at least 2 files (videos or audios)."}
 
-        # Perform merge
         result_path = merge_videos_and_audios(
             video_files, audio_files, orig_vol, music_vol, temp_dir
         )
 
-        # If error occurred
         if isinstance(result_path, str) and result_path.startswith("Error"):
-            return JSONResponse(
-                status_code=500,
-                content={"error": result_path},
-            )
+            return {"error": result_path}
 
-        # Return merged file
-        media_type = "video/mp4" if result_path.lower().endswith(".mp4") else "audio/mpeg"
+        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)
 
     finally:
-        # Clean up temp dir
         shutil.rmtree(temp_dir)
-        logger.info("Temporary files cleaned up.")
 
-# Log the API public URL for HuggingFace Spaces
+# Optional: log space public URL
 def log_api_url():
     url = os.getenv("SPACE_PUBLIC_URL")
     if url: