Fabrice-TIERCELIN commited on
Commit
a75f7f3
1 Parent(s): 0353aff

Automatic device detection

Browse files

Hi

@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_

Files changed (1) hide show
  1. app.py +12 -4
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="cuda:0"):
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("cuda")
74
- tango.stft.to("cuda")
75
- tango.model.to("cuda")
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):