Hilley/ChatTTS-OpenVoice

1import torch2import torch.utils.data3from librosa.filters import mel as librosa_mel_fn4 5MAX_WAV_VALUE = 32768.06 7 8def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):9    """10    PARAMS11    ------12    C: compression factor13    """14    return torch.log(torch.clamp(x, min=clip_val) * C)15 16 17def dynamic_range_decompression_torch(x, C=1):18    """19    PARAMS20    ------21    C: compression factor used to compress22    """23    return torch.exp(x) / C24 25 26def spectral_normalize_torch(magnitudes):27    output = dynamic_range_compression_torch(magnitudes)28    return output29 30 31def spectral_de_normalize_torch(magnitudes):32    output = dynamic_range_decompression_torch(magnitudes)33    return output34 35 36mel_basis = {}37hann_window = {}38 39 40def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):41    if torch.min(y) < -1.1:42        print("min value is ", torch.min(y))43    if torch.max(y) > 1.1:44        print("max value is ", torch.max(y))45 46    global hann_window47    dtype_device = str(y.dtype) + "_" + str(y.device)48    wnsize_dtype_device = str(win_size) + "_" + dtype_device49    if wnsize_dtype_device not in hann_window:50        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(51            dtype=y.dtype, device=y.device52        )53 54    y = torch.nn.functional.pad(55        y.unsqueeze(1),56        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),57        mode="reflect",58    )59    y = y.squeeze(1)60 61    spec = torch.stft(62        y,63        n_fft,64        hop_length=hop_size,65        win_length=win_size,66        window=hann_window[wnsize_dtype_device],67        center=center,68        pad_mode="reflect",69        normalized=False,70        onesided=True,71        return_complex=False,72    )73 74    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)75    return spec76 77 78def spectrogram_torch_conv(y, n_fft, sampling_rate, hop_size, win_size, center=False):79    # if torch.min(y) < -1.:80    #     print('min value is ', torch.min(y))81    # if torch.max(y) > 1.:82    #     print('max value is ', torch.max(y))83 84    global hann_window85    dtype_device = str(y.dtype) + '_' + str(y.device)86    wnsize_dtype_device = str(win_size) + '_' + dtype_device87    if wnsize_dtype_device not in hann_window:88        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)89 90    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')91    92    # ******************** original ************************#93    # y = y.squeeze(1)94    # spec1 = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],95    #                   center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)96 97    # ******************** ConvSTFT ************************#98    freq_cutoff = n_fft // 2 + 199    fourier_basis = torch.view_as_real(torch.fft.fft(torch.eye(n_fft)))100    forward_basis = fourier_basis[:freq_cutoff].permute(2, 0, 1).reshape(-1, 1, fourier_basis.shape[1])101    forward_basis = forward_basis * torch.as_tensor(librosa.util.pad_center(torch.hann_window(win_size), size=n_fft)).float()102 103    import torch.nn.functional as F104 105    # if center:106    #     signal = F.pad(y[:, None, None, :], (n_fft // 2, n_fft // 2, 0, 0), mode = 'reflect').squeeze(1)107    assert center is False108 109    forward_transform_squared = F.conv1d(y, forward_basis.to(y.device), stride = hop_size)110    spec2 = torch.stack([forward_transform_squared[:, :freq_cutoff, :], forward_transform_squared[:, freq_cutoff:, :]], dim = -1)111 112 113    # ******************** Verification ************************#114    spec1 = torch.stft(y.squeeze(1), n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],115                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)116    assert torch.allclose(spec1, spec2, atol=1e-4)117 118    spec = torch.sqrt(spec2.pow(2).sum(-1) + 1e-6)119    return spec120 121 122def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):123    global mel_basis124    dtype_device = str(spec.dtype) + "_" + str(spec.device)125    fmax_dtype_device = str(fmax) + "_" + dtype_device126    if fmax_dtype_device not in mel_basis:127        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)128        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(129            dtype=spec.dtype, device=spec.device130        )131    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)132    spec = spectral_normalize_torch(spec)133    return spec134 135 136def mel_spectrogram_torch(137    y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False138):139    if torch.min(y) < -1.0:140        print("min value is ", torch.min(y))141    if torch.max(y) > 1.0:142        print("max value is ", torch.max(y))143 144    global mel_basis, hann_window145    dtype_device = str(y.dtype) + "_" + str(y.device)146    fmax_dtype_device = str(fmax) + "_" + dtype_device147    wnsize_dtype_device = str(win_size) + "_" + dtype_device148    if fmax_dtype_device not in mel_basis:149        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)150        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(151            dtype=y.dtype, device=y.device152        )153    if wnsize_dtype_device not in hann_window:154        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(155            dtype=y.dtype, device=y.device156        )157 158    y = torch.nn.functional.pad(159        y.unsqueeze(1),160        (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)),161        mode="reflect",162    )163    y = y.squeeze(1)164 165    spec = torch.stft(166        y,167        n_fft,168        hop_length=hop_size,169        win_length=win_size,170        window=hann_window[wnsize_dtype_device],171        center=center,172        pad_mode="reflect",173        normalized=False,174        onesided=True,175        return_complex=False,176    )177 178    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)179 180    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)181    spec = spectral_normalize_torch(spec)182 183    return spec