Update app.py

app.py

@@ -1,39 +1,38 @@
-import diffusers
-import huggingface_hub
-print(diffusers.__version__)
-print(huggingface_hub.__version__)
-
+import os
+import uuid
+import threading
+import queue
+from typing import Dict, Optional
 
-from fastapi import FastAPI
 import torch
-import uuid
-import os
 import soundfile as sf
+from fastapi import FastAPI, BackgroundTasks, HTTPException
+from fastapi.responses import FileResponse
 from diffusers import DiffusionPipeline
+from huggingface_hub import login
 
-
+# -------------------- App --------------------
 
 app = FastAPI()
 
 OUTPUT_DIR = "outputs"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# -------- Device & dtype handling (HF-safe) --------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 
 pipe = None  # lazy-loaded
 
+# -------------------- HF Login --------------------
 
-from huggingface_hub import login
 login(token=os.environ["HF_TOKEN"])
 
+# -------------------- Pipeline --------------------
 
 def load_pipeline():
     global pipe
     if pipe is None:
         print("Loading Stable Audio Open pipeline...")
-
         pipe = DiffusionPipeline.from_pretrained(
             "stabilityai/stable-audio-open-1.0",
             torch_dtype=DTYPE,
@@ -41,58 +40,127 @@ def load_pipeline():
             revision="main",
             low_cpu_mem_usage=False,
         )
-
         pipe.enable_attention_slicing()
         pipe.to(DEVICE)
-
         print("Pipeline loaded on", DEVICE)
 
-
-
 @app.on_event("startup")
 def startup():
     load_pipeline()
+    start_worker()
 
-# -------- API --------
-@app.post("/generate")
-async def generate(prompt: str, duration: float = 10.0):
-    load_pipeline()
+# -------------------- Job State --------------------
 
-    # ✅ Stable Audio Open expects duration as a pipeline attribute
-    pipe.audio_length_in_s = float(duration)
+class Job:
+    def __init__(self, prompt: str, duration: float):
+        self.id = uuid.uuid4().hex
+        self.prompt = prompt
+        self.duration = duration
+        self.progress = 0
+        self.status = "queued"  # queued | running | done | error
+        self.filepath: Optional[str] = None
+        self.error: Optional[str] = None
 
-    with torch.no_grad():
-        output = pipe(
-            prompt=prompt,
-            guidance_scale=7.5,
-            num_inference_steps=150,
-        )
+jobs: Dict[str, Job] = {}
+job_queue: queue.Queue[Job] = queue.Queue()
+
+# -------------------- Worker Thread --------------------
+
+def worker_loop():
+    while True:
+        job: Job = job_queue.get()
+        try:
+            job.status = "running"
+            job.progress = 5
+
+            load_pipeline()
 
-    audio = output.audios[0]
+            pipe.audio_length_in_s = float(job.duration)
 
-    # Convert torch → numpy
-    if isinstance(audio, torch.Tensor):
-        audio = audio.detach().cpu().numpy()
-    
-    # If shape is (channels, samples), transpose it
-    if audio.ndim == 2 and audio.shape[0] < audio.shape[1]:
-        audio = audio.T
-    
-    # Ensure float32
-    audio = audio.astype("float32")
-    
-    # Clamp to valid range
-    audio = audio.clip(-1.0, 1.0)
+            with torch.no_grad():
+                output = pipe(
+                    prompt=job.prompt,
+                    guidance_scale=7.5,
+                    num_inference_steps=150,
+                )
 
+            job.progress = 90
 
-    filename = f"{uuid.uuid4().hex}.wav"
-    path = os.path.join(OUTPUT_DIR, filename)
+            audio = output.audios[0]
 
-    sf.write(path, audio, samplerate=44100)
+            if isinstance(audio, torch.Tensor):
+                audio = audio.detach().cpu().numpy()
+
+            if audio.ndim == 2 and audio.shape[0] < audio.shape[1]:
+                audio = audio.T
+
+            audio = audio.astype("float32").clip(-1.0, 1.0)
+
+            filename = f"{job.id}.wav"
+            path = os.path.join(OUTPUT_DIR, filename)
+            sf.write(path, audio, samplerate=44100)
+
+            job.filepath = path
+            job.progress = 100
+            job.status = "done"
+
+        except Exception as e:
+            job.status = "error"
+            job.error = str(e)
+
+        finally:
+            job_queue.task_done()
+
+def start_worker():
+    t = threading.Thread(target=worker_loop, daemon=True)
+    t.start()
+
+# -------------------- API --------------------
+
+@app.post("/generate")
+def generate(prompt: str, duration: float = 10.0):
+    job = Job(prompt=prompt, duration=duration)
+    jobs[job.id] = job
+    job_queue.put(job)
+
+    return {
+        "job_id": job.id,
+        "status": job.status,
+    }
+
+@app.get("/status/{job_id}")
+def status(job_id: str):
+    job = jobs.get(job_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Job not found")
 
     return {
-        "prompt": prompt,
-        "duration": duration,
-        "file": filename,
+        "job_id": job.id,
+        "status": job.status,
+        "progress": job.progress,
+        "error": job.error,
+        "ready": job.status == "done",
     }
 
+def cleanup_file(path: str):
+    try:
+        os.remove(path)
+    except Exception:
+        pass
+
+@app.get("/download/{job_id}")
+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 != "done" or not job.filepath:
+        raise HTTPException(status_code=400, detail="File not ready")
+
+    background_tasks.add_task(cleanup_file, job.filepath)
+
+    return FileResponse(
+        path=job.filepath,
+        media_type="audio/wav",
+        filename=os.path.basename(job.filepath),
+    )