Update module

app.py

@@ -73,10 +73,11 @@ import spaces
 # Load model once globally
 #ckpt_path = "ckpts/ep215_0906.oat.ckpt"
 #model = AVSEModule.load_from_checkpoint(ckpt_path)
-state_dict = torch.load("ckpts/ep215_0906.oat.ckpt")
-model.load_state_dict(state_dict, strict=True)
-model.to("cuda")
-model.eval()
+avse_model = AVSEModule()
+avse_state_dict = torch.load("ckpts/ep215_0906.oat.ckpt")
+avse_model.load_state_dict(avse_state_dict, strict=True)
+avse_model.to("cuda")
+avse_model.eval()
 
 @spaces.GPU
 def run_avse_inference(video_path, audio_path):
@@ -97,7 +98,7 @@ def run_avse_inference(video_path, audio_path):
     }
 
     with torch.no_grad():
-        estimated = model.enhance(data).reshape(-1).cpu().numpy()
+        estimated = avse_model.enhance(data).reshape(-1).cpu().numpy()
 
     # Save result
     tmp_wav = audio_path.replace(".wav", "_enhanced.wav")