| """SSL-Mamba network (inference-only port). |
| |
| Source: github.com/hoanmyTran/spoofing_ssl_mamba @ 9afc33de7f74a9b9b3423a9028c10e1d18c1c6b8 (MIT), |
| ``model.py`` ``SSLClassifier``. Checkpoint: HF hoanmyTran/ssl_spoofing_mamba @ |
| 4e72dfcef4f03460e08f68a4a5fbcf3705a3414e ``interspeech2025.pt`` (Apache-2.0; flat state_dict), |
| converted once to ``model.safetensors`` with the training-only ``loss_fn.weight`` dropped. |
| |
| Edits vs the source (all inference-only; module names unchanged so keys load strictly): |
| 1. pytorch_lightning / torchmetrics scaffolding, loss and training/eval loops removed. |
| 2. XLS-R 300M built from the bundled config (``xlsr_300m_config.json`` = the |
| facebook/wav2vec2-xls-r-300m config) instead of ``from_pretrained``: no Hub weight |
| download; the checkpoint overwrites every weight (the front-end is fully fine-tuned). |
| 3. ``mamba_ssm.Mamba`` -> ``_mamba_v1.Mamba`` (pure PyTorch, same parameter names, same |
| math; chunked closed-form selective scan, see that file). mamba_ssm / causal_conv1d are |
| not installed and Triton is unreliable on sm_120. |
| 4. ``multiconv_cgmlp`` espnet imports -> ``_espnet_shim`` (LayerNorm eps 1e-12, activations); |
| ``pooling.py`` (MultiHeadAttentionPooling) vendored from the authors' sibling repo |
| (dissimilarity_deepfake_detection, as in benchmarks/SSL-Feature-Gating/), not in this repo. |
| 5. ``nn.SiLU(inplace=True)`` / ``nn.SELU(inplace=True)`` -> non-inplace (export-friendly; |
| identical values). |
| Forward is otherwise verbatim: 25 hidden states -> Linear(1024,128) -> SiLU -> sum over layers |
| -> MultiConv cgMLP (7/15/23/31) -> 12 x [x = x + LayerNorm(Mamba(x))] (unidirectional) |
| -> MHAP (4 heads, mean+std) -> Linear 256->128 -> SELU -> Linear 128->2. |
| """ |
| from __future__ import annotations |
|
|
| import os |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from transformers import Wav2Vec2Config, Wav2Vec2Model |
|
|
| from _mamba_v1 import Mamba |
| from multiconv_cgmlp import MultiConvolutionalGatingMLP |
| from pooling import MultiHeadAttentionPooling |
|
|
| _DIR = os.path.dirname(os.path.abspath(__file__)) |
| _CFG = os.path.join(_DIR, "xlsr_300m_config.json") |
| WINDOW = 64600 |
|
|
|
|
| def pad_repeat(x: np.ndarray, max_len: int = WINDOW) -> np.ndarray: |
| """Source dataset.pad: first max_len samples; shorter clips tile-repeated.""" |
| x = np.asarray(x, dtype=np.float32).reshape(-1) |
| n = x.shape[0] |
| if n >= max_len: |
| return x[:max_len] |
| return np.tile(x, max_len // n + 1)[:max_len].astype(np.float32) |
|
|
|
|
| class SSLMambaNet(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.feature_extractor = Wav2Vec2Model(Wav2Vec2Config.from_json_file(_CFG)) |
| self.cgmlp = MultiConvolutionalGatingMLP( |
| size=128, linear_units=1024, arch_type="concat_fusion", kernel_sizes="7,15,23,31", |
| merge_conv_kernel=31, use_non_linear=True, dropout_rate=0.1, |
| use_linear_after_conv=True, activation="silu", gate_activation="silu") |
| self.feature_projection = nn.Linear(1024, 128) |
| self.blocks = nn.ModuleList( |
| nn.Sequential(Mamba(d_model=128, d_state=16, d_conv=4, expand=2), nn.LayerNorm(128)) |
| for _ in range(12)) |
| self.last_pooling = MultiHeadAttentionPooling(128) |
| self.silu = nn.SiLU() |
| self.classifier = nn.Sequential(nn.Linear(256, 128), nn.SELU(), nn.Linear(128, 2)) |
|
|
| def set_ref_scan(self, on: bool) -> None: |
| for m in self.modules(): |
| if isinstance(m, Mamba): |
| m.use_ref_scan = on |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.feature_extractor(x, output_hidden_states=True) |
| hs = torch.stack(x.hidden_states, dim=1) |
| hs = self.silu(self.feature_projection(hs)) |
| x = torch.sum(hs, 1) |
| x = self.cgmlp(x) |
| for b in self.blocks: |
| x = x + b(x) |
| x = self.last_pooling(x.permute(0, 2, 1)) |
| return self.classifier(x.squeeze(-1)) |
|
|
|
|
| def load_state(path: str) -> dict: |
| from safetensors.torch import load_file |
| return load_file(path) |
|
|
