new commit

app.py

@@ -1,44 +1,29 @@
-from score import load_model
-from predict import loadWav
+import os
 import torch
 import torch.nn.functional as F
+from ssl_ecapa_model import SSL_ECAPA_TDNN
+from score import loadModel
+from predict import loadWav
 import gradio as gr
-import time
 
-model = load_model("wavlm_ecapa.model")
+model = loadModel('voxsim_wavlm_ecapa.model')
 model.eval()
 
 def calc_voxsim(inp_path, ref_path):
-    start = time.time()
-    inp_wavs, inp_wav = loadWav(inp_path)
-    ref_wavs, ref_wav = loadWav(ref_path)
-    print("loadWav time: ", time.time() - start)
-
-    inp_wavs = torch.FloatTensor(inp_wavs)
-    inp_wav = torch.FloatTensor(inp_wav)
-    ref_wavs = torch.FloatTensor(ref_wavs)
-    ref_wav = torch.FloatTensor(ref_wav)
-    print("torch.FloatTensor time: ", time.time() - start)
+    inp_wav = loadWav(inp_path, max_frames=0)
+    ref_wav = loadWav(ref_path, max_frames=0)
 
     with torch.no_grad():
-        input_emb_1 = F.normalize(model.forward(inp_wavs), p=2, dim=1)
-        print("input_emb_1 time: ", time.time() - start)
-        input_emb_2 = F.normalize(model.forward(inp_wav), p=2, dim=1)
-        print("input_emb_2 time: ", time.time() - start)
-        ref_emb_1 = F.normalize(model.forward(ref_wavs), p=2, dim=1)
-        print("ref_emb_1 time: ", time.time() - start)
-        ref_emb_2 = F.normalize(model.forward(ref_wav), p=2, dim=1)
-        print("ref_emb_2 time: ", time.time() - start)
+        input_emb = F.normalize(model.forward(inp_wav), p=2, dim=1)
+        ref_emb = F.normalize(model.forward(ref_wav), p=2, dim=1)
 
-        score_1 = torch.mean(torch.matmul(input_emb_1, ref_emb_1.T))
-        score_2 = torch.mean(torch.matmul(input_emb_2, ref_emb_2.T))
-        score = (score_1 + score_2) / 2
-        print("score time: ", time.time() - start)
+        score = torch.matmul(input_emb, ref_emb.T)
         return score.detach().cpu().numpy()
 
 description = """
 Voice similarity demo using wavlm-ecapa model, which is trained on Voxsim dataset.
 This demo only accepts .wav format. Best at 16 kHz sampling rate.
+The inference process of this Spaces demo is suboptimal due to the limitations of a basic CPU. To obtain an accurate score, refer to the "[voxsim_trainer](https://github.com/kaistmm/voxsim_trainer)" repository and run the code via the CLI.
 
 Paper is available [here](https://arxiv.org/abs/2407.18505)
 """
