| """Speaker encoders — the "voice fingerprint" half of Vanta.
|
|
|
| Two implementations behind one contract: (B, T) waveform at 16 kHz in, (B, 192)
|
| embedding out.
|
|
|
| TrainedSpeakerEncoder our own ECAPA-TDNN, trained from scratch by
|
| scripts/train_speaker_encoder.py. This is what
|
| production runs.
|
| SpeakerEncoder speechbrain's VoxCeleb-trained ECAPA. Kept as the
|
| baseline `compare_encoders.py` measures against, and
|
| as the encoder for the inactive SepFormer backend.
|
|
|
| Use `SpeakerEncoder.load(checkpoint=...)` rather than constructing either
|
| directly — passing a checkpoint selects ours, omitting it selects the
|
| pretrained one.
|
|
|
| Both are frozen during TSE training: fine-tuning a speaker encoder while the
|
| separator learns tends to destabilize the identity space, and we want the
|
| fingerprint to stay recognizable.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| from pathlib import Path
|
|
|
| import torch
|
| import torch.nn as nn
|
|
|
| from vanta.config import DATA_DIR
|
|
|
| ECAPA_EMBED_DIM = 192
|
|
|
|
|
| class SpeakerEncoder(nn.Module):
|
| def __init__(
|
| self,
|
| savedir: Path | None = None,
|
| freeze: bool = True,
|
| run_opts: dict | None = None,
|
| ):
|
| super().__init__()
|
|
|
|
|
| from speechbrain.inference.speaker import EncoderClassifier
|
| from speechbrain.utils.fetching import LocalStrategy
|
|
|
| savedir = savedir or (DATA_DIR / "_models" / "ecapa_voxceleb")
|
| savedir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
| self.encoder = EncoderClassifier.from_hparams(
|
| source="speechbrain/spkrec-ecapa-voxceleb",
|
| savedir=str(savedir),
|
| run_opts=run_opts or {},
|
| local_strategy=LocalStrategy.COPY,
|
| )
|
| self.embed_dim = ECAPA_EMBED_DIM
|
| self.freeze = freeze
|
| if freeze:
|
| for p in self.encoder.parameters():
|
| p.requires_grad_(False)
|
| self.encoder.eval()
|
|
|
| def forward(self, wav: torch.Tensor) -> torch.Tensor:
|
| """wav: (B, T) at 16 kHz. Returns (B, 192) speaker embeddings."""
|
| if wav.dim() == 3:
|
| wav = wav.squeeze(1)
|
|
|
| if self.freeze:
|
| with torch.no_grad():
|
| emb = self.encoder.encode_batch(wav)
|
| else:
|
| emb = self.encoder.encode_batch(wav)
|
| return emb.squeeze(1)
|
|
|
| def train(self, mode: bool = True):
|
|
|
| super().train(mode)
|
| if self.freeze:
|
| self.encoder.eval()
|
| return self
|
|
|
| @staticmethod
|
| def load(
|
| checkpoint: "Path | str | None" = None,
|
| freeze: bool = True,
|
| run_opts: dict | None = None,
|
| ) -> nn.Module:
|
| """Return a speaker encoder: ours if a checkpoint is given, else pretrained.
|
|
|
| Both expose the same contract — (B, T) waveform at 16 kHz in, (B, 192)
|
| embedding out — so the separator is indifferent to which one it gets.
|
| Note that a separator trained against one encoder's embedding space is
|
| not automatically compatible with another's; see TrainedSpeakerEncoder.
|
| """
|
| if checkpoint is None:
|
| return SpeakerEncoder(freeze=freeze, run_opts=run_opts)
|
| return TrainedSpeakerEncoder(checkpoint, freeze=freeze)
|
|
|
| def _apply(self, fn, *args, **kwargs):
|
|
|
|
|
|
|
|
|
|
|
| out = super()._apply(fn, *args, **kwargs)
|
| try:
|
| dev = next(self.encoder.parameters()).device
|
| self.encoder.device = dev
|
| if hasattr(self.encoder, "mods"):
|
| self.encoder.mods.to(dev)
|
| except StopIteration:
|
| pass
|
| return out
|
|
|
|
|
| class TrainedSpeakerEncoder(nn.Module):
|
| """Our own ECAPA-TDNN, trained from scratch (scripts/train_speaker_encoder.py).
|
|
|
| Drop-in replacement for the pretrained wrapper above: same (B, T) -> (B, 192)
|
| contract, so `Vanta` can take either.
|
|
|
| Compatibility caveat, and it matters: a separator learns to read the specific
|
| embedding space it trained against. Both encoders emit 192 numbers, but the
|
| *meaning* of those numbers differs — swapping encoders under a separator that
|
| never saw this one is a domain shift, not a like-for-like exchange. Expect to
|
| re-train (or at least fine-tune) the separator after switching.
|
| """
|
|
|
| def __init__(self, checkpoint: "Path | str", freeze: bool = True):
|
| super().__init__()
|
| from vanta.models.ecapa_tdnn import EcapaTdnn
|
|
|
| ck = torch.load(str(checkpoint), map_location="cpu", weights_only=False)
|
| self.encoder = EcapaTdnn(
|
| embed_dim=ck.get("embed_dim", ECAPA_EMBED_DIM),
|
| channels=ck.get("channels", 512),
|
| )
|
| self.encoder.load_state_dict(ck["model_state"])
|
| self.embed_dim = ck.get("embed_dim", ECAPA_EMBED_DIM)
|
| del ck
|
| import gc
|
| gc.collect()
|
| self.freeze = freeze
|
| if freeze:
|
| for p in self.encoder.parameters():
|
| p.requires_grad_(False)
|
| self.encoder.eval()
|
|
|
| def forward(self, wav: torch.Tensor) -> torch.Tensor:
|
| if wav.dim() == 3:
|
| wav = wav.squeeze(1)
|
| if self.freeze:
|
| with torch.no_grad():
|
| return self.encoder(wav)
|
| return self.encoder(wav)
|
|
|
| def train(self, mode: bool = True):
|
| super().train(mode)
|
| if self.freeze:
|
| self.encoder.eval()
|
| return self
|
|
|
