| """Small causal Nemotron-L15 -> ECAPA-192 predictors; no ASR/teacher weights.""" |
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class SpeakerPredictor(nn.Module): |
| def __init__(self, architecture="gru", width=128, input_dim=1024, output_dim=192, centered=False): |
| super().__init__() |
| self.architecture = architecture |
| self.width = width |
| self.centered = centered |
| self.norm = nn.LayerNorm(input_dim) |
| if architecture == "mlp": |
| self.encoder = nn.Sequential(nn.Linear(input_dim * 2, width), nn.GELU(), |
| nn.Linear(width, width), nn.GELU()) |
| elif architecture == "gru": |
| self.project = nn.Sequential(nn.Linear(input_dim, width), nn.GELU()) |
| self.encoder = nn.GRU(width, width, batch_first=True) |
| elif architecture == "tcn": |
| self.project = nn.Sequential(nn.Linear(input_dim, width), nn.GELU()) |
| self.encoder = nn.ModuleList([ |
| nn.Sequential(nn.Conv1d(width, width, 3, groups=width, dilation=d), |
| nn.GELU(), nn.Conv1d(width, width, 1), nn.GELU()) |
| for d in (1, 2, 4)]) |
| else: |
| raise ValueError(architecture) |
| self.embedding = nn.Linear(width, output_dim) |
| self.activity = nn.Linear(width, 3) |
| self.register_buffer("teacher_mean",torch.zeros(output_dim)) |
| self.register_buffer("teacher_radius",torch.ones(())) |
|
|
| def forward(self, features, lengths=None): |
| x = self.norm(features) |
| if self.architecture == "mlp": |
| h = self.encoder(torch.cat([x.mean(1), x.std(1, unbiased=False)], -1)) |
| elif self.architecture == "gru": |
| h, _ = self.encoder(self.project(x)) |
| h = h[:, -1] |
| else: |
| h = self.project(x).transpose(1, 2) |
| for dilation, layer in zip((1, 2, 4), self.encoder): |
| h = h + layer(F.pad(h, (2 * dilation, 0))) |
| h = h[:, :, -1] |
| direction = F.normalize(self.embedding(h),dim=-1) |
| embedding = F.normalize(self.teacher_mean+self.teacher_radius*direction,dim=-1) if self.centered else direction |
| return embedding, self.activity(h) |
|
|
|
|
| def load_state(model, state): |
| """Compatibility with the first uncentered pilot checkpoints.""" |
| missing, unexpected = model.load_state_dict(state,strict=False) |
| if unexpected or set(missing)-{"teacher_mean","teacher_radius"}: |
| raise ValueError(f"Unexpected checkpoint mismatch: missing={missing}, unexpected={unexpected}") |
|
|
|
|
| def parameter_count(model): |
| return sum(p.numel() for p in model.parameters() if p.requires_grad) |
|
|
