Test on lengths

app.py

@@ -82,6 +82,9 @@ avse_model.load_state_dict(avse_state_dict, strict=True)
 avse_model.to("cuda")
 avse_model.eval()
 
+CHUNK_SIZE_AUDIO = 48000  # 3 sec at 16kHz
+CHUNK_SIZE_VIDEO = 75     # 25fps × 3 sec
+
 @spaces.GPU
 def run_avse_inference(video_path, audio_path):
     estimated = run_avse(video_path, audio_path)
@@ -101,15 +104,39 @@ def run_avse_inference(video_path, audio_path):
     ]).astype(np.float32)
     bg_frames /= 255.0
 
+    audio_chunks = [
+        noisy[i:i + CHUNK_SIZE_AUDIO]
+        for i in range(0, len(noisy), CHUNK_SIZE_AUDIO)
+    ]
+
+    video_chunks = [
+        bg_frames[i:i + CHUNK_SIZE_VIDEO]
+        for i in range(0, len(bg_frames), CHUNK_SIZE_VIDEO)
+    ]
+
+    min_len = min(len(audio_chunks), len(video_chunks))  # sync length
+
 
     # Combine into input dict (match what model.enhance expects)
-    data = {
-        "noisy_audio": noisy,
-        "video_frames": bg_frames[np.newaxis, ...]
-    }
+    #data = {
+    #    "noisy_audio": noisy,
+    #    "video_frames": bg_frames[np.newaxis, ...]
+    #}
+
+    #with torch.no_grad():
+    #    estimated = avse_model.enhance(data).reshape(-1)
+    estimated_chunks = []
 
     with torch.no_grad():
-        estimated = avse_model.enhance(data).reshape(-1)
+        for i in range(min_len):
+            chunk_data = {
+                "noisy_audio": audio_chunks[i],
+                "video_frames": video_chunks[i][np.newaxis, ...]
+            }
+            est = avse_model.enhance(chunk_data).reshape(-1)
+            estimated_chunks.append(est)
+
+    estimated = np.concatenate(estimated_chunks, axis=0)
 
     # Save result
     tmp_wav = audio_path.replace(".wav", "_enhanced.wav")
@@ -135,6 +162,10 @@ def extract_resampled_audio(video_path, target_sr=16000):
     torchaudio.save(resampled_audio_path, waveform, sample_rate=target_sr)
     return resampled_audio_path
 
+@spaces.GPU
+def yolo_detection(frame, verbose=False):
+    return model(frame, verbose=verbose)[0]
+
 @spaces.GPU
 def extract_faces(video_file):
     cap = cv2.VideoCapture(video_file)
@@ -147,7 +178,8 @@ def extract_faces(video_file):
             break
 
         # Inference
-        results = model(frame, verbose=False)[0]
+        #results = model(frame, verbose=False)[0]
+        results = yolo_detection(frame, verbose=False)
         for box in results.boxes:
             # version 1
             # x1, y1, x2, y2 = map(int, box.xyxy[0])