Update handler.py

handler.py

@@ -24,6 +24,8 @@
 from pyannote.audio import Pipeline, Audio
 import torch
 import os
+import threading
+import time
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -47,23 +49,41 @@ class EndpointHandler:
         # initialize audio reader
         self._io = Audio()
 
+
     def __call__(self, data):
         inputs = data.pop("inputs", data)
         waveform = torch.tensor(inputs["waveform"])
         sample_rate = inputs["sample_rate"]
-
         parameters = data.pop("parameters", dict())
-        diarization = self._pipeline(
-            {"waveform": waveform, "sample_rate": sample_rate}, **parameters
-        )
 
-        processed_diarization = [
-            {
-                "speaker": speaker,
-                "start": f"{turn.start:.3f}",
-                "end": f"{turn.end:.3f}",
-            }
-            for turn, _, speaker in diarization.itertracks(yield_label=True)
-        ]
+        # Container for storing diarization result
+        diarization_result = {}
+
+        def diarize():
+            nonlocal diarization_result
+            diarization = self._pipeline(
+                {"waveform": waveform, "sample_rate": sample_rate}, **parameters
+            )
+            diarization_result = [
+                {
+                    "speaker": speaker,
+                    "start": f"{turn.start:.3f}",
+                    "end": f"{turn.end:.3f}",
+                }
+                for turn, _, speaker in diarization.itertracks(yield_label=True)
+            ]
+
+        # Running diarization in a separate thread
+        diarization_thread = threading.Thread(target=diarize)
+        diarization_thread.start()
+
+        # Wait for the diarization to complete or timeout
+        diarization_thread.join(timeout=298)
+
+        # Check if the thread is still alive (indicating a timeout occurred)
+        if diarization_thread.is_alive():
+            print("Diarization timed out")
+            # Handle the timeout case, maybe by raising an error or a warning
+            raise TimeoutError("Diarization process exceeded time limit.")
 
-        return {"diarization": processed_diarization}
+        return {"diarization": diarization_result}