Update app.py

app.py

@@ -1,441 +1,95 @@
 import os
 import uuid
-import tempfile
-import logging
 import shutil
-from typing import List, Optional, Literal
-
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException, BackgroundTasks
-from fastapi.responses import FileResponse # JSONResponse removed as not used now
+import logging
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, BackgroundTasks
+from fastapi.responses import FileResponse
+from spleeter.separator import Separator
 from pydub import AudioSegment
-from pydub.exceptions import CouldntDecodeError
+from starlette.middleware.cors import CORSMiddleware
 
-# --- Spleeter (AI Vocal Removal) Imports ---
-# Wrap in try-except to handle potential import errors gracefully
-try:
-    from spleeter.separator import Separator
-    from spleeter.utils import logging as spleeter_logging
-    spleeter_available = True
-    # Optional: Configure Spleeter logging level (e.g., ERROR to reduce noise)
-    # spleeter_logging.set_level(spleeter_logging.ERROR)
-except ImportError:
-    spleeter_available = False
-    Separator = None # Define Separator as None if import fails
-    logging.warning("Spleeter library not found or failed to import.")
-    logging.warning("AI Vocal Removal endpoint (/ai/remove-vocals) will be disabled.")
-    logging.warning("Install spleeter: pip install spleeter")
+# Setup
+app = FastAPI(title="AI Audio Editor API", description="FastAPI audio editor with vocal remover", version="1.0")
 
-
-# --- Configuration & Setup ---
-TEMP_DIR = tempfile.gettempdir()
+# Directories
+TEMP_DIR = "temp"
 os.makedirs(TEMP_DIR, exist_ok=True)
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# --- Global Spleeter Separator Initialization ---
-# Load the model once on startup for better request performance.
-# This increases startup time and initial memory usage significantly.
-# Choose the model: 2stems (vocals/accompaniment), 4stems (v/drums/bass/other), 5stems (v/d/b/piano/other)
-# Using 'spleeter:2stems' - downloads model on first use if not cached.
-spleeter_separator: Optional[Separator] = None
-if spleeter_available:
-    try:
-        logger.info("Initializing Spleeter Separator (Model: spleeter:2stems)... This may download model files.")
-        # MWF = Multi-channel Wiener Filtering (can improve quality but slower)
-        spleeter_separator = Separator('spleeter:2stems', mwf=False)
-        logger.info("Spleeter Separator initialized successfully.")
-    except Exception as e:
-        logger.error(f"FATAL: Failed to initialize Spleeter Separator: {e}", exc_info=True)
-        logger.error("AI Vocal Removal endpoint will likely fail.")
-        spleeter_separator = None # Ensure it's None if init failed
+# Logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("audio_editor")
 
