vits-jax/mel_processing.py at main · hyunwoo3235/vits-jax · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import jax.numpy as jnp
from audax.core.functional import melscale_fbanks, spectrogram, apply_melscale

MAX_WAV_VALUE = 32768.0


def dynamic_range_compression_jax(x, C=1, clip_val=1e-5):
    """
    PARAMS
    ------
    C: compression factor
    """
    return jnp.log(jnp.clip(x, a_min=clip_val, a_max=None) * C)


def dynamic_range_decompression_jax(x, C=1):
    """
    PARAMS
    ------
    C: compression factor used to compress
    """
    return jnp.exp(x) / C


def spectral_normalize_jax(magnitudes):
    output = dynamic_range_compression_jax(magnitudes)
    return output


def spectral_de_normalize_jax(magnitudes):
    output = dynamic_range_decompression_jax(magnitudes)
    return output


mel_basis = None
hann_window = None


def mel_setup(n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax):
    global hann_window, mel_basis
    if hann_window is None:
        hann_window = jnp.hanning(win_size)
    if mel_basis is None:
        mel_basis = melscale_fbanks(
            n_freqs=n_fft // 2 + 1,
            n_mels=num_mels,
            sample_rate=sampling_rate,
            f_min=fmin,
            f_max=fmax,
            norm="slaney",
            mel_scale="slaney",
        )


def spectrogram_jax(y, n_fft, hop_size, win_size, center=False):
    global hann_window
    if hann_window is None:
        hann_window = jnp.hanning(win_size)

    y = jnp.pad(
        y,
        ((0, 0), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), (0, 0)),
        mode="reflect",
    )

    spec = spectrogram(
        y,
        pad=0,
        window=hann_window,
        n_fft=n_fft,
        hop_length=hop_size,
        win_length=win_size,
        power=1.0,
        normalized=False,
        center=center,
        onesided=True,
    )
    return spec


def spec_to_mel_jax(spec, n_fft, num_mels, sampling_rate, fmin, fmax):
    global mel_basis
    if mel_basis is None:
        mel_basis = melscale_fbanks(
            n_freqs=n_fft // 2 + 1,
            n_mels=num_mels,
            sample_rate=sampling_rate,
            f_min=fmin,
            f_max=fmax,
            norm="slaney",
            mel_scale="slaney",
        )
    mel = apply_melscale(spec, mel_basis)
    mel = spectral_normalize_jax(mel)

    return mel


def mel_spectrogram_jax(
    y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False
):
    global hann_window, mel_basis
    if hann_window is None:
        hann_window = jnp.hanning(win_size)
    if mel_basis is None:
        mel_basis = melscale_fbanks(
            n_freqs=n_fft // 2 + 1,
            n_mels=num_mels,
            sample_rate=sampling_rate,
            f_min=fmin,
            f_max=fmax,
            norm="slaney",
            mel_scale="slaney",
        )

    y = jnp.pad(
        y,
        ((0, 0), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), (0, 0)),
        mode="reflect",
    )

    spec = spectrogram(
        y,
        pad=0,
        window=hann_window,
        n_fft=n_fft,
        hop_length=hop_size,
        win_length=win_size,
        power=1.0,
        normalized=False,
        center=center,
        onesided=True,
    )

    mel = apply_melscale(spec, mel_basis)
    mel = spectral_normalize_jax(mel)

    return mel