Update app.py

app.py

@@ -1,10 +1,20 @@
+import os
 import gradio as gr
 import torch
-from torchaudio.sox_effects import apply_effects_file
+import pydub
+import torchaudio
+from torchaudio.sox_effects import apply_effects_tensor
+import numpy as np
 from transformers import AutoFeatureExtractor, AutoModelForAudioXVector
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+def load_audio(file_name):
+    audio = pydub.AudioSegment.from_file(file_name)
+    arr = np.array(audio.get_array_of_samples(), dtype=np.float32)
+    arr = arr / (1 << (8 * audio.sample_width - 1))
+    return arr.astype(np.float32), audio.frame_rate
+
 STYLE = """
 <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" integrity="sha256-YvdLHPgkqJ8DVUxjjnGVlMMJtNimJ6dYkowFFvp4kKs=" crossorigin="anonymous">
 """
@@ -44,7 +54,7 @@ EFFECTS = [
 
 THRESHOLD = 0.85
 
-model_name = "microsoft/unispeech-sat-base-plus-sv"
+model_name = "microsoft/wavlm-base-plus-sv"
 feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 model = AutoModelForAudioXVector.from_pretrained(model_name).to(device)
 cosine_sim = torch.nn.CosineSimilarity(dim=-1)
@@ -53,9 +63,12 @@ cosine_sim = torch.nn.CosineSimilarity(dim=-1)
 def similarity_fn(path1, path2):
     if not (path1 and path2):
         return '<b style="color:red">ERROR: Please record audio for *both* speakers!</b>'
-
-    wav1, _ = apply_effects_file(path1, EFFECTS)
-    wav2, _ = apply_effects_file(path2, EFFECTS)
+    
+    wav1, sr1 = load_audio(path1)
+    print(wav1, wav1.shape, wav1.dtype)
+    wav1, _ = apply_effects_tensor(torch.tensor(wav1).unsqueeze(0), sr1, EFFECTS)
+    wav2, sr2 = load_audio(path2)
+    wav2, _ = apply_effects_tensor(torch.tensor(wav2).unsqueeze(0), sr2, EFFECTS)
     print(wav1.shape, wav2.shape)
 
     input1 = feature_extractor(wav1.squeeze(0), return_tensors="pt", sampling_rate=16000).input_values.to(device)
@@ -89,8 +102,8 @@ description = (
 )
 article = (
     "<p style='text-align: center'>"
-    "<a href='https://huggingface.co/microsoft/unispeech-sat-large-sv' target='_blank'>🎙️ Learn more about UniSpeech-SAT</a> | "
-    "<a href='https://arxiv.org/abs/2110.05752' target='_blank'>📚 UniSpeech-SAT paper</a> | "
+    "<a href='https://huggingface.co/microsoft/wavlm-base-plus-sv' target='_blank'>🎙️ Learn more about WavLM</a> | "
+    "<a href='https://arxiv.org/abs/2110.13900' target='_blank'>📚 WavLM paper</a> | "
     "<a href='https://www.danielpovey.com/files/2018_icassp_xvectors.pdf' target='_blank'>📚 X-Vector paper</a>"
     "</p>"
 )
@@ -103,7 +116,7 @@ interface = gr.Interface(
     fn=similarity_fn,
     inputs=inputs,
     outputs=output,
-    title="Voice Authentication with UniSpeech-SAT + X-Vectors",
+    title="Voice Authentication with WavLM + X-Vectors",
     description=description,
     article=article,
     layout="horizontal",
@@ -112,4 +125,4 @@ interface = gr.Interface(
     live=False,
     examples=examples,
 )
-interface.launch(enable_queue=True)
+interface.launch(enable_queue=True)