vad trimming for ref wavs

.gitignore

@@ -0,0 +1 @@
+__pycache__

chatterbox/src/chatterbox/tts.py

@@ -2,10 +2,12 @@ from dataclasses import dataclass
 from pathlib import Path
 
 import librosa
+import numpy as np
 import torch
 import perth
 import torch.nn.functional as F
 from huggingface_hub import hf_hub_download
+from silero_vad import load_silero_vad, get_speech_timestamps
 
 from .models.t3 import T3
 from .models.s3tokenizer import S3_SR, drop_invalid_tokens
@@ -121,6 +123,7 @@ class ChatterboxTTS:
         self.device = device
         self.conds = conds
         self.watermarker = perth.PerthImplicitWatermarker()
+        self.silero_vad = load_silero_vad()
 
     @classmethod
     def from_local(cls, ckpt_dir, device) -> 'ChatterboxTTS':
@@ -162,11 +165,32 @@ class ChatterboxTTS:
 
         return cls.from_local(Path(local_path).parent, device)
 
+    def trim_excess_silence(self, wav, sr):
+        "Trim excess silence from speech. Input must be a multiple of 16kHz."
+        assert sr % 16_000 == 0, "Silero requires an integer multiple of 16kHz"
+
+        # Get VAD as sample-level bool array
+        silero_regions = get_speech_timestamps(wav, self.silero_vad, sampling_rate=sr)
+        vad = np.zeros_like(wav)
+        for region in silero_regions:
+            vad[region["start"]:region["end"]] = 1
+
+        # Dilate VAD
+        max_silence_ms = 400
+        cfilter = np.ones(int(sr * max_silence_ms / (2 * 1000)))
+        dilated_vad = np.convolve(vad, cfilter, mode="same") > 0
+
+        # Trim out silence
+        return wav[dilated_vad]
+
     def prepare_conditionals(self, wav_fpath, exaggeration=0.5):
-        ## Load reference wav
-        s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
+        # Load reference wav at high SR and trim silence
+        ref_wav, highres_sr = librosa.load(wav_fpath, sr=48_000)
+        ref_wav = self.trim_excess_silence(ref_wav, highres_sr)
 
-        ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR)
+        # Resample down
+        s3gen_ref_wav = librosa.resample(ref_wav, orig_sr=highres_sr, target_sr=S3GEN_SR)
+        ref_16k_wav = librosa.resample(ref_wav, orig_sr=highres_sr, target_sr=S3_SR)
 
         s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
         s3gen_ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)