Update README.md

README.md

@@ -88,10 +88,15 @@ import torchaudio
 import fairseq
 from huggingface_hub import PyTorchModelHubMixin
 
-# === Detect all audio files in a specified folder ===
+# This is the only part of the script you need to modify.
+# Set this to the path where your audio files are stored.
 folder_path = "/path/to/folder/contains/wavs/"
 audio_formats = (".mp3", ".wav", ".flac", ".m4a")
 
+# === Set device (use GPU if available) ===
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
 # === Download Fairseq checkpoint if not present ===
 # The downloaded checkpoint is used for building front-end architecture,
 # its weights will be replaced by the model.safetensors file in this repo.
@@ -111,16 +116,15 @@ class SSLModel(torch.nn.Module):
         super().__init__()
         # The downloaded .pt file is used here
         model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([ssl_path])
-        self.model = model[0]
+        self.model = model[0].to(device)
 
     def extract_feat(self, input_data):
         # If input has shape (B, T, 1), squeeze the last dim
         if input_data.ndim == 3:
             input_data = input_data[:, :, 0]
-
         # Extract features
         with torch.no_grad():
-            features = self.model(input_data, mask=False, features_only=True)['x']
+            features = self.model(input_data.to(device), mask=False, features_only=True)['x']
         return features
 
 # === Function for reading and pre-processing waveforms ===
@@ -135,7 +139,7 @@ def load_wav_and_preprocess(wav_path, target_sr=16000):
     with torch.no_grad():
         wav = torch.nn.functional.layer_norm(wav, wav.shape)
     # Add batch dimension and return
-    return wav.unsqueeze(0)
+    return wav.unsqueeze(0).to(device)
 
 # === The actual deepfake detection model using SSL frontend + FC backend ===
 class DeepfakeDetector(torch.nn.Module, PyTorchModelHubMixin):
@@ -164,6 +168,7 @@ class DeepfakeDetector(torch.nn.Module, PyTorchModelHubMixin):
 
 # === Load AntiDeepfake model from Hugging Face===
 model = DeepfakeDetector.from_pretrained("nii-yamagishilab/wav2vec-large-anti-deepfake")
+model.to(device)
 model.eval()
 
 # === Inference on a folder of audio files ===