| """Feature extractor for the ECAPA-TDNN speaker encoder. |
| |
| Converts raw audio waveforms into log-mel spectrograms suitable for the |
| ECAPA-TDNN speaker encoder model. |
| """ |
|
|
| import numpy as np |
| import torch |
| from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin |
|
|
|
|
| class EcapaTdnnFeatureExtractor(FeatureExtractionMixin): |
| r""" |
| Feature extractor for ECAPA-TDNN speaker encoder models. |
| |
| Converts raw audio waveforms to 128-bin log-mel spectrograms matching the |
| Qwen3-TTS preprocessing pipeline. |
| |
| Args: |
| sample_rate (`int`, defaults to 24000): |
| Target sample rate in Hz. Audio will be resampled if needed. |
| n_fft (`int`, defaults to 1024): |
| FFT window size. |
| hop_length (`int`, defaults to 256): |
| Hop length between STFT frames. |
| n_mels (`int`, defaults to 128): |
| Number of mel-frequency bins. |
| fmin (`float`, defaults to 0): |
| Minimum frequency for mel filterbank. |
| fmax (`float`, defaults to 12000): |
| Maximum frequency for mel filterbank. |
| """ |
|
|
| model_input_names = ["input_values"] |
|
|
| def __init__( |
| self, |
| sample_rate=24000, |
| n_fft=1024, |
| hop_length=256, |
| n_mels=128, |
| fmin=0, |
| fmax=12000, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.sample_rate = sample_rate |
| self.sampling_rate = sample_rate |
| self.n_fft = n_fft |
| self.hop_length = hop_length |
| self.n_mels = n_mels |
| self.fmin = fmin |
| self.fmax = fmax |
|
|
| def __call__(self, raw_speech, sampling_rate=None, return_tensors="pt", **kwargs): |
| """ |
| Process raw audio waveform(s) into log-mel spectrogram features. |
| |
| Args: |
| raw_speech (`np.ndarray`, `list[np.ndarray]`, or file path `str`): |
| Raw audio waveform(s) as float32 numpy array(s), or a file path. |
| sampling_rate (`int`, *optional*): |
| Sample rate of the input audio. Resampled to ``self.sample_rate`` |
| if different. |
| return_tensors (`str`, defaults to ``"pt"``): |
| Return type — ``"pt"`` for PyTorch tensors. |
| |
| Returns: |
| ``BatchFeature`` with ``input_values`` key containing the log-mel |
| spectrogram tensor of shape ``(batch, time, n_mels)``. |
| """ |
| |
| if isinstance(raw_speech, str): |
| import librosa |
| raw_speech, sampling_rate = librosa.load(raw_speech, sr=None, mono=True) |
|
|
| if isinstance(raw_speech, np.ndarray) and raw_speech.ndim == 1: |
| raw_speech = [raw_speech] |
|
|
| features = [] |
| for audio in raw_speech: |
| if isinstance(audio, str): |
| import librosa |
| audio, sampling_rate = librosa.load(audio, sr=None, mono=True) |
|
|
| mel = self._compute_mel(audio, sampling_rate or self.sample_rate) |
| features.append(mel) |
|
|
| |
| max_len = max(f.shape[1] for f in features) |
| padded = [] |
| for f in features: |
| if f.shape[1] < max_len: |
| f = torch.nn.functional.pad(f, (0, 0, 0, max_len - f.shape[1])) |
| padded.append(f) |
|
|
| input_values = torch.cat(padded, dim=0) |
| return BatchFeature({"input_values": input_values}) |
|
|
| def _compute_mel(self, audio, sr): |
| """Compute 128-bin log-mel spectrogram matching Qwen3-TTS requirements.""" |
| import librosa |
| from librosa.filters import mel as librosa_mel_fn |
|
|
| if isinstance(audio, torch.Tensor): |
| audio = audio.numpy() |
|
|
| if sr != self.sample_rate: |
| audio = librosa.resample( |
| audio.astype(np.float32), orig_sr=sr, target_sr=self.sample_rate |
| ) |
|
|
| y = torch.from_numpy(audio).unsqueeze(0).float() |
| mel_basis = torch.from_numpy( |
| librosa_mel_fn( |
| sr=self.sample_rate, |
| n_fft=self.n_fft, |
| n_mels=self.n_mels, |
| fmin=self.fmin, |
| fmax=self.fmax, |
| ) |
| ).float() |
|
|
| padding = (self.n_fft - self.hop_length) // 2 |
| y = torch.nn.functional.pad( |
| y.unsqueeze(1), (padding, padding), mode="reflect" |
| ).squeeze(1) |
| hann = torch.hann_window(self.n_fft) |
| spec = torch.stft( |
| y, |
| self.n_fft, |
| hop_length=self.hop_length, |
| win_length=self.n_fft, |
| window=hann, |
| center=False, |
| return_complex=True, |
| ) |
| spec = torch.abs(spec) |
| mel = torch.matmul(mel_basis, spec) |
| mel = torch.log(torch.clamp(mel, min=1e-5)) |
| return mel.transpose(1, 2) |
|
|
