Automatic device detection (#2)

- Automatic device detection (a75f7f304fe956d700b84e43c8200b833b477069)


Co-authored-by: Fabrice TIERCELIN <[email protected]>

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", device="cuda:0"):
+    def __init__(self, name="declare-lab/tango2", 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):