Upload

app.py

@@ -39,19 +39,18 @@ hop_size    = stft_cfg["hop_size"]
 win_size    = stft_cfg["win_size"]
 compress_ff = model_cfg["compress_factor"]
 
+
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda"
+model  = SEMamba(cfg).to(device)
+sdict  = torch.load(ckpt, map_location=device)
+model.load_state_dict(sdict["generator"])
+model.eval()
+
+
    
 @spaces.GPU
 def enhance(filepath):
-    global model
-    if model is None:
-        print("loading model")
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model  = SEMamba(cfg).to(device)
-        sdict  = torch.load(ckpt, map_location=device)
-        model.load_state_dict(sdict["generator"])
-        model.eval()
-        print("Finished.")
-
     # load & (if needed) resample to model SR
     wav, orig_sr = librosa.load(filepath, sr=None)
     if orig_sr != SR:
@@ -62,7 +61,8 @@ def enhance(filepath):
     x = (x*norm).unsqueeze(0)
     # STFT → model → ISTFT
     amp,pha,_ = mag_phase_stft(x, **stft_cfg, compress_factor=model_cfg["compress_factor"])
-    amp2,pha2 = model(amp, pha)
+    with torch.no_grad():
+        amp2, pha2, comp = model(amp, pha)
     out = mag_phase_istft(amp2, pha2, **stft_cfg, compress_factor=model_cfg["compress_factor"])
     out = (out/norm).squeeze().cpu().numpy()
     # back to original rate