"""WavLM-large + AASIST speech deepfake detector (inference only).""" import torch import torch.nn as nn import torch.nn.functional as F from transformers import WavLMModel class _GraphAttentionLayer(nn.Module): def __init__(self, in_dim, out_dim, temperature): super().__init__() self.att_proj = nn.Linear(in_dim, out_dim) self.att_weight = nn.Parameter(torch.empty(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.temp = temperature def forward(self, x): N = x.size(1) x_i = x.unsqueeze(2).expand(-1, -1, N, -1) att = torch.matmul(torch.tanh(self.att_proj(x_i * x_i.transpose(1, 2))), self.att_weight) / self.temp att = F.softmax(att, dim=-2) out = self.proj_with_att(torch.matmul(att.squeeze(-1), x)) + self.proj_without_att(x) s = out.shape return F.selu(self.bn(out.reshape(-1, s[-1])).reshape(s)) class _HtrgGraphAttentionLayer(nn.Module): def __init__(self, in_dim, out_dim, temperature): 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 = nn.Parameter(torch.empty(out_dim, 1)) self.att_weight22 = nn.Parameter(torch.empty(out_dim, 1)) self.att_weight12 = nn.Parameter(torch.empty(out_dim, 1)) self.att_weightM = nn.Parameter(torch.empty(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.temp = temperature def forward(self, x1, x2, master): n1 = x1.size(1) x = torch.cat([self.proj_type1(x1), self.proj_type2(x2)], dim=1) N = x.size(1) x_i = x.unsqueeze(2).expand(-1, -1, N, -1) att_map = torch.tanh(self.att_proj(x_i * x_i.transpose(1, 2))) board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1) board[:, :n1, :n1] = torch.matmul(att_map[:, :n1, :n1], self.att_weight11) board[:, n1:, n1:] = torch.matmul(att_map[:, n1:, n1:], self.att_weight22) board[:, :n1, n1:] = torch.matmul(att_map[:, :n1, n1:], self.att_weight12) board[:, n1:, :n1] = torch.matmul(att_map[:, n1:, :n1], self.att_weight12) board = F.softmax(board / self.temp, dim=-2) out = self.proj_with_att(torch.matmul(board.squeeze(-1), x)) + self.proj_without_att(x) s = out.shape out = F.selu(self.bn(out.reshape(-1, s[-1])).reshape(s)) att_m = F.softmax(torch.matmul(torch.tanh(self.att_projM(x * master)), self.att_weightM) / self.temp, dim=-2) master = self.proj_with_attM(torch.matmul(att_m.squeeze(-1).unsqueeze(1), x)) + self.proj_without_attM(master) return out.narrow(1, 0, n1), out.narrow(1, n1, N - n1), master class _GraphPool(nn.Module): def __init__(self, k, in_dim): super().__init__() self.k = k self.proj = nn.Linear(in_dim, 1) def forward(self, h): scores = torch.sigmoid(self.proj(h)) n_keep = max(int(h.size(1) * self.k), 1) idx = torch.topk(scores, n_keep, dim=1)[1].expand(-1, -1, h.size(-1)) return torch.gather(h * scores, 1, idx) class AASISTPool(nn.Module): def __init__(self, feat_dim, num_layers, gat_dims=(128, 64), pool_ratios=(0.5, 0.7, 0.5, 0.5), temps=(2.0, 2.0, 100.0, 100.0)): super().__init__() g = list(gat_dims) self.proj = nn.Linear(feat_dim, g[0]) self.pos_S = nn.Parameter(torch.randn(1, num_layers, g[0])) self.master1 = nn.Parameter(torch.randn(1, 1, g[0])) self.master2 = nn.Parameter(torch.randn(1, 1, g[0])) self.GAT_S = _GraphAttentionLayer(g[0], g[0], temps[0]) self.GAT_T = _GraphAttentionLayer(g[0], g[0], temps[1]) self.HtrgGAT_11 = _HtrgGraphAttentionLayer(g[0], g[1], temps[2]) self.HtrgGAT_12 = _HtrgGraphAttentionLayer(g[1], g[1], temps[3]) self.HtrgGAT_21 = _HtrgGraphAttentionLayer(g[0], g[1], temps[2]) self.HtrgGAT_22 = _HtrgGraphAttentionLayer(g[1], g[1], temps[3]) self.pool_S = _GraphPool(pool_ratios[0], g[0]) self.pool_T = _GraphPool(pool_ratios[1], g[0]) self.pool_hS1 = _GraphPool(pool_ratios[2], g[1]) self.pool_hT1 = _GraphPool(pool_ratios[2], g[1]) self.pool_hS2 = _GraphPool(pool_ratios[2], g[1]) self.pool_hT2 = _GraphPool(pool_ratios[2], g[1]) self.output_dim = 5 * g[1] def forward(self, hidden_states): projected = self.proj(torch.stack(hidden_states, dim=1)) e_S = torch.max(torch.abs(projected), dim=2)[0] + self.pos_S e_T = torch.max(torch.abs(projected), dim=1)[0] out_S = self.pool_S(self.GAT_S(e_S)) out_T = self.pool_T(self.GAT_T(e_T)) out_T1, out_S1, m1 = self.HtrgGAT_11(out_T, out_S, self.master1) out_S1, out_T1 = self.pool_hS1(out_S1), self.pool_hT1(out_T1) a_T, a_S, a_m = self.HtrgGAT_12(out_T1, out_S1, m1) out_T1, out_S1, m1 = out_T1 + a_T, out_S1 + a_S, m1 + a_m out_T2, out_S2, m2 = self.HtrgGAT_21(out_T, out_S, self.master2) out_S2, out_T2 = self.pool_hS2(out_S2), self.pool_hT2(out_T2) a_T, a_S, a_m = self.HtrgGAT_22(out_T2, out_S2, m2) out_T2, out_S2, m2 = out_T2 + a_T, out_S2 + a_S, m2 + a_m out_T, out_S, master = torch.max(out_T1, out_T2), torch.max(out_S1, out_S2), torch.max(m1, m2) return torch.cat([torch.max(torch.abs(out_T), 1)[0], out_T.mean(1), torch.max(torch.abs(out_S), 1)[0], out_S.mean(1), master.squeeze(1)], dim=1) class MLPBlock(nn.Module): def __init__(self, dim): super().__init__() self.net = nn.Sequential(nn.Linear(dim, dim * 2), nn.GELU(), nn.Dropout(0.1), nn.Linear(dim * 2, dim), nn.GELU(), nn.Dropout(0.1)) def forward(self, x): return self.net(x) class Classifier(nn.Module): def __init__(self, input_dim, hidden_dim, num_layers, num_labels=1): super().__init__() self.input_proj = nn.Linear(input_dim, hidden_dim) if input_dim != hidden_dim else nn.Identity() self.mlp_blocks = nn.ModuleList([MLPBlock(hidden_dim) for _ in range(num_layers)]) self.layer_norms = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(num_layers)]) self.output_head = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_dim, num_labels)) def forward(self, x): x = self.input_proj(x) for mlp, ln in zip(self.mlp_blocks, self.layer_norms): x = x + mlp(ln(x)) return self.output_head(x) class DeepfakeDetector(nn.Module): def __init__(self, hidden_dim=320, classifier_layers=3, aasist_gat_dims=(128, 64)): super().__init__() self.wavlm = WavLMModel.from_pretrained("microsoft/wavlm-large") self.pool = AASISTPool(self.wavlm.config.hidden_size, self.wavlm.config.num_hidden_layers, aasist_gat_dims) self.classifier = Classifier(self.pool.output_dim, hidden_dim, classifier_layers) @torch.no_grad() def forward(self, waveform): # (B, T) 16 kHz -> logit (B,); sigmoid(logit)=P(real) hs = list(self.wavlm(waveform, output_hidden_states=True).hidden_states[1:]) return self.classifier(self.pool(hs)).squeeze(-1)