app.py

import os
import uuid
import shutil
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 starlette.middleware.cors import CORSMiddleware

# Setup
app = FastAPI(title="AI Audio Editor API", description="FastAPI audio editor with vocal remover", version="1.0")

# Directories
TEMP_DIR = "temp"
os.makedirs(TEMP_DIR, exist_ok=True)

# Logger
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("audio_editor")

# CORS (optional for web frontend support)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# 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_file(path: str):
    try:
        os.remove(path)
        logger.info(f"Deleted temp file: {path}")
    except Exception as e:
        logger.error(f"Cleanup failed: {e}")

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

# Root endpoint
@app.get("/", tags=["Root"])
def read_root():
    return {"message": "Welcome to the AI Audio Editor API!"}

# AI vocal remover endpoint
@app.post("/remove_vocals", tags=["AI"])
async def remove_vocals(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(..., description="Audio file for AI vocal removal."),
    output_format: str = Form("mp3", description="Output format (mp3, wav, etc.)")
):
    logger.info(f"Processing file for vocal removal: {file.filename}")

    input_path = save_upload_file(file)
    background_tasks.add_task(cleanup_file, input_path)

    try:
        # Output folder for spleeter
        out_dir = os.path.join(TEMP_DIR, uuid.uuid4().hex)
        os.makedirs(out_dir, exist_ok=True)

        # Use spleeter
        separator = Separator("spleeter:2stems")
        separator.separate_to_file(input_path, out_dir)

        # 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")

        if not os.path.exists(instrumental_path):
            raise FileNotFoundError("Instrumental not generated.")

        # 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: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

    finally:
        shutil.rmtree(out_dir, ignore_errors=True)