Update app.py

app.py

@@ -1,12 +1,16 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, BackgroundTasks, HTTPException
+from fastapi.responses import FileResponse
+import threading
+import queue
+import uuid
+import os
+import time
 
 import audiocraft
-print("Audiocraft version:", audiocraft.__version__)
-
 from audiocraft.models import MusicGen
 from audiocraft.data.audio import audio_write
-import uuid
-import os
+
+print("Audiocraft version:", audiocraft.__version__)
 
 app = FastAPI()
 
@@ -17,28 +21,140 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 print("Loading MusicGen model...")
 model = MusicGen.get_pretrained(MODEL_NAME)
 model.set_generation_params(
-    duration=30,   # seconds
+    duration=30,
     temperature=1.0,
     top_k=250,
     top_p=0.0
 )
 
+# -------------------------
+# Job system (in-memory)
+# -------------------------
+
+job_queue = queue.Queue()
+jobs = {}  # job_id -> metadata
+
+
+class JobStatus:
+    QUEUED = "queued"
+    PROCESSING = "processing"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+def worker():
+    """Background worker that processes queued jobs"""
+    while True:
+        job_id = job_queue.get()
+        job = jobs.get(job_id)
+
+        if not job:
+            job_queue.task_done()
+            continue
+
+        try:
+            jobs[job_id]["status"] = JobStatus.PROCESSING
+
+            wav = model.generate([job["prompt"]])[0]
+
+            filename = f"{job_id}.wav"
+            path = os.path.join(OUTPUT_DIR, filename)
+
+            audio_write(
+                path,
+                wav.cpu(),
+                model.sample_rate,
+                strategy="loudness",
+                loudness_compressor=True
+            )
+
+            jobs[job_id].update({
+                "status": JobStatus.COMPLETED,
+                "file_path": path
+            })
+
+        except Exception as e:
+            jobs[job_id].update({
+                "status": JobStatus.FAILED,
+                "error": str(e)
+            })
+
+        job_queue.task_done()
+
+
+# Start worker thread
+threading.Thread(target=worker, daemon=True).start()
+
+# -------------------------
+# API endpoints
+# -------------------------
+
 @app.post("/generate")
 async def generate(prompt: str):
-    wav = model.generate([prompt])[0]
+    job_id = uuid.uuid4().hex
 
-    filename = f"{uuid.uuid4().hex}"
-    path = os.path.join(OUTPUT_DIR, filename)
+    jobs[job_id] = {
+        "status": JobStatus.QUEUED,
+        "prompt": prompt,
+        "created_at": time.time(),
+        "file_path": None
+    }
 
-    audio_write(
-        path,
-        wav.cpu(),
-        model.sample_rate,
-        strategy="loudness",
-        loudness_compressor=True
-    )
+    job_queue.put(job_id)
 
     return {
-        "prompt": prompt,
-        "file": f"{filename}.wav"
+        "job_id": job_id,
+        "status_url": f"/status/{job_id}"
     }
+
+
+@app.get("/status/{job_id}")
+async def status(job_id: str):
+    job = jobs.get(job_id)
+
+    if not job:
+        raise HTTPException(status_code=404, detail="Job not found")
+
+    response = {
+        "job_id": job_id,
+        "status": job["status"]
+    }
+
+    if job["status"] == JobStatus.COMPLETED:
+        response["download_url"] = f"/download/{job_id}"
+
+    if job["status"] == JobStatus.FAILED:
+        response["error"] = job.get("error")
+
+    return response
+
+
+def delete_file(path: str, delay: int = 10):
+    """Deletes file after response is sent"""
+    time.sleep(delay)
+    if os.path.exists(path):
+        os.remove(path)
+
+
+@app.get("/download/{job_id}")
+async def download(job_id: str, background_tasks: BackgroundTasks):
+    job = jobs.get(job_id)
+
+    if not job:
+        raise HTTPException(status_code=404, detail="Job not found")
+
+    if job["status"] != JobStatus.COMPLETED:
+        raise HTTPException(status_code=400, detail="Job not completed")
+
+    path = job["file_path"]
+
+    if not path or not os.path.exists(path):
+        raise HTTPException(status_code=404, detail="File not found")
+
+    background_tasks.add_task(delete_file, path)
+
+    return FileResponse(
+        path,
+        media_type="audio/wav",
+        filename=os.path.basename(path)
+    )