Update app.py

app.py

@@ -10,8 +10,16 @@ from audioldm.variational_autoencoder import AutoencoderKL
 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 Tango:
-    def __init__(self, name="declare-lab/tango2-full", device="cuda:0"):
+    def __init__(self, name="declare-lab/tango2-full", device=device_selection):
         
         path = snapshot_download(repo_id=name)
         
@@ -70,9 +78,9 @@ class Tango:
 # Initialize TANGO
 
 tango = Tango(device="cpu")
-tango.vae.to("cuda")
-tango.stft.to("cuda")
-tango.model.to("cuda")
+tango.vae.to(device_type)
+tango.stft.to(device_type)
+tango.model.to(device_type)
     
 @spaces.GPU(duration=60)
 def gradio_generate(prompt, steps, guidance):
@@ -106,8 +114,6 @@ description_text = """
 <p><a href="https://huggingface.co/spaces/declare-lab/tango2/blob/main/app.py?duplicate=true"> <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a> For faster inference without waiting in queue, you may duplicate the space and upgrade to a GPU in the settings. <br/><br/>
 Generate audio using Tango2 by providing a text prompt. Tango2 was built from Tango and was trained on <a href="https://huggingface.co/datasets/declare-lab/audio-alpaca">Audio-alpaca</a>
 <br/><br/> This is the demo for Tango2 for text to audio generation: <a href="https://arxiv.org/abs/2404.09956">Read our paper.</a>
-<br>
-This version is trained on an extended version of Audio-alpaca.
 <p/>
 """
 # Gradio input and output components