ORI-Muchim/BlueArchiveTTS

1import math2import os3from packaging import version4import random5import torch6from torch import nn7import torch.nn.functional as F8import torch.utils.data9import numpy as np10import librosa11import librosa.util as librosa_util12from librosa.util import normalize, pad_center, tiny13from scipy.signal import get_window14from scipy.io.wavfile import read15from librosa.filters import mel as librosa_mel_fn16 17MAX_WAV_VALUE = 32768.018 19 20def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):21    """22    PARAMS23    ------24    C: compression factor25    """26    return torch.log(torch.clamp(x, min=clip_val) * C)27 28 29def dynamic_range_decompression_torch(x, C=1):30    """31    PARAMS32    ------33    C: compression factor used to compress34    """35    return torch.exp(x) / C36 37 38def spectral_normalize_torch(magnitudes):39    output = dynamic_range_compression_torch(magnitudes)40    return output41 42 43def spectral_de_normalize_torch(magnitudes):44    output = dynamic_range_decompression_torch(magnitudes)45    return output46 47 48mel_basis = {}49hann_window = {}50 51 52def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):53    if torch.min(y) < -1.:54        print('min value is ', torch.min(y))55    if torch.max(y) > 1.:56        print('max value is ', torch.max(y))57 58    global hann_window59    dtype_device = str(y.dtype) + '_' + str(y.device)60    wnsize_dtype_device = str(win_size) + '_' + dtype_device61    if wnsize_dtype_device not in hann_window:62        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)63 64    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')65    y = y.squeeze(1)66 67    if version.parse(torch.__version__) >= version.parse("2"):68        spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],69                          center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)70    else:71        spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],72                      center=center, pad_mode='reflect', normalized=False, onesided=True)73        74    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)75    return spec76 77 78def spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax):79    global mel_basis80    dtype_device = str(spec.dtype) + '_' + str(spec.device)81    fmax_dtype_device = str(fmax) + '_' + dtype_device82    if fmax_dtype_device not in mel_basis:83        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)84        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=spec.dtype, device=spec.device)85    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)86    spec = spectral_normalize_torch(spec)87    return spec88 89 90def mel_spectrogram_torch(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):91    if torch.min(y) < -1.:92        print('min value is ', torch.min(y))93    if torch.max(y) > 1.:94        print('max value is ', torch.max(y))95 96    global mel_basis, hann_window97    dtype_device = str(y.dtype) + '_' + str(y.device)98    fmax_dtype_device = str(fmax) + '_' + dtype_device99    wnsize_dtype_device = str(win_size) + '_' + dtype_device100    if fmax_dtype_device not in mel_basis:101        mel = librosa_mel_fn(sampling_rate, n_fft, num_mels, fmin, fmax)102        mel_basis[fmax_dtype_device] = torch.from_numpy(mel).to(dtype=y.dtype, device=y.device)103    if wnsize_dtype_device not in hann_window:104        hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)105 106    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')107    y = y.squeeze(1)108 109    if version.parse(torch.__version__) >= version.parse("2"):110        spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],111                          center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=False)112    else:113        spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],114                      center=center, pad_mode='reflect', normalized=False, onesided=True)115    '''116    #- reserve : from https://github.com/jaywalnut310/vits/issues/15#issuecomment-1084148441117    with autocast(enabled=False):118        y = y.float()119        spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],120                        center=center, pad_mode='reflect', normalized=False, onesided=True)121    '''122 123    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)124 125    spec = torch.matmul(mel_basis[fmax_dtype_device], spec)126    spec = spectral_normalize_torch(spec)127 128    return spec129