Upload new model

app.py

@@ -61,6 +61,7 @@ import tempfile
 from ultralytics import YOLO
 from moviepy import ImageSequenceClip
 from scipy.io import wavfile
+from avse_code import run_avse
 
 # Load face detector
 model = YOLO("yolov8n-face.pt").cuda()  # assumes CUDA available
@@ -75,18 +76,23 @@ import spaces
 #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/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()