change mel func

orator/src/orator/models/voice_encoder/melspec.py

@@ -1,75 +1,78 @@
 from functools import lru_cache
 
+from scipy import signal
 import numpy as np
-import torch
-from torchaudio.transforms import MelSpectrogram
-
-from .config import VoiceEncConfig
-
-
-class ResembleMelSpectrogram(torch.nn.Module):
-    def __init__(self, hp=VoiceEncConfig()):
-        """
-        Torch implementation of Resemble's mel extraction.
-        Note that the values are NOT identical to librosa's implementation due to floating point precisions, however
-        the results are very very close. One test file gave an L1 error of just 0.005%, full results:
-            Librosa mel max:  0.871768
-            Torch mel max:    0.871768
-            Librosa mel mean: 0.316302
-            Torch mel mean:   0.316289
-            Max diff:         0.061105
-            Mean diff:        1.453384e-05
-            Percent error:    0.004595%
-        """
-        super().__init__()
-        self.melspec = MelSpectrogram(
-            hp.sample_rate,
-            n_fft=hp.n_fft,
-            win_length=hp.win_size,
-            hop_length=hp.hop_size,
-            f_min=hp.fmin,
-            f_max=hp.fmax,
-            n_mels=hp.num_mels,
-            power=1,
-            normalized=False,
-            # NOTE: Folowing librosa's default.
-            pad_mode="constant",
-            norm="slaney",
-            mel_scale="slaney",
-        )
-        self.register_buffer(
-            "stft_magnitude_min",
-            torch.FloatTensor([hp.stft_magnitude_min])
-        )
-        self.min_level_db = 20 * np.log10(hp.stft_magnitude_min)
-        self.preemphasis = hp.preemphasis
-        self.hop_size = hp.hop_size
-
-    def forward(self, wav, pad=True):
-        """
-        Args:
-            wav: [B, T]
-        """
-        if self.preemphasis > 0:
-            wav = torch.nn.functional.pad(wav, [1, 0], value=0)
-            wav = wav[..., 1:] - self.preemphasis * wav[..., :-1]
-
-        mel = self.melspec(wav)
-
-        mel = self._amp_to_db(mel)
-        mel_normed = self._normalize(mel)
-        assert not pad or mel_normed.shape[-1] == 1 + \
-            wav.shape[-1] // self.hop_size   # Sanity check
-        return mel_normed   # (M, T)
-
-    def _normalize(self, s, headroom_db=15):
-        s = (s - self.min_level_db) / (-self.min_level_db + headroom_db)
-        return s
-
-    def _amp_to_db(self, x):
-        return 20 * torch.maximum(self.stft_magnitude_min, x).log10()
+import librosa
 
 
 @lru_cache()
-def melspectrogram():
-    return ResembleMelSpectrogram()
+def mel_basis(hp):
+    assert hp.fmax <= hp.sample_rate // 2
+    return librosa.filters.mel(
+        sr=hp.sample_rate,
+        n_fft=hp.n_fft,
+        n_mels=hp.num_mels,
+        fmin=hp.fmin,
+        fmax=hp.fmax)  # -> (nmel, nfreq)
+
+
+def preemphasis(wav, hp):
+    assert hp.preemphasis != 0
+    wav = signal.lfilter([1, -hp.preemphasis], [1], wav)
+    wav = np.clip(wav, -1, 1)
+    return wav
+
+
+def melspectrogram(wav, hp, pad=True):
+    # Run through pre-emphasis
+    if hp.preemphasis > 0:
+        wav = preemphasis(wav, hp)
+        assert np.abs(wav).max() - 1 < 1e-07
+
+    # Do the stft
+    spec_complex = _stft(wav, hp, pad=pad)
+
+    # Get the magnitudes
+    spec_magnitudes = np.abs(spec_complex)
+
+    if hp.mel_power != 1.0:
+        spec_magnitudes **= hp.mel_power
+
+    # Get the mel and convert magnitudes->db
+    mel = np.dot(mel_basis(hp), spec_magnitudes)
+    if hp.mel_type == "db":
+        mel = _amp_to_db(mel, hp)
+
+    # Normalise the mel from db to 0,1
+    if hp.normalized_mels:
+        mel = _normalize(mel, hp).astype(np.float32)
+
+    assert not pad or mel.shape[1] == 1 + len(wav) // hp.hop_size   # Sanity check
+    return mel   # (M, T)
+
+
+def _stft(y, hp, pad=True):
+    # NOTE: after 0.8, pad mode defaults to constant, setting this to reflect for
+    #   historical consistency and streaming-version consistency
+    return librosa.stft(
+        y,
+        n_fft=hp.n_fft,
+        hop_length=hp.hop_size,
+        win_length=hp.win_size,
+        center=pad,
+        pad_mode="reflect",
+    )
+
+
+def _amp_to_db(x, hp):
+    return 20 * np.log10(np.maximum(hp.stft_magnitude_min, x))
+
+
+def _db_to_amp(x):
+    return np.power(10.0, x * 0.05)
+
+
+def _normalize(s, hp, headroom_db=15):
+    min_level_db = 20 * np.log10(hp.stft_magnitude_min)
+    s = (s - min_level_db) / (-min_level_db + headroom_db)
+    return s

orator/src/orator/models/voice_encoder/voice_encoder.py

@@ -269,8 +269,6 @@ class VoiceEncoder(nn.Module):
         if "rate" not in kwargs:
             kwargs["rate"] = 1.3  # Resemble's default value.
 
-        mel_func = melspectrogram()
-        mels = [mel_func(torch.from_numpy(w)
-        [None])[0].T for w in wavs]
+        mels = [melspectrogram(w, self.hp).T for w in wavs]
 
         return self.embeds_from_mels(mels, as_spk=as_spk, batch_size=batch_size, **kwargs)