| """LaST-Net (Length-Aware Layer and Scale-Adaptive Temporal Network), inference only. |
| |
| Vendored from github.com/JunXue-tech/LaST-Net @ 3838f0c7108313f771f6c1e208ac8f1e3bde60b0 (MIT, Copyright (c) 2026 |
| LaST-Net authors; see LICENSE.source): |
| model.py duration_features, LengthCalibratedLayerSelector, TemporalResidualBlock, |
| LengthSelectiveKernelBlock, LengthPriorResidualExchange, AttentiveStatisticsPooling, LaSTNet |
| load_dataset.py crop_or_repeat (prefix crop / np.tile repeat), `available` seconds = min(len, samples) / 16000 |
| infer.py/eval.py score = log_softmax(logits)[:, 1] ("bonafide_log_score"; label bonafide = 1) |
| |
| Edits for the Arena wrapper (math unchanged): |
| 1. SSLModel builds the installed fairseq 0.12.2 Wav2Vec2Model from the vendored model cfg of xlsr2_300m.pt |
| (xlsr2_300m_model_cfg.json, merged onto Wav2Vec2Config defaults as in W2V2-AASIST/_net.py) instead of |
| load_model_ensemble_and_task([xlsr2_300m.pt]); the fine-tuned checkpoint overwrites every tensor (strict load). |
| extract_feat is verbatim: model(wav, mask=False, features_only=True, layer=23)["layer_results"]. |
| 2. LaSTNet.__init__ takes no ssl_path; forward returns the (B, 2) logits only (return_aux dropped). |
| 3. crop_np = crop_or_repeat (plus an empty-clip -> silence guard, which the source rejects with ValueError). |
| 4. nn.init.zeros_ calls dropped (every tensor is overwritten by the strict load); code reflowed. |
| Everything else is copied verbatim (the training-only layer-dropout branch is inert in eval()). |
| """ |
| import json |
| import math |
| import os |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| _DIR = os.path.dirname(os.path.abspath(__file__)) |
| XLSR_CFG = os.path.join(_DIR, "xlsr2_300m_model_cfg.json") |
| SAMPLE_RATE = 16000 |
| SECONDS = 4 |
| WINDOW = SECONDS * SAMPLE_RATE |
|
|
|
|
| def crop_np(wav, samples=WINDOW): |
| """load_dataset.crop_or_repeat + the `available` seconds that eval.py/infer.py pass to the model.""" |
| x = np.asarray(wav, dtype=np.float32).reshape(-1) |
| if x.shape[0] == 0: |
| return np.zeros(samples, dtype=np.float32), 0.0 |
| available = min(len(x), samples) / SAMPLE_RATE |
| if len(x) < samples: |
| x = np.tile(x, math.ceil(samples / len(x))) |
| return np.ascontiguousarray(x[:samples], dtype=np.float32), available |
|
|
|
|
| def build_xlsr(): |
| from fairseq.dataclass.utils import merge_with_parent |
| from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Config, Wav2Vec2Model |
| from omegaconf import OmegaConf |
|
|
| cfg = OmegaConf.create(json.load(open(XLSR_CFG))) |
| return Wav2Vec2Model.build_model(merge_with_parent(Wav2Vec2Config(), cfg, remove_missing=True), task=None) |
|
|
|
|
| class SSLModel(nn.Module): |
| """XLS-R frontend returning all 24 transformer-layer representations.""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.model = build_xlsr() |
|
|
| def extract_feat(self, wav): |
| output = self.model(wav, mask=False, features_only=True, layer=23) |
| return output["layer_results"] |
|
|
|
|
| def duration_features(wav, available_seconds=None): |
| requested = wav.new_full((len(wav),), wav.shape[-1] / 16000.0) |
| available = requested if available_seconds is None else available_seconds.to(wav) |
| available = available.clamp(0.001, 6.0) |
| return torch.stack( |
| ( |
| requested.clamp_max(6.0) / 6.0, |
| torch.log1p(requested.clamp_max(6.0)) / torch.log(requested.new_tensor(7.0)), |
| available / requested.clamp_min(0.001), |
| ), |
| dim=1, |
| ) |
|
|
|
|
| class LengthCalibratedLayerSelector(nn.Module): |
| """Fuse the 24 XLS-R layers using length and content-conditioned weights.""" |
|
|
| def __init__(self, dimension=1024, temperature=0.8, uniform_floor=0.05, layer_dropout=0.10): |
| super().__init__() |
| self.temperature = temperature |
| self.uniform_floor = uniform_floor |
| self.layer_dropout = layer_dropout |
| self.content_scorer = nn.Sequential( |
| nn.LayerNorm(2 * dimension), nn.Linear(2 * dimension, 128), nn.GELU(), nn.Linear(128, 1), |
| ) |
| self.length_scorer = nn.Sequential(nn.Linear(3, 64), nn.GELU(), nn.Linear(64, 24)) |
| self.layer_bias = nn.Parameter(torch.zeros(24)) |
|
|
| def forward(self, layers, lengths): |
| if len(layers) != 24: |
| raise RuntimeError("LaST-Net expects all 24 XLS-R layers") |
| sequences = [layer[0].transpose(0, 1) for layer in layers] |
| descriptors = [] |
| for sequence in sequences: |
| mean = sequence.mean(dim=1) |
| std = (sequence - mean.unsqueeze(1)).square().mean(dim=1).sqrt() |
| descriptors.append(torch.cat((mean, std), dim=-1)) |
| descriptors = torch.stack(descriptors, dim=1) |
|
|
| content_scores = self.content_scorer(descriptors.float()).squeeze(-1) |
| length_scores = self.length_scorer(lengths.float()) |
| scores = (self.layer_bias.unsqueeze(0) + content_scores + length_scores) / self.temperature |
|
|
| if self.training and self.layer_dropout > 0.0: |
| if len(scores) % 2 == 0: |
| mask = torch.rand(len(scores) // 2, 24, device=scores.device) < self.layer_dropout |
| mask = mask.repeat(2, 1) |
| else: |
| mask = torch.rand_like(scores) < self.layer_dropout |
| scores = scores.masked_fill(mask, -1e4) |
|
|
| conditional = F.softmax(scores, dim=1) |
| weights = (1.0 - self.uniform_floor) * conditional + self.uniform_floor / 24.0 |
| fused = sum(sequence * weights[:, index, None, None].to(sequence) for index, sequence in enumerate(sequences)) |
| return fused, weights |
|
|
|
|
| class TemporalResidualBlock(nn.Module): |
| def __init__(self, dimension, dilation): |
| super().__init__() |
| self.norm = nn.LayerNorm(dimension) |
| self.depthwise = nn.Conv1d(dimension, dimension, 5, padding=2 * dilation, dilation=dilation, |
| groups=dimension, bias=False) |
| self.pointwise = nn.Conv1d(dimension, 2 * dimension, 1) |
| self.output = nn.Linear(dimension, dimension) |
| self.scale = nn.Parameter(torch.full((dimension,), 0.1)) |
|
|
| def forward(self, x): |
| y = self.depthwise(self.norm(x).transpose(1, 2)) |
| value, gate = self.pointwise(y).chunk(2, dim=1) |
| update = (value * torch.sigmoid(gate)).transpose(1, 2) |
| return x + self.scale * self.output(update) |
|
|
|
|
| class LengthSelectiveKernelBlock(nn.Module): |
| """Length-guided local temporal filtering at three receptive fields.""" |
|
|
| def __init__(self, dimension): |
| super().__init__() |
| self.norm = nn.LayerNorm(dimension) |
| self.convs = nn.ModuleList([ |
| nn.Conv1d(dimension, dimension, 5, padding=2 * dilation, dilation=dilation, groups=dimension, bias=False) |
| for dilation in (1, 2, 4) |
| ]) |
| self.content_gate = nn.Linear(dimension, 2) |
| self.pointwise = nn.Conv1d(dimension, 2 * dimension, 1) |
| self.output = nn.Linear(dimension, dimension) |
| self.scale = nn.Parameter(torch.full((dimension,), 0.1)) |
|
|
| def forward(self, x, lengths): |
| normalized = self.norm(x) |
| channels = normalized.transpose(1, 2) |
| features = [conv(channels) for conv in self.convs] |
| duration = lengths[:, :1].float() |
| residual = 0.1 * torch.tanh(self.content_gate(normalized.mean(dim=1).float())) |
| gate2 = torch.sigmoid((duration - 0.35) / 0.15 + residual[:, :1]).to(x) |
| gate4 = torch.sigmoid((duration - 0.65) / 0.15 + residual[:, 1:]).to(x) |
| mixed = ( |
| features[0] + gate2.unsqueeze(-1) * features[1] + gate4.unsqueeze(-1) * features[2] |
| ) / (1.0 + gate2 + gate4).unsqueeze(-1) |
| value, gate = self.pointwise(mixed).chunk(2, dim=1) |
| update = (value * torch.sigmoid(gate)).transpose(1, 2) |
| return x + self.scale * self.output(update) |
|
|
|
|
| class LengthPriorResidualExchange(nn.Module): |
| """Process three temporal scales and exchange length-relevant evidence.""" |
|
|
| def __init__(self, dimension, selective_local=False): |
| super().__init__() |
| self.selective_local = selective_local |
| self.blocks = nn.ModuleList([ |
| LengthSelectiveKernelBlock(dimension) if selective_local else TemporalResidualBlock(dimension, 1), |
| TemporalResidualBlock(dimension, 2), |
| TemporalResidualBlock(dimension, 4), |
| ]) |
| self.edge_bias = nn.Parameter(torch.zeros(6)) |
| self.content_gate = nn.Linear(dimension, 6) |
| self.fusion_scale = nn.Parameter(torch.full((3, dimension), 0.1)) |
|
|
| @staticmethod |
| def resize(x, size): |
| if x.shape[1] == size: |
| return x |
| return F.interpolate(x.transpose(1, 2), size=size, mode="linear", align_corners=False).transpose(1, 2) |
|
|
| def forward(self, scales, lengths): |
| processed = [] |
| for index, (block, scale) in enumerate(zip(self.blocks, scales)): |
| if index == 0 and self.selective_local: |
| processed.append(block(scale, lengths)) |
| else: |
| processed.append(block(scale)) |
|
|
| duration = lengths[:, 0].clamp(0.0, 1.0) |
| source_prior = torch.stack((1.0 - duration, 4.0 * duration * (1.0 - duration), duration), dim=1) |
| content = sum(x.mean(dim=1) for x in processed) / 3.0 |
| content_residual = 0.1 * torch.tanh(self.content_gate(content.float())) |
|
|
| outputs = [] |
| all_weights = [] |
| edge_index = 0 |
| for target in range(3): |
| sources = [source for source in range(3) if source != target] |
| logits = [] |
| for local_index, source in enumerate(sources): |
| index = edge_index + local_index |
| logits.append(self.edge_bias[index] + 2.0 * (source_prior[:, source] - 0.5) |
| + content_residual[:, index]) |
| positive = F.softplus(torch.stack(logits, dim=1)).to(processed[target]) |
| weights = positive / positive.sum(dim=1, keepdim=True).clamp_min(1e-6) |
| incoming = sum( |
| weights[:, index, None, None] * self.resize(processed[source], processed[target].shape[1]) |
| for index, source in enumerate(sources) |
| ) |
| outputs.append(processed[target] + self.fusion_scale[target] * incoming) |
| all_weights.append(weights) |
| edge_index += 2 |
| return outputs, torch.cat(all_weights, dim=1) |
|
|
|
|
| class AttentiveStatisticsPooling(nn.Module): |
| def __init__(self, dimension=192): |
| super().__init__() |
| self.attention = nn.Sequential(nn.Linear(dimension, 128), nn.Tanh(), nn.Linear(128, 1)) |
|
|
| def forward(self, x): |
| weights = F.softmax(self.attention(x).float(), dim=1).to(x) |
| mean = (weights * x).sum(dim=1) |
| second = (weights * x.square()).sum(dim=1) |
| std = (second - mean.square()).clamp_min(1e-5).sqrt() |
| return torch.cat((mean, std), dim=-1) |
|
|
|
|
| class LaSTNet(nn.Module): |
| architecture = "LaST-Net" |
|
|
| |
| def __init__(self, temperature=0.8, uniform_floor=0.05, layer_dropout=0.10, dimension=192): |
| super().__init__() |
| self.ssl_model = SSLModel() |
| self.router = LengthCalibratedLayerSelector(temperature=temperature, uniform_floor=uniform_floor, |
| layer_dropout=layer_dropout) |
| self.input_projection = nn.Sequential(nn.Linear(1024, dimension), nn.LayerNorm(dimension)) |
| self.exchanges = nn.ModuleList([ |
| LengthPriorResidualExchange(dimension, selective_local=True), |
| LengthPriorResidualExchange(dimension, selective_local=False), |
| ]) |
| self.pools = nn.ModuleList([AttentiveStatisticsPooling(dimension) for _ in range(3)]) |
| self.classifier = nn.Sequential( |
| nn.LayerNorm(6 * dimension), nn.Linear(6 * dimension, 256), nn.GELU(), nn.Dropout(0.2), nn.Linear(256, 2), |
| ) |
|
|
| def forward(self, wav, available_seconds=None): |
| layers = self.ssl_model.extract_feat(wav) |
| lengths = duration_features(wav, available_seconds) |
| fused, _ = self.router(layers, lengths) |
| x0 = self.input_projection(fused) |
| scales = [ |
| x0, |
| F.avg_pool1d(x0.transpose(1, 2), 2, 2).transpose(1, 2), |
| F.avg_pool1d(x0.transpose(1, 2), 4, 4).transpose(1, 2), |
| ] |
| for exchange in self.exchanges: |
| scales, _ = exchange(scales, lengths) |
| embedding = torch.cat([pool(scale) for pool, scale in zip(self.pools, scales)], dim=-1) |
| return self.classifier(embedding) |
|
|
