Update pipeline.py

pipeline.py

@@ -24,8 +24,8 @@ logger = logging.getLogger(__name__)
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 HF_TOKEN = os.environ.get("HF_TOKEN", "")
 WHISPER_MODEL = os.environ.get("WHISPER_MODEL", "base")
+NUM_SPEAKERS = int(os.environ.get("NUM_SPEAKERS", "0"))  # 0 = auto-detect
 
-# Login once at import time — works with all pyannote/huggingface_hub versions
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
@@ -80,9 +80,9 @@ class SpeechAnnotationPipeline:
                     "HF_TOKEN not set. Add it in Space Settings > Variables and secrets."
                 )
             logger.info("Loading pyannote diarization pipeline...")
-            # No token param needed — login() at top handles auth
             self._diarization_pipeline = PyannotePipeline.from_pretrained(
-                "pyannote/speaker-diarization-3.1"
+                "pyannote/speaker-diarization-3.1",
+                use_auth_token=HF_TOKEN,
             )
             if DEVICE == "cuda":
                 self._diarization_pipeline = self._diarization_pipeline.to(torch.device("cuda"))
@@ -94,10 +94,22 @@ class SpeechAnnotationPipeline:
         result = model.transcribe(audio_path, verbose=False)
         return result["segments"]
 
-    def _diarize(self, audio_path):
+    def _diarize(self, audio_path, num_speakers=0, min_speakers=None, max_speakers=None):
         logger.info("Step 2/2 - Diarizing speakers...")
         pipeline = self._load_diarization()
-        diarization = pipeline(audio_path)
+
+        # Build diarization params
+        params = {}
+        if num_speakers and num_speakers > 0:
+            params["num_speakers"] = num_speakers
+            logger.info(f"Using fixed num_speakers={num_speakers}")
+        else:
+            if min_speakers:
+                params["min_speakers"] = min_speakers
+            if max_speakers:
+                params["max_speakers"] = max_speakers
+
+        diarization = pipeline(audio_path, **params)
         segments = []
         for turn, _, speaker in diarization.itertracks(yield_label=True):
             segments.append({"start": turn.start, "end": turn.end, "speaker": speaker})
@@ -124,10 +136,10 @@ class SpeechAnnotationPipeline:
             ))
         return results
 
-    def process(self, audio_path):
+    def process(self, audio_path, num_speakers=0, min_speakers=None, max_speakers=None):
         t0 = time.time()
         transcript = self._transcribe(audio_path)
-        diarization = self._diarize(audio_path)
+        diarization = self._diarize(audio_path, num_speakers, min_speakers, max_speakers)
         annotations = self._assign_speakers(transcript, diarization)
         logger.info(f"Done in {time.time()-t0:.1f}s — {len(annotations)} segments")
         return annotations