Fix stereo duration handling and robust ECAPA model loading

app/pipeline.py

@@ -1,6 +1,6 @@
 """
 Speaker Diarization Pipeline
-Combines: Voice Activity Detection → Segmentation → ECAPA-TDNN Embeddings → AHC Clustering
+Combines: Voice Activity Detection -> Segmentation -> ECAPA-TDNN Embeddings -> AHC Clustering
 """
 
 import torch
@@ -97,6 +97,16 @@ class DiarizationPipeline:
             return "cuda" if torch.cuda.is_available() else "cpu"
         return device
 
+    def _to_mono_1d(self, audio: torch.Tensor) -> torch.Tensor:
+        """Convert waveform to a mono 1D tensor for duration and preprocessing."""
+        if audio.dim() == 1:
+            return audio
+        if audio.dim() >= 2:
+            if audio.shape[0] == 1:
+                return audio[0]
+            return audio.mean(dim=0)
+        return audio.reshape(-1)
+
     def _load_vad(self):
         if self._vad_pipeline is not None:
             return
@@ -120,13 +130,13 @@ class DiarizationPipeline:
         frame_samples = int(frame_duration * self.SAMPLE_RATE)
         audio_np = audio.numpy()
         frames = [
-            audio_np[i : i + frame_samples]
+            audio_np[i: i + frame_samples]
             for i in range(0, len(audio_np) - frame_samples, frame_samples)
         ]
 
         energies_db = []
-        for f in frames:
-            rms = np.sqrt(np.mean(f ** 2) + 1e-10)
+        for frame in frames:
+            rms = np.sqrt(np.mean(frame ** 2) + 1e-10)
             energies_db.append(20 * np.log10(rms))
 
         is_speech = np.array(energies_db) > threshold_db
@@ -198,25 +208,38 @@ class DiarizationPipeline:
     ) -> DiarizationResult:
         """Run full diarization pipeline on audio."""
         import time
+
         t_start = time.time()
 
         if isinstance(audio, (str, Path)):
             waveform, sample_rate = self.load_audio(audio)
-            audio_tensor = waveform.squeeze(0)
+            audio_tensor = self._to_mono_1d(waveform)
         else:
             assert sample_rate is not None, "sample_rate required when passing tensor"
-            audio_tensor = audio.squeeze(0) if audio.dim() > 1 else audio
+            audio_tensor = self._to_mono_1d(audio)
 
-        audio_duration = len(audio_tensor) / sample_rate
+        num_samples = int(audio_tensor.numel())
+        audio_duration = num_samples / float(sample_rate)
         logger.info(f"Processing {audio_duration:.1f}s audio at {sample_rate}Hz")
 
