| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
|
|
|
|
| @dataclass |
| class VocosFbankConfig: |
| sampling_rate: int = 24000 |
| n_mels: int = 100 |
| n_fft: int = 1024 |
| hop_length: int = 256 |
|
|
|
|
| def compute_num_frames(num_samples: int, hop_length: int) -> int: |
| return int((num_samples + hop_length // 2) // hop_length) |
|
|
|
|
| class LocalVocosFbank: |
| def __init__(self) -> None: |
| self.config = VocosFbankConfig() |
| self.window = torch.hann_window(self.config.n_fft) |
| self.mel_basis = _create_mel_filterbank( |
| sample_rate=self.config.sampling_rate, |
| n_fft=self.config.n_fft, |
| n_mels=self.config.n_mels, |
| ) |
|
|
| def extract(self, samples: torch.Tensor, sampling_rate: int) -> torch.Tensor: |
| if sampling_rate != self.config.sampling_rate: |
| raise ValueError( |
| f"Mismatched sampling rate: expected {self.config.sampling_rate}, got {sampling_rate}" |
| ) |
| if samples.ndim == 1: |
| samples = samples.unsqueeze(0) |
| if samples.ndim != 2: |
| raise ValueError(f"Expected waveform shape [C, T], got {tuple(samples.shape)}") |
| if samples.shape[0] == 2: |
| samples = samples.mean(dim=0, keepdim=True) |
|
|
| stft = torch.stft( |
| samples, |
| n_fft=self.config.n_fft, |
| hop_length=self.config.hop_length, |
| win_length=self.config.n_fft, |
| window=self.window.to(samples.device), |
| center=True, |
| pad_mode="reflect", |
| return_complex=True, |
| ) |
| spec = stft.abs() |
| mel = torch.matmul(self.mel_basis.to(samples.device).t(), spec).clamp(min=1e-7).log() |
| mel = mel.reshape(-1, mel.shape[-1]).t() |
| num_frames = compute_num_frames(samples.shape[1], self.config.hop_length) |
|
|
| if mel.shape[0] > num_frames: |
| mel = mel[:num_frames] |
| elif mel.shape[0] < num_frames: |
| mel = torch.nn.functional.pad( |
| mel.unsqueeze(0), |
| (0, 0, 0, num_frames - mel.shape[0]), |
| mode="replicate", |
| ).squeeze(0) |
| return mel |
|
|
|
|
| def _hz_to_mel(freq: torch.Tensor) -> torch.Tensor: |
| return 2595.0 * torch.log10(1.0 + freq / 700.0) |
|
|
|
|
| def _mel_to_hz(mels: torch.Tensor) -> torch.Tensor: |
| return 700.0 * (torch.pow(10.0, mels / 2595.0) - 1.0) |
|
|
|
|
| def _create_mel_filterbank(sample_rate: int, n_fft: int, n_mels: int) -> torch.Tensor: |
| n_freqs = n_fft // 2 + 1 |
| all_freqs = torch.linspace(0, sample_rate // 2, n_freqs) |
| m_min = _hz_to_mel(torch.tensor(0.0)) |
| m_max = _hz_to_mel(torch.tensor(float(sample_rate // 2))) |
| m_pts = torch.linspace(m_min, m_max, n_mels + 2) |
| f_pts = _mel_to_hz(m_pts) |
|
|
| f_diff = f_pts[1:] - f_pts[:-1] |
| slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1) |
| down_slopes = -slopes[:, :-2] / f_diff[:-1] |
| up_slopes = slopes[:, 2:] / f_diff[1:] |
| fb = torch.maximum(torch.zeros(1), torch.minimum(down_slopes, up_slopes)) |
|
|
| if (fb.max(dim=0).values == 0.0).any(): |
| raise ValueError("Mel filterbank has empty filters") |
| return fb |
|
|
|
|
| def _resample_linear(wav: torch.Tensor, orig_freq: int, new_freq: int) -> torch.Tensor: |
| if orig_freq == new_freq: |
| return wav |
| old_len = wav.shape[-1] |
| new_len = max(1, int(round(old_len * new_freq / orig_freq))) |
| old_pos = np.arange(old_len, dtype=np.float64) |
| new_pos = np.linspace(0, old_len - 1, new_len, dtype=np.float64) |
| channels = [] |
| for channel in wav.cpu().numpy(): |
| channels.append(np.interp(new_pos, old_pos, channel).astype(np.float32)) |
| return torch.from_numpy(np.stack(channels, axis=0)) |
|
|
|
|
| def load_prompt_wav(prompt_wav: str | Path, sampling_rate: int) -> torch.Tensor: |
| wav_np, sr = sf.read(str(prompt_wav), always_2d=True, dtype="float32") |
| wav = torch.from_numpy(wav_np.T.copy()) |
| if sr != sampling_rate: |
| wav = _resample_linear(wav, orig_freq=sr, new_freq=sampling_rate) |
| return wav |
|
|
|
|
| def rms_norm(wav: torch.Tensor, target_rms: float): |
| wav_rms = torch.sqrt(torch.mean(torch.square(wav))) |
| if wav_rms < target_rms: |
| wav = wav * target_rms / wav_rms |
| return wav, wav_rms |
|
|
|
|
| def load_local_vocos(vocoder_dir: str | Path): |
| from scripts.local_vocos import LocalVocos |
|
|
| vocoder_dir = Path(vocoder_dir) |
| vocoder = LocalVocos() |
| try: |
| state_dict = torch.load( |
| str(vocoder_dir / "pytorch_model.bin"), |
| weights_only=True, |
| map_location="cpu", |
| ) |
| except TypeError: |
| state_dict = torch.load(str(vocoder_dir / "pytorch_model.bin"), map_location="cpu") |
| state_dict = { |
| key: value |
| for key, value in state_dict.items() |
| if key.startswith(("backbone.", "head.")) |
| } |
| vocoder.load_state_dict(state_dict) |
| return vocoder |
|
|
