"""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 # hidden_states[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) # ----------------------------------------------------------------------------- XLS-R front-end 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 # allowed: dropped layers 5..23 + the pre-training-only heads 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) # ----------------------------------------------------------------------------- AASIST head (model.py) 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 # source: bn1/selu computed for non-first blocks, but conv1 is applied to x (kept verbatim; the # inplace SELU acts on bn1's fresh output, so x is untouched) 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)) # (B, T, 1024) h = h.unsqueeze(1).transpose(2, 3) # (B, 1, 1024, T), as generate_score.py return self.head(h)[1] # ----------------------------------------------------------------------------- checkpoint conversion class _Stub(nn.Module): """Stand-in for the source's model.* classes while unpickling (state only, no source code).""" def forward(self, *a, **k): # pragma: no cover 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()} # ----------------------------------------------------------------------------- weight acquisition 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/. 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//, converted once with the allow-list unpickler, and cached as /head_.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)