Abd756/StemSense

1from fastapi import FastAPI, BackgroundTasks, HTTPException, Form2from fastapi.middleware.cors import CORSMiddleware3from fastapi.responses import FileResponse4from pydantic import BaseModel5from typing import Optional6import uuid7import os8import shutil9from datetime import datetime10 11# Import our StemSense modules12from core.downloader import AudioDownloader13from core.stems import StemSeparator14from core.analyzer import AudioAnalyzer15from core.packager import Packager16from config import EXPORT_DIR17 18app = FastAPI(19    title="StemSense API",20    description="AI Audio Separation & Analysis Service",21    version="1.0.0"22)23 24# 🛡️ CORS Middleware (Gateway) configuration25# This allows our Next.js frontend (port 3000) to talk to this API (port 8000)26origins = [27    "http://localhost:3000",28    "http://127.0.0.1:3000",29    "https://*.hf.space",30    31]32 33app.add_middleware(34    CORSMiddleware,35    allow_origins=origins,36    allow_credentials=True,37    allow_methods=["*"], # Allow all methods (GET, POST, etc.)38    allow_headers=["*"], # Allow all headers39)40 41# In-memory storage for task statuses (In a real app, use Redis or a DB)42tasks = {}43 44class ProcessRequest(BaseModel):45    input: str46 47class TaskStatus(BaseModel):48    task_id: str49    status: str50    result_file: Optional[str] = None51    error: Optional[str] = None52    created_at: str53 54# Helper function to run the heavy processing in the background55def run_full_workflow(task_id: str, query: str):56    tasks[task_id]["status"] = "downloading"57    58    downloader = AudioDownloader()59    separator = StemSeparator()60    analyzer = AudioAnalyzer()61    packager = Packager()62 63    try:64        # 1. Download65        audio_path = downloader.download(query)66        if not audio_path:67            tasks[task_id]["status"] = "failed"68            tasks[task_id]["error"] = "Download failed"69            return70 71        track_name = os.path.splitext(os.path.basename(audio_path))[0]72        73        # 2. Separate74        tasks[task_id]["status"] = "separating"75        stems_dir = separator.separate(audio_path)76        if not stems_dir:77            tasks[task_id]["status"] = "failed"78            tasks[task_id]["error"] = "Stem separation failed"79            return80 81        # 3. Analyze82        tasks[task_id]["status"] = "analyzing"83        analysis_results = analyzer.analyze(audio_path)84 85        # 4. Package86        tasks[task_id]["status"] = "packaging"87        zip_path = packager.create_package(88            track_name=track_name,89            original_file=audio_path,90            stems_dir=stems_dir,91            analysis_data=analysis_results or {"note": "analysis failed"}92        )93 94        if zip_path:95            tasks[task_id]["status"] = "completed"96            tasks[task_id]["result_file"] = os.path.basename(zip_path)97        else:98            tasks[task_id]["status"] = "failed"99            tasks[task_id]["error"] = "Packaging failed"100 101    except Exception as e:102        tasks[task_id]["status"] = "failed"103        tasks[task_id]["error"] = str(e)104 105@app.get("/")106async def root():107    return {"message": "Welcome to StemSense API. Use POST /process to start."}108 109@app.post("/process", response_model=dict)110async def process_audio(background_tasks: BackgroundTasks, input: str = Form(...)):111    """112    Submit a song name or YouTube URL for processing via Form Data.113    """114    task_id = str(uuid.uuid4())115    tasks[task_id] = {116        "task_id": task_id,117        "status": "queued",118        "result_file": None,119        "error": None,120        "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")121    }122    123    # Start the background task124    background_tasks.add_task(run_full_workflow, task_id, input)125    126    return {"task_id": task_id, "message": "Job submitted successfully"}127 128@app.get("/tasks/{task_id}", response_model=TaskStatus)129async def get_status(task_id: str):130    """131    Check the status of a processing task.132    """133    if task_id not in tasks:134        raise HTTPException(status_code=404, detail="Task not found")135    return tasks[task_id]136 137@app.get("/download/{filename}")138async def download_file(filename: str):139    """140    Download a completed ZIP package.141    """142    file_path = os.path.join(EXPORT_DIR, filename)143    if not os.path.exists(file_path):144        raise HTTPException(status_code=404, detail="File not found")145    146    return FileResponse(147        path=file_path,148        filename=filename,149        media_type='application/zip'150    )151 152if __name__ == "__main__":153    import uvicorn154    uvicorn.run(app, host="0.0.0.0", port=8000)155