Update app.py

app.py

@@ -3,18 +3,20 @@ import numpy as np
 import torch
 import torchaudio
 import transformers
+import spaces
 from huggingface_hub import hf_hub_download
 
 fe_path = hf_hub_download("sarulab-speech/sidon-v0.1", filename="feature_extractor.pt")
 decoder_path = hf_hub_download("sarulab-speech/sidon-v0.1", filename="decoder.pt")
 
-fe = torch.jit.load(fe_path)
-decoder = torch.jit.load(decoder_path)
+fe = torch.jit.load(fe_path).to('cuda')
+decoder = torch.jit.load(decoder_path).to('cuda')
 preprocessor = transformers.SeamlessM4TFeatureExtractor.from_pretrained(
     "facebook/w2v-bert-2.0"
 )
 
 
+@spaces.GPU
 def denoise_speech(audio):
     if audio is None:
         return None
@@ -39,7 +41,7 @@ def denoise_speech(audio):
     for chunk in wav_16k.view(-1).split(16000 * 20):
         inputs = preprocessor(
             torch.nn.functional.pad(chunk, (40, 40)), return_tensors="pt"
-        )
+        ).to('cuda')
         with torch.inference_mode():
             feature = fe(inputs["input_features"])["last_hidden_state"]
             if feature_cache is not None:
@@ -50,7 +52,7 @@ def denoise_speech(audio):
                 restored_wav = decoder(feature.transpose(1, 2))
                 restored_wav = restored_wav[:, :, 50 * 3 :]
             feature_cache = feature[:, -5:, :]
-        restoreds.append(restored_wav)
+        restoreds.append(restored_wav.cpu())
     restored_wav = torch.cat(restoreds, dim=-1)
 
     return 48_000, (restored_wav.view(-1, 1).numpy() * 32767).astype(np.int16)