@@ -46,8 +31,8 @@ Paper is available [here](https://arxiv.org/abs/2407.18505)
 iface = gr.Interface(
     fn=calc_voxsim,
     inputs=(
-        gr.Audio(label="Input Audio"),
-        gr.Audio(label="Reference Audio")
+        gr.Audio(label="Input Audio", type='filepath'),
+        gr.Audio(label="Reference Audio", type='filepath')
     ),
     outputs="text",
     title="voice similarity with VoxSim",

predict.py

@@ -2,10 +2,9 @@ import argparse
 import pathlib
 import tqdm
 from torch.utils.data import Dataset, DataLoader
-import librosa
-import numpy
-from score import Score
+from score import loadWav, Score
 import torch
+import os
 
 import warnings
 warnings.filterwarnings("ignore")
@@ -13,93 +12,61 @@ warnings.filterwarnings("ignore")
 
 def get_arg():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--bs", required=False, default=None, type=int)
-    parser.add_argument("--mode", required=True, choices=["predict_file", "predict_dir"], type=str)
-    parser.add_argument("--ckpt_path", required=False, default="wavlm_ecapa.model", type=pathlib.Path)
-    parser.add_argument("--inp_dir", required=False, default=None, type=pathlib.Path)
-    parser.add_argument("--ref_dir", required=False, default=None, type=pathlib.Path)
-    parser.add_argument("--inp_path", required=False, default=None, type=pathlib.Path)
-    parser.add_argument("--ref_path", required=False, default=None, type=pathlib.Path)
-    parser.add_argument("--out_path", required=True, type=pathlib.Path)
-    parser.add_argument("--num_workers", required=False, default=0, type=int)
+    parser.add_argument("--mode", required=True, choices=["predict_file", "predict_dir"], type=str, help="predict mode")
+    parser.add_argument("--ckpt_path", required=False, default="voxsim_wavlm_ecapa.model", type=pathlib.Path, help="path to the model checkpoint")
+    parser.add_argument("--inp_dir", required=False, default=None, type=pathlib.Path, help="input directory when predict_dir mode")
+    parser.add_argument("--ref_dir", required=False, default=None, type=pathlib.Path, help="reference directory when predict_dir mode")
+    parser.add_argument("--inp_path", required=False, default=None, type=pathlib.Path, help="input file when predict_file mode")
+    parser.add_argument("--ref_path", required=False, default=None, type=pathlib.Path, help="reference file when predict_file mode")
+    parser.add_argument("--out_path", required=True, type=pathlib.Path, help="output path")
+    parser.add_argument("--num_workers", required=False, default=4, type=int, help="number of workers for dataloader")
     return parser.parse_args()
 
 
-def loadWav(filename, max_frames: int = 400):
-
-    # Maximum audio length
-    max_audio = max_frames * 160 + 240
-
-    # Read wav file and convert to torch tensor
-    if type(filename) == tuple:
-        sr, audio = filename
-        audio = librosa.util.normalize(audio)
-        print(numpy.linalg.norm(audio))
-    else:
-        audio, sr = librosa.load(filename, sr=16000)
-    audio_org = audio.copy()
-    
-    audiosize = audio.shape[0]
-
-    if audiosize <= max_audio:
-        shortage = max_audio - audiosize + 1
-        audio       = numpy.pad(audio, (0, shortage), 'wrap')
-        audiosize   = audio.shape[0]
-
-    startframe = numpy.linspace(0,audiosize-max_audio,num=10)
-    
-    feats = []
-    for asf in startframe:
-        feats.append(audio[int(asf):int(asf)+max_audio])
-
-    feat = numpy.stack(feats,axis=0).astype(numpy.float32)
-
-    return torch.FloatTensor(feat), torch.FloatTensor(numpy.stack([audio_org],axis=0).astype(numpy.float32))
-
-
 class AudioDataset(Dataset):
     def __init__(self, inp_dir_path: pathlib.Path, ref_dir_path: pathlib.Path, max_frames: int = 400):
-        self.inp_wavlist = list(inp_dir_path.glob("*.wav"))
-        self.ref_wavlist = list(ref_dir_path.glob("*.wav"))
-        assert len(self.inp_wavlist) == len(self.ref_wavlist)
+        self.inp_dir_path = inp_dir_path
+        self.ref_dir_path = ref_dir_path
+        self.inp_wavlist = [file for file in os.listdir(inp_dir_path) if file.endswith(".wav")]
+        inp_wavset = set(self.inp_wavlist)
+        ref_wavset = set([file for file in os.listdir(ref_dir_path) if file.endswith(".wav")])
+        diff = inp_wavset - ref_wavset
+        if diff:
+            diff = list(diff)
+            diff.sort()
+            raise ValueError(f"Files {diff} are in inp_dir but not in ref_dir.")
         self.inp_wavlist.sort()
-        self.ref_wavlist.sort()
-        _, self.sr = librosa.load(self.inp_wavlist[0], sr=None)
+
         self.max_audio = max_frames * 160 + 240
     
     def __len__(self):
         return len(self.inp_wavlist)
 
     def __getitem__(self, idx):
-        inp_wavs, inp_wav = loadWav(self.inp_wavlist[idx])
-        ref_wavs, ref_wav = loadWav(self.ref_wavlist[idx])
+        inp_wavs, inp_wav = loadWav(os.path.join(self.inp_dir_path, self.inp_wavlist[idx]))
+        ref_wavs, ref_wav = loadWav(os.path.join(self.ref_dir_path, self.inp_wavlist[idx]))
         return inp_wavs, inp_wav, ref_wavs, ref_wav
 
 def main():
     args = get_arg()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     if args.mode == "predict_file":
-        assert args.inp_path is not None
-        assert args.ref_path is not None
-        assert args.inp_dir is None
-        assert args.ref_dir is None
+        assert args.inp_path is not None, "inp_path is required when mode is predict_file."
+        assert args.ref_path is not None, "ref_path is required when mode is predict_file."
         assert args.inp_path.exists()
-        assert args.inp_path.is_file()
         assert args.ref_path.exists()
+        assert args.inp_path.is_file()
         assert args.ref_path.is_file()
         inp_wavs, inp_wav = loadWav(args.inp_path)
         ref_wavs, ref_wav = loadWav(args.ref_path)
         scorer = Score(ckpt_path=args.ckpt_path, device=device)
         score = scorer.score(inp_wavs, inp_wav, ref_wavs, ref_wav)
-        print("Voxsim score: ", score[0])
+        print("VoxSIM score: ", score)
         with open(args.out_path, "w") as fw:
-            fw.write(str(score[0]))
+            fw.write(str(score))
     else:
         assert args.inp_dir is not None, "inp_dir is required when mode is predict_dir."
         assert args.ref_dir is not None, "ref_dir is required when mode is predict_dir."
-        assert args.bs is not None, "bs is required when mode is predict_dir."
-        assert args.inp_path is None, "inp_path should be None"
-        assert args.ref_path is None, "ref_path should be None"
         assert args.inp_dir.exists()
         assert args.ref_dir.exists()
         assert args.inp_dir.is_dir()
@@ -107,17 +74,18 @@ def main():
         dataset = AudioDataset(args.inp_dir, args.ref_dir)
         loader = DataLoader(
             dataset,
-            batch_size=args.bs,
+            batch_size=1,
             shuffle=False,
             num_workers=args.num_workers)
         scorer = Score(ckpt_path=args.ckpt_path, device=device)
-        with open(args.out_path, 'w'):
-            pass
-        for batch in tqdm.tqdm(loader):
-            scores = score.score(batch.to(device))
-            with open(args.out_path, 'a') as fw:
-                for s in scores:
-                    fw.write(str(s) + "\n")
+        avg_score = []
+        with open(args.out_path, 'w') as fw:
+            for batch in tqdm.tqdm(loader):
+                inp_wavs, inp_wav, ref_wavs, ref_wav = batch
+                score = scorer.score(inp_wavs, inp_wav, ref_wavs, ref_wav)
+                avg_score.append(score)
+                fw.write(str(score) + "\n")
+        print("Average VoxSIM score: ", sum(avg_score)/len(avg_score))
         print("save to ", args.out_path)
 
 if __name__ == "__main__":

score.py

@@ -1,35 +1,49 @@
 import os
+import numpy
+import librosa
 import torch
 import torch.nn.functional as F
 from ssl_ecapa_model import SSL_ECAPA_TDNN
 from huggingface_hub import hf_hub_download
 
 
-def load_model(ckpt_path):
-    model = SSL_ECAPA_TDNN(feat_dim=1024, emb_dim=256, feat_type='wavlm_large')
-    load_parameters(model, ckpt_path)
-    return model
+def loadWav(filename, max_frames: int = 400, num_eval: int = 10):
+
+    # Maximum audio length
+    max_audio = max_frames * 160 + 240
+
+    # Read wav file and convert to torch tensor
+    audio, sr = librosa.load(filename, sr=16000)
+    audio_org = audio.copy()
+    
+    audiosize = audio.shape[0]
+
+    if audiosize <= max_audio:
+        shortage = max_audio - audiosize + 1
+        audio       = numpy.pad(audio, (0, shortage), 'wrap')
+        audiosize   = audio.shape[0]
+
+    startframe = numpy.linspace(0,audiosize-max_audio, num=num_eval)
+    
+    feats = []
+    if max_frames == 0:
+        feats.append(audio)
+        feat = numpy.stack(feats,axis=0).astype(numpy.float32)
+        return torch.FloatTensor(feat)
+    else:
+        for asf in startframe:
+            feats.append(audio[int(asf):int(asf)+max_audio])
+        feat = numpy.stack(feats,axis=0).astype(numpy.float32)
+        return torch.FloatTensor(feat), torch.FloatTensor(numpy.stack([audio_org],axis=0).astype(numpy.float32))
 
 
-def load_parameters(model, ckpt_path):
-    model_state = model.state_dict()
+def loadModel(ckpt_path):
+    model = SSL_ECAPA_TDNN(feat_dim=1024, emb_dim=256, feat_type='wavlm_large')
     if not os.path.isfile(ckpt_path):
         print("Downloading model from Hugging Face Hub...")
-        new_ckpt_path = hf_hub_download(repo_id="junseok520/voxsim-models", filename=ckpt_path, local_dir="./")
-        ckpt_path = new_ckpt_path
-    loaded_state = torch.load(ckpt_path, map_location='cpu', weights_only=True)
-    
-    for name, param in loaded_state.items():
-        if name.startswith('__S__.'):
-            if name[6:] in model_state:
-                model_state[name[6:]].copy_(param)
-            else:
-                print("{} is not in the model.".format(name[6:]))
-        else:
-            if name in model_state:
-                model_state[name].copy_(param)
-            else:
-                print("{} is not in the model.".format(name))
+        ckpt_path = hf_hub_download(repo_id="junseok520/voxsim-models", filename=ckpt_path, local_dir="./")
+    model.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+    return model
 
 
 class Score:
@@ -37,7 +51,7 @@ class Score:
 
     def __init__(
         self,
-        ckpt_path: str = "wavlm_ecapa.pt",
+        ckpt_path: str = "voxsim_wavlm_ecapa.model",
         device: str = "gpu"):
         """
         Args:
@@ -47,43 +61,15 @@ class Score:
         """
         print(f"Using device: {device}")
         self.device = device
-        self.model = load_model(ckpt_path).to(self.device)
+        self.model = loadModel(ckpt_path).to(self.device)
         self.model.eval()
     
     def score(self, inp_wavs: torch.tensor, inp_wav: torch.tensor, ref_wavs: torch.tensor, ref_wav: torch.tensor) -> torch.tensor:
-        """
-        Args:
-            wavs: audio waveform to be evaluated. When len(wavs) == 1 or 2,
-                the model processes the input as a single audio clip. The model
-                performs batch processing when len(wavs) == 3. 
-        """
-        # if len(wavs.shape) == 1:
-        #     out_wavs = wavs.unsqueeze(0).unsqueeze(0)
-        # elif len(wavs.shape) == 2:
-        #     out_wavs = wavs.unsqueeze(0)
-        # elif len(wavs.shape) == 3:
-        #     out_wavs = wavs
-        # else:
-        #     raise ValueError('Dimension of input tensor needs to be <= 3.')
-
-        if len(inp_wavs.shape) == 2:
-            bs = 1
-        elif len(inp_wavs.shape) == 3:
-            bs = inp_wavs.shape[0]
-        else:
-            raise ValueError('Dimension of input tensor needs to be <= 3.')
 
         inp_wavs = inp_wavs.reshape(-1, inp_wavs.shape[-1]).to(self.device)
         inp_wav = inp_wav.reshape(-1, inp_wav.shape[-1]).to(self.device)
         ref_wavs = ref_wavs.reshape(-1, ref_wavs.shape[-1]).to(self.device)
         ref_wav = ref_wav.reshape(-1, ref_wav.shape[-1]).to(self.device)
-        
-        # assert inp_wavs.shape[1] == 10
-        # assert ref_wavs.shape[1] == 10
-        # assert inp_wav.shape[1] == 1
-        # assert ref_wav.shape[1] == 1
-
-        # import pdb; pdb.set_trace()
 
         with torch.no_grad():
             input_emb_1 = F.normalize(self.model.forward(inp_wavs), p=2, dim=1).detach()
@@ -92,15 +78,15 @@ class Score:
             ref_emb_2 = F.normalize(self.model.forward(ref_wav), p=2, dim=1).detach()
 
             emb_size = input_emb_1.shape[-1]
-            input_emb_1 = input_emb_1.reshape(bs, -1, emb_size)
-            input_emb_2 = input_emb_2.reshape(bs, -1, emb_size)
-            ref_emb_1 = ref_emb_1.reshape(bs, -1, emb_size)
-            ref_emb_2 = ref_emb_2.reshape(bs, -1, emb_size)
+            input_emb_1 = input_emb_1.reshape(-1, emb_size)
+            input_emb_2 = input_emb_2.reshape(-1, emb_size)
+            ref_emb_1 = ref_emb_1.reshape(-1, emb_size)
+            ref_emb_2 = ref_emb_2.reshape(-1, emb_size)
 
-            score_1 = torch.mean(torch.bmm(input_emb_1, ref_emb_1.transpose(1,2)), dim=(1,2))
-            score_2 = torch.mean(torch.bmm(input_emb_2, ref_emb_2.transpose(1,2)), dim=(1,2))
+            score_1 = torch.mean(torch.matmul(input_emb_1, ref_emb_1.T))
+            score_2 = torch.mean(torch.matmul(input_emb_2, ref_emb_2.T))
             score = (score_1 + score_2) / 2
-            score = score.detach().cpu().numpy()
+            score = score.detach().cpu().item()
 
             return score
 

ssl_ecapa_model.py

@@ -4,14 +4,68 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchaudio.transforms as trans
+from torch.nn.utils.rnn import pad_sequence
+from s3prl.upstream.interfaces import UpstreamBase
+from s3prl.upstream.wavlm.WavLM import WavLM, WavLMConfig
 
-urls = {
-    'hubert_large_ll60k': "https://dl.fbaipublicfiles.com/hubert/hubert_large_ll60k.pt",
-    'xls_r_300m': "https://dl.fbaipublicfiles.com/fairseq/wav2vec/xlsr2_300m.pt",
-    'unispeech_sat': "https://huggingface.co/s3prl/converted_ckpts/resolve/main/unispeech_sat_large.pt",
-    'wavlm_base_plus': "https://huggingface.co/s3prl/converted_ckpts/resolve/main/wavlm_base_plus.pt",
-    'wavlm_large': "https://huggingface.co/s3prl/converted_ckpts/resolve/main/wavlm_large.pt",
-}
+
+''' WavLM UpstreamExpert without loading pretrained checkpoint
+'''
+
+
+class UpstreamExpert(UpstreamBase):
+    def __init__(self, cfg, **kwargs):
+        super().__init__(**kwargs)
+
+        self.cfg = WavLMConfig(torch.load(cfg))
+        self.model = WavLM(self.cfg)
+
+        self.model.feature_grad_mult = 0.0
+        self.model.encoder.layerdrop = 0.0
+
+        if len(self.hooks) == 0:
+            module_name = "self.model.encoder.layers"
+            for module_id in range(len(eval(module_name))):
+                self.add_hook(
+                    f"{module_name}[{module_id}]",
+                    lambda input, output: input[0].transpose(0, 1),
+                )
+            self.add_hook("self.model.encoder", lambda input, output: output[0])
+
+        self._init_layerdrop = self.model.encoder.layerdrop
+
+    @property
+    def layer_drop(self):
+        return self.model.encoder.layerdrop
+
+    def set_layer_drop(self, layerdrop: float = None):
+        if isinstance(layerdrop, float):
+            self.model.encoder.layerdrop = layerdrop
+        elif layerdrop is None:
+            self.model.encoder.layerdrop = self._init_layerdrop
+        else:
+            raise ValueError("layerdrop can only be float or None")
+
+    def get_downsample_rates(self, key: str) -> int:
+        return 320
+
+    def forward(self, wavs):
+        if self.cfg.normalize:
+            wavs = [F.layer_norm(wav, wav.shape) for wav in wavs]
+
+        device = wavs[0].device
+        wav_lengths = torch.LongTensor([len(wav) for wav in wavs]).to(device)
+        wav_padding_mask = ~torch.lt(
+            torch.arange(max(wav_lengths)).unsqueeze(0).to(device),
+            wav_lengths.unsqueeze(1),
+        )
+        padded_wav = pad_sequence(wavs, batch_first=True)
+
+        features, feat_padding_mask = self.model.extract_features(
+            padded_wav,
+            padding_mask=wav_padding_mask,
+            mask=False,
+        )
 
 
 ''' Res2Conv1d + BatchNorm1d + ReLU
@@ -199,7 +253,10 @@ class SSL_ECAPA_TDNN(nn.Module):
             self.feature_extract = trans.MFCC(sample_rate=sr, n_mfcc=feat_dim, log_mels=False,
                                               melkwargs=melkwargs)
         else:
-            self.feature_extract = torch.hub.load('s3prl/s3prl', feat_type)
+            if feat_type == "wavlm_large":
+                self.feature_extract = UpstreamExpert(cfg="wavlm_large_cfg.pt")
+            else:
+                raise NotImplementedError
             
             if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(self.feature_extract.model.encoder.layers[23].self_attn, "fp32_attention"):
                 self.feature_extract.model.encoder.layers[23].self_attn.fp32_attention = False