Update app.py

app.py

@@ -3,30 +3,50 @@ import torch
 import uuid
 import os
 import soundfile as sf
-
-from diffusers.pipelines.stable_audio import StableAudioPipeline
-
+from diffusers import DiffusionPipeline
 
 app = FastAPI()
 
 OUTPUT_DIR = "outputs"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-print("Loading Stable Audio Open...")
-pipe = StableAudioPipeline.from_pretrained(
-    "stabilityai/stable-audio-open-1.0",
-    torch_dtype=torch.float16
-)
-pipe.to("cuda")
+# -------- 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
+
 
+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,
+        )
+        pipe.to(DEVICE)
+        pipe.set_progress_bar_config(disable=True)
+        print("Pipeline loaded on", DEVICE)
+
+
+# -------- API --------
 @app.post("/generate")
 async def generate(prompt: str, duration: float = 10.0):
-    audio = pipe(
-        prompt,
-        duration=duration,
-        guidance_scale=7.5,
-        num_inference_steps=200
-    ).audios[0]
+    """
+    Generate music/audio from a text prompt using Stable Audio Open
+    """
+    load_pipeline()
+
+    with torch.no_grad():
+        output = pipe(
+            prompt=prompt,
+            duration=duration,
+            guidance_scale=7.5,
+            num_inference_steps=150,  # safer for HF
+        )
+
+    audio = output.audios[0]
 
     filename = f"{uuid.uuid4().hex}.wav"
     path = os.path.join(OUTPUT_DIR, filename)
@@ -36,5 +56,5 @@ async def generate(prompt: str, duration: float = 10.0):
     return {
         "prompt": prompt,
         "duration": duration,
-        "file": filename
+        "file": filename,
     }