Update app.py

app.py

@@ -225,19 +225,22 @@ class Mustango:
 
 
 # Initialize Mustango
-mustango = Mustango(device="cpu")
-mustango.vae.to(device_type)
-mustango.stft.to(device_type)
-mustango.model.to(device_type)
-mustango.music_model.beats_model.to(device_type)
-mustango.music_model.chords_model.to(device_type)
+# mustango = Mustango(device="cpu")
+# mustango.vae.to(device_type)
+# mustango.stft.to(device_type)
+# mustango.model.to(device_type)
+# mustango.music_model.beats_model.to(device_type)
+# mustango.music_model.chords_model.to(device_type)
+        
+# if torch.cuda.is_available():
+#     mustango = Mustango(device=device_selection)
+# else:
+#     mustango = Mustango(device="CPU")
+
+mustango = Mustango(device=device_selection)
 
 mustango.model.device = device_selection
 
-# 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)