Upload new model

app.py

@@ -72,44 +72,10 @@ from model import AVSEModule
 from config import sampling_rate
 import spaces
 
-# Load model once globally
-#ckpt_path = "ckpts/ep215_0906.oat.ckpt"
-#model = AVSEModule.load_from_checkpoint(ckpt_path)
-avse_model = AVSEModule()
-#avse_state_dict = torch.load("ckpts/ep215_0906.oat.ckpt")
-avse_state_dict = torch.load("ckpts/ep220_0908.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):
-    run_avse(video_path, audio_path)
-    # Load audio
-    #noisy, _ = sf.read(audio_path, dtype='float32')  # (N, )
-    #noisy = torch.tensor(noisy).unsqueeze(0)  # (1, N)
-    noisy = wavfile.read(audio_path)[1].astype(np.float32) / (2 ** 15)
-
-    # Norm.
-    #noisy = noisy * (0.8 / np.max(np.abs(noisy)))
-
-    # Load grayscale video
-    vr = VideoReader(video_path, ctx=cpu(0))
-    frames = vr.get_batch(list(range(len(vr)))).asnumpy()
-    bg_frames = np.array([
-        cv2.cvtColor(frames[i], cv2.COLOR_RGB2GRAY) for i in range(len(frames))
-    ]).astype(np.float32)
-    bg_frames /= 255.0
-
-
-    # Combine into input dict (match what model.enhance expects)
-    data = {
-        "noisy_audio": noisy,
-        "video_frames": bg_frames[np.newaxis, ...]
-    }
-
-    with torch.no_grad():
-        estimated = avse_model.enhance(data).reshape(-1)
+    estimated = run_avse(video_path, audio_path)
 
     # Save result
     tmp_wav = audio_path.replace(".wav", "_enhanced.wav")