1import os
2import torch
3import torchaudio
4from fairseq.models.wav2vec import Wav2Vec2Model, Wav2Vec2Config
5from huggingface_hub import PyTorchModelHubMixin
6
7# This is the only part of the script you need to modify.
8# Set this to the path where your audio files are stored.
9folder_path = "/path/to/folder/contains/wavs/"
10audio_formats = (".mp3", ".wav", ".flac", ".m4a")
11
12# === Set device (use GPU if available) ===
13device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14print(f"Using device: {device}")
15
16# === Wrapper for the SSL model ===
17class SSLModel(torch.nn.Module):
18 def __init__(self):
19 super().__init__()
20 # Model config used to build SSL architecture
21 cfg = Wav2Vec2Config(
22 quantize_targets=True,
23 extractor_mode="layer_norm",
24 layer_norm_first=True,
25 final_dim=768,
26 latent_temp=(2.0, 0.1, 0.999995),
27 encoder_layerdrop=0.0,
28 dropout_input=0.0,
29 dropout_features=0.0,
30 dropout=0.0,
31 attention_dropout=0.0,
32 conv_bias=True,
33 encoder_layers=24,
34 encoder_embed_dim=1024,
35 encoder_ffn_embed_dim=4096,
36 encoder_attention_heads=16,
37 feature_grad_mult=1.0
38 )
39 # Initialize SSL model with random weights
40 self.model = Wav2Vec2Model(cfg)
41
42 def extract_feat(self, input_data):
43 # If input has shape (B, T, 1), squeeze the last dim
44 if input_data.ndim == 3:
45 input_data = input_data[:, :, 0]
46 # Extract features
47 with torch.no_grad():
48 features = self.model(input_data.to(device), mask=False, features_only=True)['x']
49 return features
50
51# === Function for reading and pre-processing waveforms ===
52def load_wav_and_preprocess(wav_path, target_sr=16000):
53 # Load audio file
54 wav, sr = torchaudio.load(wav_path)
55 # Convert to mono if stereo
56 wav = wav.mean(dim=0)
57 # Resample to target sampling rate
58 wav = torchaudio.functional.resample(wav, sr, new_freq=target_sr)
59 # Normalize waveform
60 with torch.no_grad():
61 wav = torch.nn.functional.layer_norm(wav, wav.shape)
62 # Add batch dimension and return
63 return wav.unsqueeze(0).to(device)
64
65# === The actual deepfake detection model using SSL frontend + FC backend ===
66class DeepfakeDetector(torch.nn.Module, PyTorchModelHubMixin):
67 def __init__(self):
68 super().__init__()
69 self.ssl_orig_output_dim = 1024
70 self.num_classes = 2
71
72 # Frontend: SSL model
73 self.m_ssl = SSLModel()
74
75 # Backend: Pooling + Classification
76 self.adap_pool1d = torch.nn.AdaptiveAvgPool1d(output_size=1)
77 self.proj_fc = torch.nn.Linear(
78 in_features=self.ssl_orig_output_dim,
79 out_features=self.num_classes,
80 )
81
82 def forward(self, wav):
83 emb = self.m_ssl.extract_feat(wav) # [B, T, D]
84 emb = emb.transpose(1, 2) # [B, D, T]
85 pooled_emb = self.adap_pool1d(emb) # [B, D, 1]
86 pooled_emb = pooled_emb.squeeze(-1) # [B, D]
87 logits = self.proj_fc(pooled_emb) # [B, 2]
88 return logits
89
90# === Load AntiDeepfake model from Hugging Face===
91model = DeepfakeDetector.from_pretrained("nii-yamagishilab/wav2vec-large-anti-deepfake-nda")
92model.to(device)
93model.eval()
94
95# === Inference on a folder of audio files ===
96results = []
97for root, _, files in os.walk(folder_path):
98 for file in files:
99 if file.lower().endswith(audio_formats):
100 input_path = os.path.join(root, file)
101 with torch.no_grad():
102 wav = load_wav_and_preprocess(input_path)
103 logits = model(wav)
104 probs = torch.nn.functional.softmax(logits, dim=1)
105 results.append((file, probs.cpu().numpy()[0]))
106
107# Sort results alphabetically by filename
108results.sort(key=lambda x: x[0])
109
110# Print formatted results
111print("\n=== Deepfake Detection Results ===")
112for file_name, prob in results:
113 print(f"{file_name}: real prob = {prob[1]:.3f}, fake prob = {prob[0]:.3f}")
