Upload model

app.py

@@ -67,16 +67,49 @@ from avse_code import run_avse
 model = YOLO("yolov8n-face.pt").cuda()  # assumes CUDA available
 
 
-
 from decord import VideoReader, cpu
 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):
     estimated = 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)
 
     # Save result
     tmp_wav = audio_path.replace(".wav", "_enhanced.wav")