Update app.py

app.py

@@ -1,8 +1,11 @@
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException
+from fastapi.responses import FileResponse
 import torch
 import os
 import uuid
 import subprocess
+import threading
+from typing import Dict
 
 app = FastAPI()
 
@@ -15,27 +18,87 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 QUALITY_PRESETS = {
-    "low": {
-        "shifts": 0,
-        "overlap": 0.25
-    },
-    "medium": {
-        "shifts": 1,
-        "overlap": 0.25
-    },
-    "high": {
-        "shifts": 2,
-        "overlap": 0.5
-    }
+    "low": {"shifts": 0, "overlap": 0.25},
+    "medium": {"shifts": 1, "overlap": 0.25},
+    "high": {"shifts": 2, "overlap": 0.5},
 }
 
+# -----------------------------
+# In-memory job store
+# -----------------------------
+jobs: Dict[str, dict] = {}
+
+# -----------------------------
+# Worker function
+# -----------------------------
+def run_demucs(job_id: str, input_path: str, quality: str):
+    try:
+        jobs[job_id]["status"] = "processing"
+        jobs[job_id]["progress"] = 10
+
+        preset = QUALITY_PRESETS[quality]
+        output_path = os.path.join(OUTPUT_DIR, job_id)
+
+        cmd = [
+            "python3", "-m", "demucs",
+            "--device", DEVICE,
+            "--shifts", str(preset["shifts"]),
+            "--overlap", str(preset["overlap"]),
+            "--out", output_path,
+            input_path
+        ]
+
+        subprocess.run(cmd, check=True)
+
+        base = os.path.splitext(os.path.basename(input_path))[0]
+        stems_dir = os.path.join(output_path, "htdemucs", base)
+
+        jobs[job_id]["stems"] = {
+            "vocals": f"{stems_dir}/vocals.wav",
+            "drums": f"{stems_dir}/drums.wav",
+            "bass": f"{stems_dir}/bass.wav",
+            "other": f"{stems_dir}/other.wav",
+        }
+
+        jobs[job_id]["progress"] = 100
+        jobs[job_id]["status"] = "completed"
+
+    except Exception as e:
+        jobs[job_id]["status"] = "failed"
+        jobs[job_id]["error"] = str(e)
+
+# -----------------------------
+# Cleanup helper
+# -----------------------------
+def cleanup_job(job_id: str):
+    job = jobs.get(job_id)
+    if not job:
+        return
+
+    for path in job.get("stems", {}).values():
+        if os.path.exists(path):
+            os.remove(path)
+
+    input_file = job.get("input_path")
+    if input_file and os.path.exists(input_file):
+        os.remove(input_file)
+
+    output_dir = os.path.join(OUTPUT_DIR, job_id)
+    if os.path.exists(output_dir):
+        subprocess.run(["rm", "-rf", output_dir])
+
+    jobs.pop(job_id, None)
+
+# -----------------------------
+# Create job
+# -----------------------------
 @app.post("/separate")
 async def separate_audio(
     file: UploadFile = File(...),
     quality: str = "medium"
 ):
     if quality not in QUALITY_PRESETS:
-        return {"error": "quality must be low, medium, or high"}
+        raise HTTPException(400, "quality must be low, medium, or high")
 
     job_id = uuid.uuid4().hex
     input_path = os.path.join(UPLOAD_DIR, f"{job_id}_{file.filename}")
@@ -43,29 +106,64 @@ async def separate_audio(
     with open(input_path, "wb") as f:
         f.write(await file.read())
 
-    preset = QUALITY_PRESETS[quality]
-    output_path = os.path.join(OUTPUT_DIR, job_id)
+    jobs[job_id] = {
+        "status": "queued",
+        "progress": 0,
+        "quality": quality,
+        "input_path": input_path,
+        "stems": None,
+    }
 
-    cmd = [
-        "python3", "-m", "demucs",
-        "--device", DEVICE,
-        "--shifts", str(preset["shifts"]),
-        "--overlap", str(preset["overlap"]),
-        "--out", output_path,
-        input_path
-    ]
+    thread = threading.Thread(
+        target=run_demucs,
+        args=(job_id, input_path, quality),
+        daemon=True
+    )
+    thread.start()
 
-    subprocess.run(cmd, check=True)
+    return {
+        "job_id": job_id,
+        "status": "queued"
+    }
 
-    stems_dir = os.path.join(output_path, "htdemucs", os.path.splitext(os.path.basename(input_path))[0])
+# -----------------------------
+# Progress polling
+# -----------------------------
+@app.get("/status/{job_id}")
+def job_status(job_id: str):
+    job = jobs.get(job_id)
+    if not job:
+        raise HTTPException(404, "Job not found")
 
     return {
         "job_id": job_id,
-        "quality": quality,
-        "stems": {
-            "vocals": f"{stems_dir}/vocals.wav",
-            "drums": f"{stems_dir}/drums.wav",
-            "bass": f"{stems_dir}/bass.wav",
-            "other": f"{stems_dir}/other.wav"
-        }
+        "status": job["status"],
+        "progress": job["progress"],
+        "stems": job.get("stems")
     }
+
+# -----------------------------
+# Download stem + auto cleanup
+# -----------------------------
+@app.get("/download/{job_id}/{stem}")
+def download_stem(
+    job_id: str,
+    stem: str,
+    background_tasks: BackgroundTasks
+):
+    job = jobs.get(job_id)
+    if not job or job["status"] != "completed":
+        raise HTTPException(404, "Job not completed")
+
+    path = job["stems"].get(stem)
+    if not path or not os.path.exists(path):
+        raise HTTPException(404, "Stem not found")
+
+    # delete after response finishes
+    background_tasks.add_task(cleanup_job, job_id)
+
+    return FileResponse(
+        path,
+        media_type="audio/wav",
+        filename=f"{stem}.wav"
+    )