"""XLSR-Mamba network, ported from https://github.com/swagshaw/XLSR-Mamba (MIT) ``model.py`` (SSLModel, Model) + ``mamba_blocks.py`` (MixerModel, create_block). DPDA-GA's ``nets/mamba.py`` + ``nets/mamba_blocks.py`` are byte-identical copies. Module / parameter names are identical to the source so the released state_dicts load strictly. Edits (inference only; see implementation-notes.md): 1. SSLModel builds the fairseq wav2vec 2.0 XLS-R 300M architecture from the base ``xlsr2_300m.pt`` config via merge_with_parent (cross-version fairseq fix, as in benchmarks/W2V2-AASIST/_net.py) instead of load_model_ensemble_and_task; the fine-tuned checkpoint then overwrites every weight. 2. extract_feat: removed the ``self.model.train()`` call (it fires on a device/dtype mismatch and would re-enable dropout) and the second, unused XLS-R forward (``layerresult``): the source ran XLS-R twice and discarded the second result. 3. Model takes the two hyper-parameters directly (emb_size, num_encoders) instead of argparse ``args``; no ``device`` / print. 4. Mamba mixer, Block, RMSNorm and the fused add+RMSNorm come from the dir-local ``_mamba_v1.py`` (pure PyTorch; mamba_ssm v1 subset) instead of the mamba_ssm CUDA / Triton kernels. Same math. 5. Only the code path the released checkpoints use is kept (MixerModel with if_bidirectional=True, fused_add_norm=True, rms_norm=True); BiBlock / BiMixerModel (unused by Model) are dropped. The quirk of that path is preserved verbatim: the final fused add+norm adds the *block input* ``hidden_states`` (not the last mixer output) to each direction's residual. """ from __future__ import annotations import os import sys from dataclasses import dataclass, field from functools import partial import numpy as np import torch import torch.nn as nn import torch.nn.functional as F _HERE = os.path.dirname(os.path.abspath(__file__)) if _HERE not in sys.path: sys.path.insert(0, _HERE) from _mamba_v1 import Block, Mamba, RMSNorm, rms_norm_fn # noqa: E402 WINDOW = 66800 # authors' training/eval cut (4.175 s @ 16 kHz) _XLSR_PATH = os.path.join(_HERE, "xlsr2_300m.pt") # symlink to W2V2-AASIST's base XLS-R def pad_repeat(x: np.ndarray, max_len: int = WINDOW) -> np.ndarray: """Source utils.pad: first max_len samples; shorter clips tiled to max_len.""" x = np.asarray(x, dtype=np.float32).reshape(-1) n = x.shape[0] if n >= max_len: return x[:max_len] reps = int(max_len / n) + 1 return np.tile(x, reps)[:max_len].astype(np.float32) def load_state(path: str) -> dict: if path.endswith(".safetensors"): from safetensors.torch import load_file return load_file(path, device="cpu") sd = torch.load(path, map_location="cpu", weights_only=True) return sd.get("state_dict", sd) @dataclass class MambaConfig: d_model: int = 64 n_layer: int = 6 vocab_size: int = 50277 ssm_cfg: dict = field(default_factory=dict) rms_norm: bool = True residual_in_fp32: bool = True fused_add_norm: bool = True pad_vocab_size_multiple: int = 8 class SSLModel(nn.Module): def __init__(self, cp_path: str = _XLSR_PATH): super().__init__() import fairseq from fairseq.dataclass.utils import merge_with_parent from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Config, Wav2Vec2Model state = fairseq.checkpoint_utils.load_checkpoint_to_cpu(cp_path) cfg = merge_with_parent(Wav2Vec2Config(), state["cfg"]["model"], remove_missing=True) ssl = Wav2Vec2Model.build_model(cfg, task=None) ssl.load_state_dict(state["model"], strict=True) self.model = ssl self.out_dim = 1024 def extract_feat(self, input_data): input_tmp = input_data[:, :, 0] if input_data.ndim == 3 else input_data return self.model(input_tmp, mask=False, features_only=True)["x"] def create_block(d_model, ssm_cfg=None, norm_epsilon=1e-5, rms_norm=False, residual_in_fp32=False, fused_add_norm=False, layer_idx=None, device=None, dtype=None): ssm_cfg = ssm_cfg or {} factory_kwargs = {"device": device, "dtype": dtype} mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs) norm_cls = partial(nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs) block = Block(d_model, mixer_cls, norm_cls=norm_cls, fused_add_norm=fused_add_norm, residual_in_fp32=residual_in_fp32) block.layer_idx = layer_idx return block class MixerModel(nn.Module): """mamba_blocks.MixerModel, bidirectional + fused_add_norm path only.""" def __init__(self, d_model, n_layer, ssm_cfg=None, norm_epsilon=1e-5, rms_norm=False, fused_add_norm=False, residual_in_fp32=False): super().__init__() assert rms_norm and fused_add_norm, "only the released configuration is ported" self.residual_in_fp32 = residual_in_fp32 self.fused_add_norm = fused_add_norm kw = dict(ssm_cfg=ssm_cfg, norm_epsilon=norm_epsilon, rms_norm=rms_norm, residual_in_fp32=residual_in_fp32, fused_add_norm=fused_add_norm) self.forward_layers = nn.ModuleList([create_block(d_model, layer_idx=i, **kw) for i in range(n_layer)]) self.backward_layers = nn.ModuleList([create_block(d_model, layer_idx=i, **kw) for i in range(n_layer)]) self.norm_f = RMSNorm(d_model, eps=norm_epsilon) self.f_attention_pool = nn.Linear(d_model, 1) self.b_attention_pool = nn.Linear(d_model, 1) self.LL = nn.Linear(d_model * 2, d_model) self.dropout = nn.Dropout(p=0.1) self.classifier = nn.Linear(d_model, 2) def forward(self, x): hidden_states = self.dropout(x) f_hidden_states = hidden_states b_hidden_states = hidden_states.flip([1]) f_residual, b_residual = None, None for layer in self.forward_layers: f_hidden_states, f_residual = layer(f_hidden_states, f_residual) for layer in self.backward_layers: b_hidden_states, b_residual = layer(b_hidden_states, b_residual) # source: "Use long-range features as the key for residual connections" -- # the fused add+norm takes `hidden_states` (the stack input), not the last mixer output. f_hidden_states = rms_norm_fn(hidden_states, self.norm_f.weight, self.norm_f.bias, eps=self.norm_f.eps, residual=f_residual, prenorm=False, residual_in_fp32=self.residual_in_fp32) b_hidden_states = rms_norm_fn(hidden_states, self.norm_f.weight, self.norm_f.bias, eps=self.norm_f.eps, residual=b_residual, prenorm=False, residual_in_fp32=self.residual_in_fp32) f_hidden_states = torch.matmul(F.softmax(self.f_attention_pool(f_hidden_states), dim=1) .transpose(-1, -2), f_hidden_states).squeeze(-2) b_hidden_states = torch.matmul(F.softmax(self.b_attention_pool(b_hidden_states), dim=1) .transpose(-1, -2), b_hidden_states).squeeze(-2) hidden_states = torch.cat((f_hidden_states, b_hidden_states), dim=1) hidden_states = self.LL(hidden_states) hidden_states = self.dropout(hidden_states) return self.classifier(hidden_states) class Model(nn.Module): """XLSR-Mamba: XLS-R 300M (fine-tuned) -> Linear 1024->emb -> BN+SELU -> bidirectional Mamba.""" def __init__(self, emb_size: int = 144, num_encoders: int = 12, cp_path: str = _XLSR_PATH): super().__init__() self.ssl_model = SSLModel(cp_path) self.LL = nn.Linear(1024, emb_size) self.first_bn = nn.BatchNorm2d(num_features=1) self.selu = nn.SELU(inplace=True) self.config = MambaConfig(d_model=emb_size, n_layer=num_encoders // 2) self.conformer = MixerModel(d_model=self.config.d_model, n_layer=self.config.n_layer, ssm_cfg=self.config.ssm_cfg, rms_norm=self.config.rms_norm, residual_in_fp32=self.config.residual_in_fp32, fused_add_norm=self.config.fused_add_norm) def set_ref_scan(self, on: bool) -> None: """Sequential selective_scan_ref (True) vs the chunked scan (False, default).""" for m in self.modules(): if isinstance(m, Mamba): m.use_ref_scan = on def forward(self, x): x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1)) x = self.LL(x_ssl_feat) # (bs, frames, emb) x = x.unsqueeze(dim=1) x = self.first_bn(x) x = self.selu(x) x = x.squeeze(dim=1) return self.conformer(x)