Update handler.py

handler.py

@@ -10,7 +10,7 @@ SAMPLE_RATE = 16000
 class EndpointHandler():
     def __init__(self, path=""):
         self.pipeline = Pipeline.from_pretrained(
-            "pyannote/speaker-diarization",
+            "pyannote/speaker-diarization@2.1", # 3.0 and later is nor supported as of yet in dec 2023
             use_auth_token=os.environ.get("HF_API_TOKEN")
         )
         self.pipeline.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
@@ -20,12 +20,12 @@ class EndpointHandler():
         Args:
             data (Dict):
                 'inputs': Base64-encoded audio bytes
-                'parameters': Additional diarization parameters, including 'num_speakers' (optional)
+                'parameters': Additional diarization parameters (currently unused)
         Return:
             Dict: Speaker diarization results
         """
         inputs = data.get("inputs")
-        parameters = data.get("parameters", {})
+        parameters = data.get("parameters", {}) # We are not using them now, since model don't take speaker count anymore
 
         # Decode the base64 audio data
         audio_data = base64.b64decode(inputs)
@@ -42,27 +42,9 @@ class EndpointHandler():
 
         pyannote_input = {"waveform": audio_tensor, "sample_rate": SAMPLE_RATE}
 
-        # Extract num_speakers from parameters, if present
-        num_speakers = parameters.pop("num_speakers", None)
-
-        # Run diarization pipeline
+        # Run diarization pipeline (without num_speakers)
         try:
-            if num_speakers is not None:
-                diarization = self.pipeline(pyannote_input, num_speakers=num_speakers)  # Adjust parameters as needed for version 2.1.1
-            else:
-                diarization = self.pipeline(pyannote_input)
-        except TypeError as e:
-            print(f"Error: TypeError: {e}")
-            if "num_speakers" in str(e):
-                print("The 'num_speakers' parameter might not be supported by this version of the pipeline.")
-                print("Trying without num_speakers...")
-                try:
-                    diarization = self.pipeline(pyannote_input)
-                except Exception as e:
-                   print(f"An error occurred even without 'num_speakers': {e}")
-                   return {"error": "Diarization failed"}
-            else:
-                return {"error": "Diarization failed with an unexpected TypeError. Check the server logs for details."}
+            diarization = self.pipeline(pyannote_input)
         except Exception as e:
             print(f"An unexpected error occurred: {e}")
             return {"error": "Diarization failed unexpectedly"}