-# --- FastAPI App Initialization ---
-app = FastAPI(
-    title="Advanced Audio Editor API",
-    description="API for audio editing (trim, concat, volume, convert) and AI Vocal Removal (using Spleeter). Requires FFmpeg.",
-    version="2.0.0",
+# CORS (optional for web frontend support)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-# --- Helper Functions (Mostly unchanged, added directory cleanup) ---
+# Helper functions
+def save_upload_file(upload_file: UploadFile) -> str:
+    extension = os.path.splitext(upload_file.filename)[-1]
+    temp_path = os.path.join(TEMP_DIR, f"{uuid.uuid4().hex}{extension}")
+    with open(temp_path, "wb") as buffer:
+        shutil.copyfileobj(upload_file.file, buffer)
+    return temp_path
 
-def cleanup_path(path: str):
-    """Safely remove a file or directory."""
+def cleanup_file(path: str):
     try:
-        if not path or not os.path.exists(path):
-            # logger.debug(f"Cleanup skipped: Path '{path}' does not exist.")
-            return
-
-        if os.path.isfile(path):
-            os.remove(path)
-            logger.info(f"Cleaned up temporary file: {path}")
-        elif os.path.isdir(path):
-            shutil.rmtree(path)
-            logger.info(f"Cleaned up temporary directory: {path}")
-        else:
-             logger.warning(f"Cleanup attempted on non-file/dir path: {path}")
-
+        os.remove(path)
+        logger.info(f"Deleted temp file: {path}")
     except Exception as e:
-        logger.error(f"Error cleaning up path {path}: {e}", exc_info=True)
+        logger.error(f"Cleanup failed: {e}")
 
-async def save_upload_file(upload_file: UploadFile) -> str:
-    """Saves an uploaded file to a temporary location and returns the path."""
-    file_extension = os.path.splitext(upload_file.filename)[1] or '.tmp'
-    # Use a subdirectory within TEMP_DIR for better organization
-    request_temp_dir = os.path.join(TEMP_DIR, f"audio_api_upload_{uuid.uuid4().hex}")
-    os.makedirs(request_temp_dir, exist_ok=True)
-    temp_file_path = os.path.join(request_temp_dir, f"input{file_extension}")
+def export_audio(audio: AudioSegment, output_format: str = "mp3") -> str:
+    output_path = os.path.join(TEMP_DIR, f"{uuid.uuid4().hex}.{output_format}")
+    audio.export(output_path, format=output_format)
+    return output_path
 
-    try:
-        with open(temp_file_path, "wb") as buffer:
-            while content := await upload_file.read(1024 * 1024):
-                buffer.write(content)
-        logger.info(f"Saved uploaded file '{upload_file.filename}' to temp path: {temp_file_path}")
-        return temp_file_path
-    except Exception as e:
-        logger.error(f"Failed to save uploaded file {upload_file.filename}: {e}", exc_info=True)
-        cleanup_path(request_temp_dir) # Cleanup directory if save fails
-        raise HTTPException(status_code=500, detail=f"Could not save uploaded file: {upload_file.filename}")
-    finally:
-        await upload_file.close()
-
-def load_audio(file_path: str) -> AudioSegment:
-    """Loads an audio file using pydub."""
-    # (Implementation unchanged)
-    try:
-        audio = AudioSegment.from_file(file_path)
-        logger.info(f"Loaded audio from: {file_path} (Duration: {len(audio)}ms)")
-        return audio
-    except CouldntDecodeError:
-        logger.warning(f"pydub couldn't decode file: {file_path}. Unsupported format or corrupted?")
-        raise HTTPException(status_code=415, detail=f"Unsupported audio format or corrupted file: {os.path.basename(file_path)}")
-    except FileNotFoundError:
-         logger.error(f"Audio file not found after saving: {file_path}")
-         raise HTTPException(status_code=500, detail="Internal error: Audio file disappeared.")
-    except Exception as e:
-        logger.error(f"Error loading audio file {file_path}: {e}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Error processing audio file: {os.path.basename(file_path)}")
-
-def export_audio(audio: AudioSegment, desired_format: str, base_filename: str = "edited_audio") -> str:
-    """Exports an AudioSegment to a temporary file with specified format and returns the path."""
-    # (Slight modification to allow base filename)
-    output_filename = f"{base_filename}_{uuid.uuid4().hex}.{desired_format.lower()}"
-    # Place export in main TEMP_DIR, not necessarily the upload sub-dir
-    output_path = os.path.join(TEMP_DIR, output_filename)
-    try:
-        logger.info(f"Exporting audio to format '{desired_format}' at {output_path}")
-        # Add bitrate argument for common formats if desired (e.g., "192k" for mp3)
-        export_params = {}
-        if desired_format.lower() == "mp3":
-             export_params['bitrate'] = "192k" # Example bitrate
-
-        audio.export(output_path, format=desired_format.lower(), **export_params)
-        return output_path
-    except Exception as e:
-        logger.error(f"Error exporting audio to format {desired_format}: {e}", exc_info=True)
-        cleanup_path(output_path)
-        raise HTTPException(status_code=500, detail=f"Failed to export audio to format '{desired_format}'.")
-
-
-# --- API Endpoints ---
-
-@app.get("/", tags=["General"])
+# Root endpoint
+@app.get("/", tags=["Root"])
 def read_root():
-    """Root endpoint providing a welcome message and feature status."""
-    features = ["Trim (/trim)", "Concatenate (/concat)", "Volume (/volume)", "Convert (/convert)"]
-    if spleeter_separator:
-        features.append("AI Vocal Removal (/ai/remove-vocals)")
-    else:
-        features.append("AI Vocal Removal (Disabled - Spleeter not available)")
-    return {
-        "message": "Welcome to the Advanced Audio Editor API.",
-        "available_features": features,
-        "important": "AI Vocal Removal is computationally intensive and may take significant time."
-        }
-
-# --- Existing Endpoints (Trim, Concat, Volume, Convert) ---
-# Minor changes: Use updated cleanup_path, ensure input cleanup uses the directory
-#                Use updated export_audio
-
-@app.post("/trim", tags=["Editing - Pydub"])
-async def trim_audio(
-    background_tasks: BackgroundTasks,
-    file: UploadFile = File(..., description="Audio file to trim."),
-    start_ms: int = Form(..., description="Start time in milliseconds."),
-    end_ms: int = Form(..., description="End time in milliseconds.")
-):
-    """Trims an audio file (uses pydub)."""
-    if start_ms < 0 or end_ms <= start_ms:
-        raise HTTPException(status_code=422, detail="Invalid start/end times.")
-
-    logger.info(f"Trim request: file='{file.filename}', start={start_ms}ms, end={end_ms}ms")
-    input_path = await save_upload_file(file)
-    input_dir = os.path.dirname(input_path)
-    background_tasks.add_task(cleanup_path, input_dir) # Schedule input dir cleanup
-
-    output_path = None # Define output_path before try block
-    try:
-        audio = load_audio(input_path)
-        trimmed_audio = audio[start_ms:end_ms]
-        logger.info(f"Audio trimmed to {len(trimmed_audio)}ms")
-
-        original_format = os.path.splitext(file.filename)[1][1:].lower() or "mp3"
-        if original_format in ["tmp", ""]: original_format = "mp3"
-
-        output_path = export_audio(trimmed_audio, original_format, base_filename=f"trimmed_{os.path.splitext(file.filename)[0]}")
-        background_tasks.add_task(cleanup_path, output_path) # Schedule output cleanup
-
-        return FileResponse(
-            path=output_path,
-            media_type=f"audio/{original_format}",
-            filename=f"trimmed_{file.filename}"
-        )
-    except Exception as e:
-        logger.error(f"Error during trim operation: {e}", exc_info=True)
-        # Ensure immediate cleanup on error if possible
-        if output_path: cleanup_path(output_path)
-        # Input dir cleanup is handled by background task unless error is critical before scheduling
-        if isinstance(e, HTTPException): raise e
-        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during trimming: {str(e)}")
-
-
-@app.post("/concat", tags=["Editing - Pydub"])
-async def concatenate_audio(
-    background_tasks: BackgroundTasks,
-    files: List[UploadFile] = File(..., description="Two or more audio files to join in order."),
-    output_format: str = Form("mp3", description="Desired output format (e.g., 'mp3', 'wav', 'ogg').")
-):
-    """Concatenates two or more audio files sequentially (uses pydub)."""
-    if len(files) < 2:
-        raise HTTPException(status_code=422, detail="Please upload at least two files to concatenate.")
+    return {"message": "Welcome to the AI Audio Editor API!"}
 
-    logger.info(f"Concatenate request: {len(files)} files, output_format='{output_format}'")
-    input_dirs = [] # Store directories to clean up
-    loaded_audios = []
-    output_path = None
-
-    try:
-        for file in files:
-            input_path = await save_upload_file(file)
-            input_dir = os.path.dirname(input_path)
-            input_dirs.append(input_dir)
-            background_tasks.add_task(cleanup_path, input_dir)
-            audio = load_audio(input_path)
-            loaded_audios.append(audio)
-
-        if not loaded_audios: raise ValueError("No audio segments loaded.")
-
-        combined_audio = loaded_audios[0]
-        for i in range(1, len(loaded_audios)):
-            combined_audio += loaded_audios[i]
-        logger.info(f"Concatenated audio length: {len(combined_audio)}ms")
-
-        first_filename_base = os.path.splitext(files[0].filename)[0]
-        output_base = f"concat_{first_filename_base}_and_{len(files)-1}_others"
-        output_path = export_audio(combined_audio, output_format, base_filename=output_base)
-        background_tasks.add_task(cleanup_path, output_path)
-
-        return FileResponse(
-            path=output_path,
-            media_type=f"audio/{output_format}",
-            filename=f"{output_base}.{output_format}"
-        )
-    except Exception as e:
-        logger.error(f"Error during concat operation: {e}", exc_info=True)
-        if output_path: cleanup_path(output_path)
-        # Input dirs cleanup handled by background tasks
-        if isinstance(e, HTTPException): raise e
-        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during concatenation: {str(e)}")
-
-
-@app.post("/volume", tags=["Editing - Pydub"])
-async def change_volume(
-    background_tasks: BackgroundTasks,
-    file: UploadFile = File(..., description="Audio file to adjust volume for."),
-    change_db: float = Form(..., description="Volume change in decibels (dB). +/- values.")
-):
-    """Adjusts audio volume (uses pydub)."""
-    logger.info(f"Volume request: file='{file.filename}', change_db={change_db}dB")
-    input_path = await save_upload_file(file)
-    input_dir = os.path.dirname(input_path)
-    background_tasks.add_task(cleanup_path, input_dir)
-    output_path = None
-
-    try:
-        audio = load_audio(input_path)
-        adjusted_audio = audio + change_db
-        logger.info(f"Volume adjusted by {change_db}dB.")
-
-        original_format = os.path.splitext(file.filename)[1][1:].lower() or "mp3"
-        if original_format in ["tmp", ""]: original_format = "mp3"
-
-        output_base = f"volume_{change_db}dB_{os.path.splitext(file.filename)[0]}"
-        output_path = export_audio(adjusted_audio, original_format, base_filename=output_base)
-        background_tasks.add_task(cleanup_path, output_path)
-
-        return FileResponse(
-            path=output_path,
-            media_type=f"audio/{original_format}",
-            filename=f"{output_base}.{original_format}" # Use correct extension
-        )
-    except Exception as e:
-        logger.error(f"Error during volume operation: {e}", exc_info=True)
-        if output_path: cleanup_path(output_path)
-        if isinstance(e, HTTPException): raise e
-        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during volume adjustment: {str(e)}")
-
-
-@app.post("/convert", tags=["Editing - Pydub"])
-async def convert_format(
+# AI vocal remover endpoint
+@app.post("/remove_vocals", tags=["AI"])
+async def remove_vocals(
     background_tasks: BackgroundTasks,
-    file: UploadFile = File(..., description="Audio file to convert."),
-    output_format: str = Form(..., description="Target audio format (e.g., 'mp3', 'wav', 'ogg', 'flac').")
+    file: UploadFile = File(..., description="Audio file for AI vocal removal."),
+    output_format: str = Form("mp3", description="Output format (mp3, wav, etc.)")
 ):
-    """Converts audio file format (uses pydub)."""
-    allowed_formats = {'mp3', 'wav', 'ogg', 'flac', 'aac', 'm4a'}
-    safe_output_format = output_format.lower()
-    if safe_output_format not in allowed_formats:
-         raise HTTPException(status_code=422, detail=f"Invalid output format. Allowed: {', '.join(allowed_formats)}")
+    logger.info(f"Processing file for vocal removal: {file.filename}")
 
-    logger.info(f"Convert request: file='{file.filename}', output_format='{safe_output_format}'")
-    input_path = await save_upload_file(file)
-    input_dir = os.path.dirname(input_path)
-    background_tasks.add_task(cleanup_path, input_dir)
-    output_path = None
+    input_path = save_upload_file(file)
+    background_tasks.add_task(cleanup_file, input_path)
 
     try:
-        audio = load_audio(input_path)
-        filename_base = os.path.splitext(file.filename)[0]
-        output_base = f"{filename_base}_converted"
+        # Output folder for spleeter
+        out_dir = os.path.join(TEMP_DIR, uuid.uuid4().hex)
+        os.makedirs(out_dir, exist_ok=True)
 
-        output_path = export_audio(audio, safe_output_format, base_filename=output_base)
-        background_tasks.add_task(cleanup_path, output_path)
+        # Use spleeter
+        separator = Separator("spleeter:2stems")
+        separator.separate_to_file(input_path, out_dir)
 
-        return FileResponse(
-            path=output_path,
-            media_type=f"audio/{safe_output_format}",
-            filename=f"{output_base}.{safe_output_format}"
-        )
-    except Exception as e:
-        logger.error(f"Error during convert operation: {e}", exc_info=True)
-        if output_path: cleanup_path(output_path)
-        if isinstance(e, HTTPException): raise e
-        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during format conversion: {str(e)}")
-
-
-# --- AI Vocal Removal Endpoint ---
-
-@app.post("/ai/remove-vocals", tags=["Editing - AI"])
-async def ai_remove_vocals(
-    background_tasks: BackgroundTasks,
-    file: UploadFile = File(..., description="Audio file containing mixed vocals and accompaniment."),
-    stem_to_return: Literal['accompaniment', 'vocals'] = Form("accompaniment", description="Which stem to return: 'accompaniment' (default) or 'vocals'."),
-    output_format: str = Form("wav", description="Output format for the separated stem (e.g., 'wav', 'mp3'). WAV recommended for quality.")
-):
-    """
-    Separates vocals from accompaniment using Spleeter (AI model).
-    NOTE: This is computationally intensive and can take significant time.
-    """
-    if not spleeter_separator:
-        logger.warning("Vocal removal endpoint called, but Spleeter is not available.")
-        raise HTTPException(status_code=503, detail="AI Vocal Removal service is unavailable (Spleeter not loaded).")
-
-    logger.info(f"AI Vocal Removal request: file='{file.filename}', return='{stem_to_return}', format='{output_format}'")
-
-    input_path = await save_upload_file(file)
-    input_dir = os.path.dirname(input_path) # Directory where input was saved
-    spleeter_output_dir = os.path.join(TEMP_DIR, f"spleeter_out_{uuid.uuid4().hex}") # Unique output dir for Spleeter
-    final_output_path = None # Path to the file that will be returned
-
-    # Schedule cleanup for both input dir and potential Spleeter output dir
-    background_tasks.add_task(cleanup_path, input_dir)
-    background_tasks.add_task(cleanup_path, spleeter_output_dir) # This will be created by Spleeter
+        # Locate instrumental file
+        base_name = os.path.splitext(os.path.basename(input_path))[0]
+        instrumental_path = os.path.join(out_dir, base_name, "accompaniment.wav")
 
-    try:
-        logger.info(f"Starting Spleeter separation for {input_path} into {spleeter_output_dir}...")
-        # Spleeter separates into the specified directory, creating <filename>/vocals.wav and <filename>/accompaniment.wav
-        # We pass the input *file* path and the desired *output directory* path.
-        spleeter_separator.separate_to_file(
-            input_path,
-            spleeter_output_dir,
-            codec='wav' # Spleeter defaults to WAV, ensuring consistent intermediate format
-        )
-        logger.info(f"Spleeter separation completed.")
-
-        # Spleeter creates a subdirectory named after the input file (without extension)
-        input_filename_base = os.path.splitext(os.path.basename(input_path))[0]
-        stem_output_folder = os.path.join(spleeter_output_dir, input_filename_base)
-
-        # Determine the path to the requested stem file (always WAV from Spleeter)
-        target_stem_filename = f"{stem_to_return}.wav"
-        raw_stem_path = os.path.join(stem_output_folder, target_stem_filename)
-
-        if not os.path.exists(raw_stem_path):
-            logger.error(f"Spleeter output stem not found: {raw_stem_path}")
-            raise HTTPException(status_code=500, detail=f"AI separation failed: Could not find the '{stem_to_return}' stem.")
-
-        # --- Optional Conversion ---
-        safe_output_format = output_format.lower()
-        if safe_output_format == "wav":
-            # No conversion needed, return the direct Spleeter output
-            # We need to move/copy it out of the spleeter dir *or* just return it directly
-            # For simplicity and better cleanup, let's return it directly.
-            # BUT FileResponse needs the final path, and background task cleans the whole spleeter_output_dir.
-            # SAFER: Copy the desired file out to the main TEMP_DIR before returning.
-            final_output_path = os.path.join(TEMP_DIR, f"{input_filename_base}_{stem_to_return}_{uuid.uuid4().hex}.wav")
-            shutil.copyfile(raw_stem_path, final_output_path)
-            logger.info(f"Copied requested WAV stem to final output path: {final_output_path}")
-            background_tasks.add_task(cleanup_path, final_output_path) # Schedule cleanup for the copy
+        if not os.path.exists(instrumental_path):
+            raise FileNotFoundError("Instrumental not generated.")
 
-        else:
-            # Convert the WAV stem to the desired format using pydub
-            logger.info(f"Loading separated '{stem_to_return}' stem for conversion to '{safe_output_format}'...")
-            audio_stem = load_audio(raw_stem_path) # Load the WAV stem
-            output_base = f"{input_filename_base}_{stem_to_return}"
-            final_output_path = export_audio(audio_stem, safe_output_format, base_filename=output_base)
-            logger.info(f"Converted stem saved to: {final_output_path}")
-            background_tasks.add_task(cleanup_path, final_output_path) # Schedule cleanup for converted file
-
-        # --- Return Result ---
-        if not final_output_path or not os.path.exists(final_output_path):
-             raise HTTPException(status_code=500, detail="Failed to prepare final output file after separation.")
-
-        return FileResponse(
-            path=final_output_path,
-            media_type=f"audio/{safe_output_format}", # Use the final format's media type
-            filename=os.path.basename(final_output_path) # Use the actual generated filename
-        )
+        # Convert to desired format
+        instrumental_audio = AudioSegment.from_file(instrumental_path)
+        output_path = export_audio(instrumental_audio, output_format)
+        background_tasks.add_task(cleanup_file, output_path)
 
+        return FileResponse(path=output_path, filename=f"instrumental_{file.filename}", media_type=f"audio/{output_format}")
+    
     except Exception as e:
-        logger.error(f"Error during AI Vocal Removal operation: {e}", exc_info=True)
-        if final_output_path: cleanup_path(final_output_path) # Attempt immediate cleanup if needed
-        # Input/Spleeter dir cleanup handled by background tasks
-        if isinstance(e, HTTPException): raise e
-        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during AI processing: {str(e)}")
+        logger.error(f"Error: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
 
-
-# --- How to Run ---
-# 1. Make sure FFmpeg is installed and accessible in your PATH.
-# 2. Save this code as `app.py`.
-# 3. Create `requirements.txt` (as shown above).
-# 4. Install dependencies: `pip install -r requirements.txt` (THIS MAY TAKE A WHILE!)
-# 5. Run the FastAPI server: `uvicorn app:app --reload`
-#
-# --- Example Usage (using curl) ---
-#
-# **AI Remove Vocals (Get Accompaniment as WAV):**
-# curl -X POST "http://127.0.0.1:8000/ai/remove-vocals" \
-#      -F "file=@my_song_mix.mp3" \
-#      -F "stem_to_return=accompaniment" \
-#      -F "output_format=wav" \
-#      --output accompaniment_output.wav
-#
-# **AI Remove Vocals (Get Vocals as MP3):**
-# curl -X POST "http://127.0.0.1:8000/ai/remove-vocals" \
-#      -F "file=@another_track.wav" \
-#      -F "stem_to_return=vocals" \
-#      -F "output_format=mp3" \
-#      --output vocals_only_output.mp3
-#
-# (Other examples for /trim, /concat, /volume, /convert remain the same as before)
+    finally:
+        shutil.rmtree(out_dir, ignore_errors=True)