Test on lengths

app.py

@@ -64,7 +64,7 @@ from scipy.io import wavfile
 from avse_code import run_avse
 
 # Load face detector
-model = YOLO("yolov8n-face.pt")  # assumes CUDA available
+model = YOLO("yolov8n-face.pt").cuda()  # assumes CUDA available
 
 
 from decord import VideoReader, cpu
@@ -79,12 +79,11 @@ 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.to("cuda")
 avse_model.eval()
 
 @spaces.GPU
 def run_avse_inference(video_path, audio_path):
-    avse_model.to("cuda")
     estimated = run_avse(video_path, audio_path)
     # Load audio
     #noisy, _ = sf.read(audio_path, dtype='float32')  # (N, )
@@ -148,7 +147,6 @@ def extract_faces(video_file):
             break
 
         # Inference
-        model = model.cuda()
         results = model(frame, verbose=False)[0]
         for box in results.boxes:
             # version 1
@@ -236,4 +234,3 @@ iface = gr.Interface(
 
 iface.launch()
 
-