+        if num_samples == 0:
+            logger.warning("Received empty audio input.")
+            return DiarizationResult(
+                segments=[],
+                num_speakers=0,
+                audio_duration=0.0,
+                processing_time=time.time() - t_start,
+                sample_rate=sample_rate,
+            )
+
         processed = self.embedder.preprocess_audio(audio_tensor, sample_rate)
 
         speech_regions = self._get_speech_regions(processed)
         if not speech_regions:
             logger.warning("No speech detected in audio.")
             return DiarizationResult(
-                segments=[], num_speakers=0,
+                segments=[],
+                num_speakers=0,
                 audio_duration=audio_duration,
                 processing_time=time.time() - t_start,
                 sample_rate=sample_rate,
@@ -232,7 +255,8 @@ class DiarizationPipeline:
         if len(embeddings) == 0:
             logger.warning("No valid embeddings extracted.")
             return DiarizationResult(
-                segments=[], num_speakers=0,
+                segments=[],
+                num_speakers=0,
                 audio_duration=audio_duration,
                 processing_time=time.time() - t_start,
                 sample_rate=sample_rate,
@@ -245,11 +269,7 @@ class DiarizationPipeline:
 
         speaker_names = {i: f"SPEAKER_{i:02d}" for i in range(self.max_speakers)}
         segments = [
-            DiarizationSegment(
-                start=start,
-                end=end,
-                speaker=speaker_names[spk_id],
-            )
+            DiarizationSegment(start=start, end=end, speaker=speaker_names[spk_id])
             for start, end, spk_id in merged
         ]
 
@@ -268,3 +288,4 @@ class DiarizationPipeline:
             processing_time=processing_time,
             sample_rate=sample_rate,
         )
+

models/embedder.py

@@ -1,10 +1,9 @@
-"""
+"""
 Speaker Embedding Extraction using ECAPA-TDNN architecture via SpeechBrain.
 Handles audio preprocessing, feature extraction, and L2-normalized embeddings.
 """
 
 import inspect
-import shutil
 from pathlib import Path
 from typing import Union, List, Tuple
 
@@ -36,24 +35,36 @@ class EcapaTDNNEmbedder:
             return "cuda" if torch.cuda.is_available() else "cpu"
         return device
 
+    def _build_hparams_kwargs(self, encoder_cls, savedir: Path, hf_cache: Path) -> dict:
+        kwargs = {
+            "source": self.MODEL_SOURCE,
+            "savedir": str(savedir),
+            "run_opts": {"device": self.device},
+        }
+
+        sig = inspect.signature(encoder_cls.from_hparams)
+        if "huggingface_cache_dir" in sig.parameters:
+            kwargs["huggingface_cache_dir"] = str(hf_cache)
+        if "local_strategy" in sig.parameters:
+            try:
+                from speechbrain.utils.fetching import LocalStrategy
+
+                kwargs["local_strategy"] = LocalStrategy.COPY
+            except Exception:
+                pass
+
+        return kwargs
+
     def _load_model(self):
         if self._model is not None:
             return
 
         try:
-            import speechbrain.utils.fetching as _fetching
-            from speechbrain.utils.fetching import LocalStrategy
-            from speechbrain.inference.classifiers import EncoderClassifier
-
-            def _patched_link(src, dst, local_strategy):
-                dst_path = Path(dst)
-                src_path = Path(src)
-                dst_path.parent.mkdir(parents=True, exist_ok=True)
-                if dst_path.exists() or dst_path.is_symlink():
-                    dst_path.unlink()
-                shutil.copy2(str(src_path), str(dst_path))
-
-            _fetching.link_with_strategy = _patched_link
+            try:
+                from speechbrain.inference.classifiers import EncoderClassifier
+            except ImportError:
+                # Backward compatibility with older SpeechBrain versions.
+                from speechbrain.pretrained import EncoderClassifier
 
             savedir = self.cache_dir / "ecapa_tdnn"
             hf_cache = self.cache_dir / "hf_cache"
@@ -63,23 +74,28 @@ class EcapaTDNNEmbedder:
             logger.info(f"Loading ECAPA-TDNN from {self.MODEL_SOURCE}...")
             logger.info(f"Savedir: {savedir}, exists: {savedir.exists()}")
 
-            kwargs = {
-                "source": self.MODEL_SOURCE,
-                "savedir": str(savedir),
-                "run_opts": {"device": self.device},
-            }
+            kwargs = self._build_hparams_kwargs(EncoderClassifier, savedir, hf_cache)
+            model = EncoderClassifier.from_hparams(**kwargs)
 
-            sig = inspect.signature(EncoderClassifier.from_hparams)
-            if "huggingface_cache_dir" in sig.parameters:
-                kwargs["huggingface_cache_dir"] = str(hf_cache)
-            if "local_strategy" in sig.parameters:
-                kwargs["local_strategy"] = LocalStrategy.COPY
+            if model is None:
+                # Some SpeechBrain/HF hub combinations ignore optional kwargs.
+                logger.warning("ECAPA load returned None; retrying with minimal from_hparams kwargs.")
+                model = EncoderClassifier.from_hparams(
+                    source=self.MODEL_SOURCE,
+                    savedir=str(savedir),
+                    run_opts={"device": self.device},
+                )
 
-            self._model = EncoderClassifier.from_hparams(**kwargs)
+            if model is None:
+                raise RuntimeError("EncoderClassifier.from_hparams returned None")
+
+            self._model = model
             self._model.eval()
             logger.success("ECAPA-TDNN model loaded successfully.")
         except ImportError as exc:
             raise ImportError("SpeechBrain not installed.") from exc
+        except Exception as exc:
+            raise RuntimeError(f"Failed to load ECAPA-TDNN model: {exc}") from exc
 
     def preprocess_audio(
         self, audio: Union[np.ndarray, torch.Tensor], sample_rate: int
@@ -157,3 +173,4 @@ class EcapaTDNNEmbedder:
             return np.empty((0, self.EMBEDDING_DIM)), []
 
         return np.stack(embeddings), valid_segments
+