Update app.py

app.py

@@ -4,23 +4,21 @@ import torchaudio
 from einops import rearrange
 from stable_audio_tools import get_pretrained_model
 from stable_audio_tools.inference.generation import generate_diffusion_cond
-from huggingface_hub import cached_download, hf_hub_url
-from transformers import AutoModelForAudioClassification
 import os
 
-token = os.environ.get("TOKEN")
-model_name = "stabilityai/stable-audio-open-1.0"
-model_config_url = hf_hub_url(repo_id=model_name, revision="main", filename="model_config.json")
-model_config = cached_download(model_config_url, use_auth_token=token)
-
-model = AutoModelForAudioClassification.from_pretrained(
-    model_name, 
-    cache_dir=None,
-    use_auth_token=token
+# Load model config from stable-audio-tools
+model, model_config = get_pretrained_model(
+    "stabilityai/stable-audio-open-1.0", config_filename="model_config.json"
 )
 sample_rate = model_config["sample_rate"]
 sample_size = model_config["sample_size"]
 
+# Load the model using the transformers library 
+token = os.environ.get("TOKEN")
+model = AutoModelForAudioClassification.from_pretrained(
+    "stabilityai/stable-audio-open-1.0", use_auth_token=token, cache_dir=None
+)
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = model.to(device)
 
@@ -32,7 +30,7 @@ def generate_music(prompt, seconds_total, bpm, genre):
     # Set up text and timing conditioning
     conditioning = [{
         "prompt": f"{bpm} BPM {genre} {prompt}",
-        "seconds_start": 0, 
+        "seconds_start": 0,
         "seconds_total": seconds_total
     }]