Update handler.py

handler.py

@@ -3,48 +3,54 @@ from pyannote.audio import Pipeline
 import torch
 import base64
 import numpy as np
-import os
 
 SAMPLE_RATE = 16000
 
 class EndpointHandler():
     def __init__(self, path=""):
-        self.pipeline = Pipeline.from_pretrained(
-            "pyannote/speaker-diarization@2.1",
-            use_auth_token=os.environ.get("HF_API_TOKEN")
-        )
-        self.pipeline.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        self.pipeline = Pipeline.from_pretrained("KIFF/pyannote-speaker-diarization-endpoint")
 
     def __call__(self, data: Dict) -> Dict:
         """
         Args:
             data (Dict):
                 'inputs': Base64-encoded audio bytes
-                'parameters': Additional diarization parameters (currently unused)
+                'parameters': Additional diarization parameters, including 'num_speakers' (optional)
         Return:
             Dict: Speaker diarization results
         """
         inputs = data.get("inputs")
-        parameters = data.get("parameters", {})  # We are not using them now
+        parameters = data.get("parameters", {})  # Default to empty dict if not provided
 
         # Decode the base64 audio data
         audio_data = base64.b64decode(inputs)
         audio_nparray = np.frombuffer(audio_data, dtype=np.int16)
 
-        # Handle multi-channel audio (convert to mono)
-        if audio_nparray.ndim > 1:
-            audio_nparray = audio_nparray.mean(axis=0)  # Average channels to create mono
-
         # Convert to PyTorch tensor
         audio_tensor = torch.from_numpy(audio_nparray).float().unsqueeze(0)
-        if audio_tensor.dim() == 1:
-            audio_tensor = audio_tensor.unsqueeze(0)
-
         pyannote_input = {"waveform": audio_tensor, "sample_rate": SAMPLE_RATE}
 
-        # Run diarization pipeline (without num_speakers)
+        # Extract num_speakers from parameters, if present
+        num_speakers = parameters.pop("num_speakers", None)
+
+        # Run diarization pipeline
         try:
-            diarization = self.pipeline(pyannote_input)  # No num_speakers parameter
+            if num_speakers is not None:
+                diarization = self.pipeline(pyannote_input, num_speakers=num_speakers, **parameters)
+            else:
+                diarization = self.pipeline(pyannote_input, **parameters)
+        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, **parameters)
+                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."}
         except Exception as e:
             print(f"An unexpected error occurred: {e}")
             return {"error": "Diarization failed unexpectedly"}