| """Codecfake W2V2-AASIST (Xie et al., arXiv 2405.04880) — inference port. |
| |
| Source: https://github.com/xieyuankun/Codecfake @ 72c32e2840eef95fbeade349425fa9c5391fa11e (NO LICENSE). |
| model.py: GraphAttentionLayer, HtrgGraphAttentionLayer, GraphPool, Residual_block, W2VAASIST |
| generate_score.py / preprocess.py: pad_dataset (first 64600 samples, repeat-tile), HF |
| Wav2Vec2FeatureExtractor (do_normalize) + frozen facebook/wav2vec2-xls-r-300m, hidden_states[5] |
| -> unsqueeze(0) -> transpose(2,3) = (B, 1, 1024, T) -> W2VAASIST -> (last_hidden, logits). |
| |
| Edits vs the source (inference-only, maths unchanged): |
| * removed `print(x.shape)` in W2VAASIST.forward and the unused imports (pytorch_model_summary, ...); |
| * Residual_block: the source computes bn1->selu for non-first blocks and then applies conv1 to `x`, |
| discarding that result; the dead computation is dropped (identical output; bn1 weights kept so the |
| checkpoint loads strictly); |
| * the frozen XLS-R front-end is truncated to its first 5 transformer layers: hidden_states[5] of the |
| 24-layer HF model is the output of layer index 4 *before* any final LayerNorm, so `XLSRLayer5` |
| runs feature encoder -> feature projection -> pos-conv -> layers[0..4] and returns that tensor |
| (bit-identical to the full model's hidden_states[5], verified in implementation-notes.md); |
| * the HF feature-extractor normalisation is reproduced per row in torch (zero mean, unit population |
| variance, eps 1e-7), so a whole batch is normalised on the GPU; |
| * checkpoints are pickled whole modules (torch.save(model)); `load_pickled_state` unpickles them with an |
| allow-list unpickler that maps the source's `model.*` classes to a bare nn.Module stub (no source |
| code runs) and returns the module's state_dict. |
| """ |
| from __future__ import annotations |
|
|
| import pickle |
| from typing import Union |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| WINDOW = 64600 |
| XLSR_REPO = "facebook/wav2vec2-xls-r-300m" |
| XLSR_REV = "1a640f32ac3e39899438a2931f9924c02f080a54" |
| N_LAYERS = 5 |
|
|
|
|
| def pad_repeat(x: np.ndarray, cut: int = WINDOW) -> np.ndarray: |
| """Source pad_dataset: first `cut` samples; shorter clips tiled (int(cut/len)+1 repeats) and cropped.""" |
| n = x.shape[0] |
| if n >= cut: |
| return x[:cut] |
| return np.tile(x, int(cut / n) + 1)[:cut] |
|
|
|
|
| def normalize(x: torch.Tensor) -> torch.Tensor: |
| """Wav2Vec2FeatureExtractor zero_mean_unit_var_norm, per row: (x - mean) / sqrt(var + 1e-7).""" |
| m = x.mean(dim=-1, keepdim=True) |
| v = x.var(dim=-1, keepdim=True, unbiased=False) |
| return (x - m) / torch.sqrt(v + 1e-7) |
|
|
|
|
| |
| class XLSRLayer5(nn.Module): |
| """Frozen XLS-R 300M, first 5 transformer layers; returns HF hidden_states[5] (B, T, 1024).""" |
|
|
| def __init__(self, w2v): |
| super().__init__() |
| self.feature_extractor = w2v.feature_extractor |
| self.feature_projection = w2v.feature_projection |
| self.pos_conv_embed = w2v.encoder.pos_conv_embed |
| self.layers = nn.ModuleList(list(w2v.encoder.layers)[:N_LAYERS]) |
|
|
| def forward(self, wav: torch.Tensor) -> torch.Tensor: |
| h = self.feature_extractor(wav).transpose(1, 2) |
| h, _ = self.feature_projection(h) |
| h = h + self.pos_conv_embed(h) |
| for layer in self.layers: |
| h = layer(h, attention_mask=None)[0] |
| return h |
|
|
|
|
| def build_xlsr(local_dir: str | None = None) -> XLSRLayer5: |
| """facebook/wav2vec2-xls-r-300m @ XLSR_REV, loaded with only N_LAYERS transformer layers (asserted).""" |
| from transformers import Wav2Vec2Config, Wav2Vec2Model |
| cfg = Wav2Vec2Config.from_pretrained(XLSR_REPO, revision=XLSR_REV) |
| assert cfg.do_stable_layer_norm and cfg.num_hidden_layers == 24 and cfg.hidden_size == 1024 |
| cfg.num_hidden_layers = N_LAYERS |
| w2v, info = Wav2Vec2Model.from_pretrained(XLSR_REPO, revision=XLSR_REV, config=cfg, |
| attn_implementation="eager", output_loading_info=True) |
| assert not info["missing_keys"], info["missing_keys"] |
| unexp = [k.removeprefix("wav2vec2.") for k in info["unexpected_keys"]] |
| bad = [k for k in unexp |
| if not (k.startswith(("encoder.layers.", "quantizer.", "project_q.", "project_hid.")) |
| and not any(k.startswith(f"encoder.layers.{i}.") for i in range(N_LAYERS)))] |
| assert sum(k.startswith("encoder.layers.") for k in unexp) == 16 * (24 - N_LAYERS) |
| assert not bad, bad |
| return XLSRLayer5(w2v) |
|
|
|
|
| |
| class GraphAttentionLayer(nn.Module): |
| def __init__(self, in_dim, out_dim, **kwargs): |
| super().__init__() |
| self.att_proj = nn.Linear(in_dim, out_dim) |
| self.att_weight = self._init_new_params(out_dim, 1) |
| self.proj_with_att = nn.Linear(in_dim, out_dim) |
| self.proj_without_att = nn.Linear(in_dim, out_dim) |
| self.bn = nn.BatchNorm1d(out_dim) |
| self.input_drop = nn.Dropout(p=0.2) |
| self.act = nn.SELU(inplace=True) |
| self.temp = kwargs.get("temperature", 1.) |
|
|
| def forward(self, x): |
| x = self.input_drop(x) |
| att_map = self._derive_att_map(x) |
| x = self._project(x, att_map) |
| x = self._apply_BN(x) |
| return self.act(x) |
|
|
| def _pairwise_mul_nodes(self, x): |
| nb_nodes = x.size(1) |
| x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1) |
| return x * x.transpose(1, 2) |
|
|
| def _derive_att_map(self, x): |
| att_map = self._pairwise_mul_nodes(x) |
| att_map = torch.tanh(self.att_proj(att_map)) |
| att_map = torch.matmul(att_map, self.att_weight) |
| att_map = att_map / self.temp |
| return F.softmax(att_map, dim=-2) |
|
|
| def _project(self, x, att_map): |
| return self.proj_with_att(torch.matmul(att_map.squeeze(-1), x)) + self.proj_without_att(x) |
|
|
| def _apply_BN(self, x): |
| org_size = x.size() |
| return self.bn(x.view(-1, org_size[-1])).view(org_size) |
|
|
| def _init_new_params(self, *size): |
| out = nn.Parameter(torch.FloatTensor(*size)) |
| nn.init.xavier_normal_(out) |
| return out |
|
|
|
|
| class HtrgGraphAttentionLayer(nn.Module): |
| def __init__(self, in_dim, out_dim, **kwargs): |
| super().__init__() |
| self.proj_type1 = nn.Linear(in_dim, in_dim) |
| self.proj_type2 = nn.Linear(in_dim, in_dim) |
| self.att_proj = nn.Linear(in_dim, out_dim) |
| self.att_projM = nn.Linear(in_dim, out_dim) |
| self.att_weight11 = self._init_new_params(out_dim, 1) |
| self.att_weight22 = self._init_new_params(out_dim, 1) |
| self.att_weight12 = self._init_new_params(out_dim, 1) |
| self.att_weightM = self._init_new_params(out_dim, 1) |
| self.proj_with_att = nn.Linear(in_dim, out_dim) |
| self.proj_without_att = nn.Linear(in_dim, out_dim) |
| self.proj_with_attM = nn.Linear(in_dim, out_dim) |
| self.proj_without_attM = nn.Linear(in_dim, out_dim) |
| self.bn = nn.BatchNorm1d(out_dim) |
| self.input_drop = nn.Dropout(p=0.2) |
| self.act = nn.SELU(inplace=True) |
| self.temp = kwargs.get("temperature", 1.) |
|
|
| def forward(self, x1, x2, master=None): |
| num_type1 = x1.size(1) |
| num_type2 = x2.size(1) |
| x1 = self.proj_type1(x1) |
| x2 = self.proj_type2(x2) |
| x = torch.cat([x1, x2], dim=1) |
| if master is None: |
| master = torch.mean(x, dim=1, keepdim=True) |
| x = self.input_drop(x) |
| att_map = self._derive_att_map(x, num_type1, num_type2) |
| master = self._update_master(x, master) |
| x = self._project(x, att_map) |
| x = self._apply_BN(x) |
| x = self.act(x) |
| x1 = x.narrow(1, 0, num_type1) |
| x2 = x.narrow(1, num_type1, num_type2) |
| return x1, x2, master |
|
|
| def _update_master(self, x, master): |
| att_map = self._derive_att_map_master(x, master) |
| return self._project_master(x, master, att_map) |
|
|
| def _pairwise_mul_nodes(self, x): |
| nb_nodes = x.size(1) |
| x = x.unsqueeze(2).expand(-1, -1, nb_nodes, -1) |
| return x * x.transpose(1, 2) |
|
|
| def _derive_att_map_master(self, x, master): |
| att_map = x * master |
| att_map = torch.tanh(self.att_projM(att_map)) |
| att_map = torch.matmul(att_map, self.att_weightM) |
| att_map = att_map / self.temp |
| return F.softmax(att_map, dim=-2) |
|
|
| def _derive_att_map(self, x, num_type1, num_type2): |
| att_map = self._pairwise_mul_nodes(x) |
| att_map = torch.tanh(self.att_proj(att_map)) |
| att_board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1) |
| att_board[:, :num_type1, :num_type1, :] = torch.matmul( |
| att_map[:, :num_type1, :num_type1, :], self.att_weight11) |
| att_board[:, num_type1:, num_type1:, :] = torch.matmul( |
| att_map[:, num_type1:, num_type1:, :], self.att_weight22) |
| att_board[:, :num_type1, num_type1:, :] = torch.matmul( |
| att_map[:, :num_type1, num_type1:, :], self.att_weight12) |
| att_board[:, num_type1:, :num_type1, :] = torch.matmul( |
| att_map[:, num_type1:, :num_type1, :], self.att_weight12) |
| att_map = att_board / self.temp |
| return F.softmax(att_map, dim=-2) |
|
|
| def _project(self, x, att_map): |
| return self.proj_with_att(torch.matmul(att_map.squeeze(-1), x)) + self.proj_without_att(x) |
|
|
| def _project_master(self, x, master, att_map): |
| x1 = self.proj_with_attM(torch.matmul(att_map.squeeze(-1).unsqueeze(1), x)) |
| return x1 + self.proj_without_attM(master) |
|
|
| def _apply_BN(self, x): |
| org_size = x.size() |
| return self.bn(x.view(-1, org_size[-1])).view(org_size) |
|
|
| def _init_new_params(self, *size): |
| out = nn.Parameter(torch.FloatTensor(*size)) |
| nn.init.xavier_normal_(out) |
| return out |
|
|
|
|
| class GraphPool(nn.Module): |
| def __init__(self, k: float, in_dim: int, p: Union[float, int]): |
| super().__init__() |
| self.k = k |
| self.sigmoid = nn.Sigmoid() |
| self.proj = nn.Linear(in_dim, 1) |
| self.drop = nn.Dropout(p=p) if p > 0 else nn.Identity() |
| self.in_dim = in_dim |
|
|
| def forward(self, h): |
| scores = self.sigmoid(self.proj(self.drop(h))) |
| return self.top_k_graph(scores, h, self.k) |
|
|
| def top_k_graph(self, scores, h, k): |
| _, n_nodes, n_feat = h.size() |
| n_nodes = max(int(n_nodes * k), 1) |
| _, idx = torch.topk(scores, n_nodes, dim=1) |
| idx = idx.expand(-1, -1, n_feat) |
| h = h * scores |
| return torch.gather(h, 1, idx) |
|
|
|
|
| class Residual_block(nn.Module): |
| def __init__(self, nb_filts, first=False): |
| super().__init__() |
| self.first = first |
| if not self.first: |
| self.bn1 = nn.BatchNorm2d(num_features=nb_filts[0]) |
| self.conv1 = nn.Conv2d(nb_filts[0], nb_filts[1], kernel_size=(2, 3), padding=(1, 1), stride=1) |
| self.selu = nn.SELU(inplace=True) |
| self.bn2 = nn.BatchNorm2d(num_features=nb_filts[1]) |
| self.conv2 = nn.Conv2d(nb_filts[1], nb_filts[1], kernel_size=(2, 3), padding=(0, 1), stride=1) |
| if nb_filts[0] != nb_filts[1]: |
| self.downsample = True |
| self.conv_downsample = nn.Conv2d(nb_filts[0], nb_filts[1], padding=(0, 1), kernel_size=(1, 3), stride=1) |
| else: |
| self.downsample = False |
|
|
| def forward(self, x): |
| identity = x |
| |
| |
| out = self.conv1(x) |
| out = self.bn2(out) |
| out = self.selu(out) |
| out = self.conv2(out) |
| if self.downsample: |
| identity = self.conv_downsample(identity) |
| out += identity |
| return out |
|
|
|
|
| class W2VAASIST(nn.Module): |
| def __init__(self): |
| super().__init__() |
| filts = [128, [1, 32], [32, 32], [32, 64], [64, 64]] |
| gat_dims = [64, 32] |
| pool_ratios = [0.5, 0.5, 0.5, 0.5] |
| temperatures = [2.0, 2.0, 100.0, 100.0] |
| self.first_bn = nn.BatchNorm2d(num_features=1) |
| self.first_bn1 = nn.BatchNorm2d(num_features=64) |
| self.drop = nn.Dropout(0.5, inplace=True) |
| self.drop_way = nn.Dropout(0.2, inplace=True) |
| self.selu = nn.SELU(inplace=True) |
| self.encoder = nn.Sequential( |
| nn.Sequential(Residual_block(nb_filts=filts[1], first=True)), |
| nn.Sequential(Residual_block(nb_filts=filts[2])), |
| nn.Sequential(Residual_block(nb_filts=filts[3])), |
| nn.Sequential(Residual_block(nb_filts=filts[4])), |
| nn.Sequential(Residual_block(nb_filts=filts[4])), |
| nn.Sequential(Residual_block(nb_filts=filts[4]))) |
| self.LL = nn.Linear(1024, 128) |
| self.attention = nn.Sequential( |
| nn.Conv2d(64, 128, kernel_size=(1, 1)), |
| nn.SELU(inplace=True), |
| nn.BatchNorm2d(128), |
| nn.Conv2d(128, 64, kernel_size=(1, 1)), |
| ) |
| self.pos_S = nn.Parameter(torch.randn(1, 42, filts[-1][-1])) |
| self.master1 = nn.Parameter(torch.randn(1, 1, gat_dims[0])) |
| self.master2 = nn.Parameter(torch.randn(1, 1, gat_dims[0])) |
| self.GAT_layer_S = GraphAttentionLayer(filts[-1][-1], gat_dims[0], temperature=temperatures[0]) |
| self.GAT_layer_T = GraphAttentionLayer(filts[-1][-1], gat_dims[0], temperature=temperatures[1]) |
| self.HtrgGAT_layer_ST11 = HtrgGraphAttentionLayer(gat_dims[0], gat_dims[1], temperature=temperatures[2]) |
| self.HtrgGAT_layer_ST12 = HtrgGraphAttentionLayer(gat_dims[1], gat_dims[1], temperature=temperatures[2]) |
| self.HtrgGAT_layer_ST21 = HtrgGraphAttentionLayer(gat_dims[0], gat_dims[1], temperature=temperatures[2]) |
| self.HtrgGAT_layer_ST22 = HtrgGraphAttentionLayer(gat_dims[1], gat_dims[1], temperature=temperatures[2]) |
| self.pool_S = GraphPool(pool_ratios[0], gat_dims[0], 0.3) |
| self.pool_T = GraphPool(pool_ratios[1], gat_dims[0], 0.3) |
| self.pool_hS1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3) |
| self.pool_hT1 = GraphPool(pool_ratios[2], gat_dims[1], 0.3) |
| self.pool_hS2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3) |
| self.pool_hT2 = GraphPool(pool_ratios[2], gat_dims[1], 0.3) |
| self.out_layer = nn.Linear(5 * gat_dims[1], 2) |
|
|
| def forward(self, x): |
| x = x.squeeze(dim=1) |
| x = x.transpose(1, 2) |
| x = self.LL(x) |
| x = x.transpose(1, 2) |
| x = x.unsqueeze(dim=1) |
| x = F.max_pool2d(x, (3, 3)) |
| x = self.first_bn(x) |
| x = self.selu(x) |
| x = self.encoder(x) |
| x = self.first_bn1(x) |
| x = self.selu(x) |
| w = self.attention(x) |
| w1 = F.softmax(w, dim=-1) |
| m = torch.sum(x * w1, dim=-1) |
| e_S = m.transpose(1, 2) + self.pos_S |
| gat_S = self.GAT_layer_S(e_S) |
| out_S = self.pool_S(gat_S) |
| w2 = F.softmax(w, dim=-2) |
| m1 = torch.sum(x * w2, dim=-2) |
| e_T = m1.transpose(1, 2) |
| gat_T = self.GAT_layer_T(e_T) |
| out_T = self.pool_T(gat_T) |
| master1 = self.master1.expand(x.size(0), -1, -1) |
| master2 = self.master2.expand(x.size(0), -1, -1) |
| out_T1, out_S1, master1 = self.HtrgGAT_layer_ST11(out_T, out_S, master=self.master1) |
| out_S1 = self.pool_hS1(out_S1) |
| out_T1 = self.pool_hT1(out_T1) |
| out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST12(out_T1, out_S1, master=master1) |
| out_T1 = out_T1 + out_T_aug |
| out_S1 = out_S1 + out_S_aug |
| master1 = master1 + master_aug |
| out_T2, out_S2, master2 = self.HtrgGAT_layer_ST21(out_T, out_S, master=self.master2) |
| out_S2 = self.pool_hS2(out_S2) |
| out_T2 = self.pool_hT2(out_T2) |
| out_T_aug, out_S_aug, master_aug = self.HtrgGAT_layer_ST22(out_T2, out_S2, master=master2) |
| out_T2 = out_T2 + out_T_aug |
| out_S2 = out_S2 + out_S_aug |
| master2 = master2 + master_aug |
| out_T1 = self.drop_way(out_T1) |
| out_T2 = self.drop_way(out_T2) |
| out_S1 = self.drop_way(out_S1) |
| out_S2 = self.drop_way(out_S2) |
| master1 = self.drop_way(master1) |
| master2 = self.drop_way(master2) |
| out_T = torch.max(out_T1, out_T2) |
| out_S = torch.max(out_S1, out_S2) |
| master = torch.max(master1, master2) |
| T_max, _ = torch.max(torch.abs(out_T), dim=1) |
| T_avg = torch.mean(out_T, dim=1) |
| S_max, _ = torch.max(torch.abs(out_S), dim=1) |
| S_avg = torch.mean(out_S, dim=1) |
| last_hidden = torch.cat([T_max, T_avg, S_max, S_avg, master.squeeze(1)], dim=1) |
| last_hidden = self.drop(last_hidden) |
| output = self.out_layer(last_hidden) |
| return last_hidden, output |
|
|
|
|
| class CodecfakeNet(nn.Module): |
| """wav (B, 64600) -> per-row normalisation -> XLS-R hidden_states[5] -> W2VAASIST logits (B, 2).""" |
|
|
| def __init__(self, xlsr: XLSRLayer5, head: W2VAASIST): |
| super().__init__() |
| self.xlsr = xlsr |
| self.head = head |
|
|
| def forward(self, wav: torch.Tensor) -> torch.Tensor: |
| h = self.xlsr(normalize(wav)) |
| h = h.unsqueeze(1).transpose(2, 3) |
| return self.head(h)[1] |
|
|
|
|
| |
| class _Stub(nn.Module): |
| """Stand-in for the source's model.* classes while unpickling (state only, no source code).""" |
|
|
| def forward(self, *a, **k): |
| raise RuntimeError("stub") |
|
|
|
|
| _ALLOWED = { |
| ("collections", "OrderedDict"), ("builtins", "set"), ("__builtin__", "set"), |
| ("torch", "FloatStorage"), ("torch", "LongStorage"), |
| ("torch._utils", "_rebuild_parameter"), ("torch._utils", "_rebuild_tensor_v2"), |
| } |
| _SOURCE_CLASSES = {"W2VAASIST", "GraphAttentionLayer", "HtrgGraphAttentionLayer", "GraphPool", "Residual_block"} |
|
|
|
|
| class _AllowListUnpickler(pickle.Unpickler): |
| def find_class(self, module, name): |
| if module == "model" and name in _SOURCE_CLASSES: |
| return _Stub |
| if module.startswith("torch.nn.modules.") or (module, name) in _ALLOWED: |
| return super().find_class(module, name) |
| raise pickle.UnpicklingError(f"blocked global {module}.{name}") |
|
|
|
|
| class _PickleModule: |
| Unpickler = _AllowListUnpickler |
| load = staticmethod(pickle.load) |
| __name__ = "codecfake_allowlist_pickle" |
|
|
|
|
| def load_pickled_state(path: str) -> dict: |
| """state_dict of the source's torch.save(W2VAASIST) file, via the allow-list unpickler.""" |
| mod = torch.load(path, map_location="cpu", weights_only=False, pickle_module=_PickleModule) |
| assert isinstance(mod, _Stub), type(mod) |
| return {k: v.detach().clone().contiguous() for k, v in mod.state_dict().items()} |
|
|
|
|
| |
| SRC_REPO = "xieyuankun/Codecfake" |
| SRC_REV = "72c32e2840eef95fbeade349425fa9c5391fa11e" |
| CKPT_SHA256 = { |
| "cotrain_w2v2aasist_CSAM": "86876e2c0eb98d6e60ecf2f1139112206adf47b1562aee72530d0c22e04c644a", |
| "codec_w2v2aasist": "5a606d1fbd578dbd214471cead283c236f4204c54d602e945924b3594f8c25a3", |
| } |
|
|
|
|
| def ckpt_url(sub: str) -> str: |
| return f"https://raw.githubusercontent.com/{SRC_REPO}/{SRC_REV}/pretrained_model/{sub}/anti-spoofing_feat_model.pt" |
|
|
|
|
| def head_state(sub: str, model_dir: str) -> dict: |
| """W2VAASIST state_dict for pretrained_model/<sub>. No upstream license -> not re-hosted: on first use the |
| authors' pickled module is downloaded from the commit-pinned GitHub raw URL (sha256-checked) into |
| ../_ckpt_dl/xieyuankun__Codecfake/<rev>/, converted once with the allow-list unpickler, and cached as |
| <model_dir>/head_<sub>.safetensors (loaded on every later run).""" |
| import hashlib |
| import os |
| import urllib.request |
| from safetensors.torch import load_file, save_file |
| cache = os.path.join(model_dir, f"head_{sub}.safetensors") |
| if not os.path.exists(cache): |
| pt = os.path.join(model_dir, "..", "_ckpt_dl", "xieyuankun__Codecfake", SRC_REV, "pretrained_model", |
| sub, "anti-spoofing_feat_model.pt") |
| if not os.path.exists(pt): |
| os.makedirs(os.path.dirname(pt), exist_ok=True) |
| urllib.request.urlretrieve(ckpt_url(sub), pt + ".part") |
| os.replace(pt + ".part", pt) |
| digest = hashlib.sha256(open(pt, "rb").read()).hexdigest() |
| if digest != CKPT_SHA256[sub]: |
| raise RuntimeError(f"{pt}: sha256 {digest} != pinned {CKPT_SHA256[sub]}") |
| save_file(load_pickled_state(pt), cache + ".part", metadata={"source": ckpt_url(sub), "sha256": digest}) |
| os.replace(cache + ".part", cache) |
| return load_file(cache) |
|
|
