Update audiosr/pipeline.py

audiosr/pipeline.py

@@ -80,6 +80,9 @@ def extract_kaldi_fbank_feature(waveform, sampling_rate, log_mel_spec):
     return {"ta_kaldi_fbank": fbank}  # [1024, 128]
 
 
+
+
+
 def make_batch_for_super_resolution(input_file, waveform=None, fbank=None):
     log_mel_spec, stft, waveform, duration, target_frame = read_audio_file(input_file)
 
@@ -111,10 +114,12 @@ def round_up_duration(duration):
     return int(round(duration / 2.5) + 1) * 2.5
 
 
+@spaces.GPU
 def build_model(ckpt_path=None, config=None, device=None, model_name="basic"):
     if device is None or device == "auto":
         if torch.cuda.is_available():
-            device = torch.device("cuda:0")
+            device = torch.Tensor([0]).cuda()
+            # device = torch.device("cuda:0")
         elif torch.backends.mps.is_available():
             device = torch.device("mps")
         else: