feat: Remove brouhaha-vad dependency to unlock Python 3.11+ support

dataspeech/gpu_enrichments/_brouhaha_compat.py

@@ -0,0 +1,84 @@
+"""Minimal stub of brouhaha model classes for checkpoint deserialization.
+
+The ylacombe/brouhaha-best checkpoint references brouhaha.models.CustomPyanNetModel.
+PyTorch Lightning needs these classes importable to load the checkpoint.
+This module provides just the model architecture definitions (no other brouhaha deps).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from pyannote.audio import Model
+from pyannote.audio.models.segmentation import PyanNet
+
+SNR_MIN = -15
+SNR_MAX = 80
+C50_MIN = -10
+C50_MAX = 60
+
+
+class ParametricSigmoid(nn.Module):
+    def __init__(self, alpha: float, beta: float) -> None:
+        super().__init__()
+        self.alpha = alpha
+        self.beta = beta
+
+    def forward(self, x: torch.Tensor):
+        return (self.beta - self.alpha) * F.sigmoid(x) + self.alpha
+
+
+class CustomClassifier(nn.Module):
+    def __init__(self, in_features, out_features: int) -> None:
+        super().__init__()
+        self.linears = nn.ModuleDict({
+            'vad': nn.Linear(in_features, out_features),
+            'snr': nn.Linear(in_features, 1),
+            'c50': nn.Linear(in_features, 1),
+        })
+
+    def forward(self, x: torch.Tensor):
+        out = dict()
+        for mode, linear in self.linears.items():
+            _output = linear(x)
+            out[mode] = _output
+        return out
+
+
+class CustomActivation(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.activations = nn.ModuleDict({
+            'vad': nn.Sigmoid(),
+            'snr': ParametricSigmoid(SNR_MAX, SNR_MIN),
+            'c50': ParametricSigmoid(C50_MAX, C50_MIN),
+        })
+
+    def forward(self, x: torch.Tensor):
+        out = list()
+        for mode, activation in self.activations.items():
+            _output = activation(x[mode])
+            out.append(_output)
+        out = torch.stack(out)
+        out = rearrange(out, "n b t o -> b t (n o)")
+        return out
+
+
+class RegressiveSegmentationModelMixin(Model):
+    def build(self):
+        nb_classif = len(set(self.specifications.classes) - set(['snr', 'c50']))
+        self.classifier = CustomClassifier(32 * 2, nb_classif)
+        self.activation = CustomActivation()
+
+
+class CustomPyanNetModel(RegressiveSegmentationModelMixin, PyanNet):
+    def build(self):
+        if self.hparams.linear["num_layers"] > 0:
+            in_features = self.hparams.linear["hidden_size"]
+        else:
+            in_features = self.hparams.lstm["hidden_size"] * (
+                2 if self.hparams.lstm["bidirectional"] else 1
+            )
+        nb_classif = len(set(self.specifications.classes) - set(['snr', 'c50']))
+        self.classifier = CustomClassifier(in_features, nb_classif)
+        self.activation = CustomActivation()

dataspeech/gpu_enrichments/_torbi_compat.py

@@ -0,0 +1,90 @@
+"""Compatibility shim for torbi on unsupported torch/CUDA versions.
+
+torbi ships prebuilt C++ binaries for specific torch+CUDA combinations.
+When no matching binary exists (e.g. torch 2.9+), import fails with
+FileNotFoundError. This module patches torbi to fall back to a pure-Python
+Viterbi implementation using librosa.sequence.viterbi — the same algorithm
+used by torbi's own reference implementation (torbi.reference.core).
+"""
+
+import sys
+
+
+def ensure_torbi():
+    """Ensure torbi is importable, patching with a librosa fallback if needed."""
+    if 'torbi' in sys.modules:
+        return
+
+    try:
+        import torbi  # noqa: F401
+        return
+    except FileNotFoundError:
+        pass
+
+    # torbi's C++ binary is unavailable for this torch/CUDA version.
+    # Create a minimal stub module providing only from_probabilities(),
+    # which is the sole function penn uses from torbi.
+    import types
+
+    import numpy as np
+    import torch
+
+    torbi_mod = types.ModuleType('torbi')
+    torbi_mod.__package__ = 'torbi'
+    sys.modules['torbi'] = torbi_mod
+
+    def from_probabilities(
+        observation,
+        batch_frames=None,
+        transition=None,
+        initial=None,
+        log_probs=False,
+        gpu=None,
+        num_threads=1,
+    ):
+        """Pure-Python Viterbi decoding via librosa (fallback for missing C++ binary).
+
+        Replicates the interface of torbi.core.from_probabilities but uses
+        librosa.sequence.viterbi for the actual decoding, matching the
+        algorithm in torbi.reference.core.from_probabilities.
+        """
+        import librosa
+
+        device = observation.device
+        batch, frames, states = observation.shape
+
+        # Convert to probability space for librosa
+        obs_probs = torch.exp(observation) if log_probs else observation
+        obs_np = obs_probs.to(torch.float32).cpu().numpy()
+
+        # Initial distribution (librosa expects probabilities, not log-probs)
+        if initial is None:
+            initial_np = np.full((states,), 1.0 / states, dtype=np.float32)
+        else:
+            init_t = torch.exp(initial) if log_probs else initial
+            initial_np = init_t.to(torch.float32).cpu().numpy()
+
+        # Transition matrix (librosa expects probabilities)
+        if transition is None:
+            trans_np = np.full(
+                (states, states), 1.0 / states, dtype=np.float32)
+        else:
+            trans_t = torch.exp(transition) if log_probs else transition
+            trans_np = trans_t.to(torch.float32).cpu().numpy()
+
+        # Decode each batch item
+        results = []
+        for i in range(batch):
+            n = batch_frames[i].item() if batch_frames is not None else frames
+            indices = librosa.sequence.viterbi(
+                obs_np[i, :n].T, trans_np, p_init=initial_np)
+            if n < frames:
+                padded = np.zeros(frames, dtype=np.int32)
+                padded[:n] = indices
+                indices = padded
+            results.append(torch.tensor(
+                indices.astype(np.int32), dtype=torch.int, device=device))
+
+        return torch.stack(results)
+
+    torbi_mod.from_probabilities = from_probabilities

dataspeech/gpu_enrichments/pitch.py

@@ -1,4 +1,6 @@
-import torch 
+import torch
+from ._torbi_compat import ensure_torbi
+ensure_torbi()
 import penn