Spaces:
Running
on
Zero
Running
on
Zero
Fabrice-TIERCELIN
commited on
Commit
•
a75f7f3
1
Parent(s):
0353aff
Automatic device detection
Browse filesHi
@abhinavkashyap92
,
This PR detects automatically the device on which the space runs. It eases the duplication.
_Successfully tested on my space Fabrice-TIERCELIN/Text-to-Audio_
app.py
CHANGED
@@ -10,8 +10,16 @@ from audioldm.variational_autoencoder import AutoencoderKL
|
|
10 |
from gradio import Markdown
|
11 |
import spaces
|
12 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
class Tango:
|
14 |
-
def __init__(self, name="declare-lab/tango2", device=
|
15 |
|
16 |
path = snapshot_download(repo_id=name)
|
17 |
|
@@ -70,9 +78,9 @@ class Tango:
|
|
70 |
# Initialize TANGO
|
71 |
|
72 |
tango = Tango(device="cpu")
|
73 |
-
tango.vae.to(
|
74 |
-
tango.stft.to(
|
75 |
-
tango.model.to(
|
76 |
|
77 |
@spaces.GPU(duration=60)
|
78 |
def gradio_generate(prompt, steps, guidance):
|
|
|
10 |
from gradio import Markdown
|
11 |
import spaces
|
12 |
|
13 |
+
# Automatic device detection
|
14 |
+
if torch.cuda.is_available():
|
15 |
+
device_type = "cuda"
|
16 |
+
device_selection = "cuda:0"
|
17 |
+
else:
|
18 |
+
device_type = "cpu"
|
19 |
+
device_selection = "cpu"
|
20 |
+
|
21 |
class Tango:
|
22 |
+
def __init__(self, name="declare-lab/tango2", device=device_selection):
|
23 |
|
24 |
path = snapshot_download(repo_id=name)
|
25 |
|
|
|
78 |
# Initialize TANGO
|
79 |
|
80 |
tango = Tango(device="cpu")
|
81 |
+
tango.vae.to(device_type)
|
82 |
+
tango.stft.to(device_type)
|
83 |
+
tango.model.to(device_type)
|
84 |
|
85 |
@spaces.GPU(duration=60)
|
86 |
def gradio_generate(prompt, steps, guidance):
|