Upload

app.py

@@ -10,9 +10,6 @@ def install_mamba():
     subprocess.run(shlex.split("pip install https://github.com/state-spaces/mamba/releases/download/v2.2.2/mamba_ssm-2.2.2+cu122torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"))
     subprocess.run(shlex.split("pip install numpy==1.26.4"))
 
-    subprocess.run(shlex.split("ls"))
-
-
 install_mamba()
 
 
@@ -36,7 +33,6 @@ from models.pcs400   import cal_pcs
 ckpt = "ckpts/SEMamba_advanced.pth"
 cfg_f = "recipes/SEMamba_advanced.yaml"
 
-
 # load config
 with open(cfg_f) as f:
     cfg = yaml.safe_load(f)
@@ -49,12 +45,17 @@ hop_size    = stft_cfg["hop_size"]
 win_size    = stft_cfg["win_size"]
 compress_ff = model_cfg["compress_factor"]
 
-# init 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()
+
+@spaces.GPU
+def load_model():
+    # init 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()
+
+load_model()
 
    
 @spaces.GPU