Update app.py

app.py

@@ -7,6 +7,10 @@ import subprocess
 import threading
 from typing import Dict
 
+import zipfile
+import tempfile
+
+
 app = FastAPI()
 
 UPLOAD_DIR = "uploads"
@@ -148,25 +152,42 @@ def job_status(job_id: str):
 # -----------------------------
 # Download stem + auto cleanup
 # -----------------------------
-@app.get("/download/{job_id}/{stem}")
-def download_stem(
+@app.get("/download/{job_id}/all")
+def download_all_stems(
     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")
+    stems = job.get("stems")
+    if not stems:
+        raise HTTPException(404, "No stems found")
 
-    # delete after response finishes
+    # create temp zip file
+    tmp = tempfile.NamedTemporaryFile(
+        delete=False,
+        suffix=".zip"
+    )
+    zip_path = tmp.name
+    tmp.close()
+
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for stem_name, stem_path in stems.items():
+            if os.path.exists(stem_path):
+                zipf.write(
+                    stem_path,
+                    arcname=f"{stem_name}.wav"
+                )
+
+    # cleanup after response is sent
     background_tasks.add_task(cleanup_job, job_id)
+    background_tasks.add_task(os.remove, zip_path)
 
     return FileResponse(
-        path,
-        media_type="audio/wav",
-        filename=f"{stem}.wav"
+        zip_path,
+        media_type="application/zip",
+        filename=f"{job_id}_stems.zip"
     )
+