Update app.py

app.py

@@ -22,8 +22,16 @@ from gradio import Markdown
 
 import spaces
 
+# Automatic device detection
+if torch.cuda.is_available():
+    device_type = "cuda"
+    device_selection = "cuda:0"
+else:
+    device_type = "cpu"
+    device_selection = "cpu"
+    
 class MusicFeaturePredictor:
-    def __init__(self, path, device="cuda:0", cache_dir=None, local_files_only=False):
+    def __init__(self, path, device=device_selection, cache_dir=None, local_files_only=False):
         self.beats_tokenizer = AutoTokenizer.from_pretrained(
             "microsoft/deberta-v3-large",
             use_fast=False,
@@ -147,7 +155,7 @@ class Mustango:
     def __init__(
         self,
         name="declare-lab/mustango",
-        device="cuda:0",
+        device=device_selection,
         cache_dir=None,
         local_files_only=False,
     ):
@@ -217,10 +225,11 @@ class Mustango:
 
 
 # Initialize Mustango
-if torch.cuda.is_available():
-    mustango = Mustango()
-else:
-    mustango = Mustango(device="cpu")
+mustango = Mustango(device="cpu")
+# if torch.cuda.is_available():
+#     mustango = Mustango()
+# else:
+#     mustango = Mustango(device="cpu")
     
 # output_wave = mustango.generate("This techno song features a synth lead playing the main melody.", 5, 3, disable_progress=False)
 @spaces.GPU(duration=120)