Update app.py

app.py

@@ -9,7 +9,8 @@ import torch
 import uuid
 import os
 import soundfile as sf
-from diffusers import DiffusionPipeline
+from diffusers import StableAudioPipeline
+
 
 app = FastAPI()
 
@@ -24,6 +25,8 @@ pipe = None  # lazy-loaded
 
 
 from huggingface_hub import login
+login(token=os.environ["HF_TOKEN"])
+
 
 def load_pipeline():
     global pipe
@@ -35,13 +38,10 @@ def load_pipeline():
         if token:
             login(token=token)
 
-        pipe = DiffusionPipeline.from_pretrained(
-            "stabilityai/stable-audio-open-1.0",
-            torch_dtype=DTYPE,
-            use_auth_token=True,  # IMPORTANT
-        )
-        pipe.to(DEVICE)
-        pipe.set_progress_bar_config(disable=True)
+        pipe = StableAudioPipeline.from_pretrained("stabilityai/stable-audio-open-1.0", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,)
+        pipe.to(device)
+        pipe.enable_attention_slicing()
+
         print("Pipeline loaded on", DEVICE)