mirror of
https://github.com/fumiama/Retrieval-based-Voice-Conversion-WebUI.git
synced 2026-06-05 09:10:25 +08:00
optimize(rmvpe): move mel&stft into rvc
This commit is contained in:
源文雨
@@ -1,6 +1,6 @@
|
||||
from io import BytesIO
|
||||
import os
|
||||
from typing import List, Optional, Tuple
|
||||
from typing import List, Optional, Tuple, Union
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
@@ -25,136 +25,7 @@ import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class STFT(torch.nn.Module):
|
||||
def __init__(
|
||||
self, filter_length=1024, hop_length=512, win_length=None, window="hann"
|
||||
):
|
||||
"""
|
||||
This module implements an STFT using 1D convolution and 1D transpose convolutions.
|
||||
This is a bit tricky so there are some cases that probably won't work as working
|
||||
out the same sizes before and after in all overlap add setups is tough. Right now,
|
||||
this code should work with hop lengths that are half the filter length (50% overlap
|
||||
between frames).
|
||||
|
||||
Keyword Arguments:
|
||||
filter_length {int} -- Length of filters used (default: {1024})
|
||||
hop_length {int} -- Hop length of STFT (restrict to 50% overlap between frames) (default: {512})
|
||||
win_length {[type]} -- Length of the window function applied to each frame (if not specified, it
|
||||
equals the filter length). (default: {None})
|
||||
window {str} -- Type of window to use (options are bartlett, hann, hamming, blackman, blackmanharris)
|
||||
(default: {'hann'})
|
||||
"""
|
||||
super(STFT, self).__init__()
|
||||
self.filter_length = filter_length
|
||||
self.hop_length = hop_length
|
||||
self.win_length = win_length if win_length else filter_length
|
||||
self.window = window
|
||||
self.forward_transform = None
|
||||
self.pad_amount = int(self.filter_length / 2)
|
||||
fourier_basis = np.fft.fft(np.eye(self.filter_length))
|
||||
|
||||
cutoff = int((self.filter_length / 2 + 1))
|
||||
fourier_basis = np.vstack(
|
||||
[np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
|
||||
)
|
||||
forward_basis = torch.FloatTensor(fourier_basis)
|
||||
inverse_basis = torch.FloatTensor(np.linalg.pinv(fourier_basis))
|
||||
|
||||
assert filter_length >= self.win_length
|
||||
# get window and zero center pad it to filter_length
|
||||
fft_window = get_window(window, self.win_length, fftbins=True)
|
||||
fft_window = pad_center(fft_window, size=filter_length)
|
||||
fft_window = torch.from_numpy(fft_window).float()
|
||||
|
||||
# window the bases
|
||||
forward_basis *= fft_window
|
||||
inverse_basis = (inverse_basis.T * fft_window).T
|
||||
|
||||
self.register_buffer("forward_basis", forward_basis.float())
|
||||
self.register_buffer("inverse_basis", inverse_basis.float())
|
||||
self.register_buffer("fft_window", fft_window.float())
|
||||
|
||||
def transform(self, input_data, return_phase=False):
|
||||
"""Take input data (audio) to STFT domain.
|
||||
|
||||
Arguments:
|
||||
input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
|
||||
|
||||
Returns:
|
||||
magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
|
||||
num_frequencies, num_frames)
|
||||
phase {tensor} -- Phase of STFT with shape (num_batch,
|
||||
num_frequencies, num_frames)
|
||||
"""
|
||||
input_data = F.pad(
|
||||
input_data,
|
||||
(self.pad_amount, self.pad_amount),
|
||||
mode="reflect",
|
||||
)
|
||||
forward_transform = input_data.unfold(
|
||||
1, self.filter_length, self.hop_length
|
||||
).permute(0, 2, 1)
|
||||
forward_transform = torch.matmul(self.forward_basis, forward_transform)
|
||||
cutoff = int((self.filter_length / 2) + 1)
|
||||
real_part = forward_transform[:, :cutoff, :]
|
||||
imag_part = forward_transform[:, cutoff:, :]
|
||||
magnitude = torch.sqrt(real_part**2 + imag_part**2)
|
||||
if return_phase:
|
||||
phase = torch.atan2(imag_part.data, real_part.data)
|
||||
return magnitude, phase
|
||||
else:
|
||||
return magnitude
|
||||
|
||||
def inverse(self, magnitude, phase):
|
||||
"""Call the inverse STFT (iSTFT), given magnitude and phase tensors produced
|
||||
by the ```transform``` function.
|
||||
|
||||
Arguments:
|
||||
magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
|
||||
num_frequencies, num_frames)
|
||||
phase {tensor} -- Phase of STFT with shape (num_batch,
|
||||
num_frequencies, num_frames)
|
||||
|
||||
Returns:
|
||||
inverse_transform {tensor} -- Reconstructed audio given magnitude and phase. Of
|
||||
shape (num_batch, num_samples)
|
||||
"""
|
||||
cat = torch.cat(
|
||||
[magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
|
||||
)
|
||||
fold = torch.nn.Fold(
|
||||
output_size=(1, (cat.size(-1) - 1) * self.hop_length + self.filter_length),
|
||||
kernel_size=(1, self.filter_length),
|
||||
stride=(1, self.hop_length),
|
||||
)
|
||||
inverse_transform = torch.matmul(self.inverse_basis, cat)
|
||||
inverse_transform = fold(inverse_transform)[
|
||||
:, 0, 0, self.pad_amount : -self.pad_amount
|
||||
]
|
||||
window_square_sum = (
|
||||
self.fft_window.pow(2).repeat(cat.size(-1), 1).T.unsqueeze(0)
|
||||
)
|
||||
window_square_sum = fold(window_square_sum)[
|
||||
:, 0, 0, self.pad_amount : -self.pad_amount
|
||||
]
|
||||
inverse_transform /= window_square_sum
|
||||
return inverse_transform
|
||||
|
||||
def forward(self, input_data):
|
||||
"""Take input data (audio) to STFT domain and then back to audio.
|
||||
|
||||
Arguments:
|
||||
input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
|
||||
|
||||
Returns:
|
||||
reconstruction {tensor} -- Reconstructed audio given magnitude and phase. Of
|
||||
shape (num_batch, num_samples)
|
||||
"""
|
||||
self.magnitude, self.phase = self.transform(input_data, return_phase=True)
|
||||
reconstruction = self.inverse(self.magnitude, self.phase)
|
||||
return reconstruction
|
||||
|
||||
from rvc.f0.mel import MelSpectrogram
|
||||
|
||||
from time import time as ttime
|
||||
|
||||
@@ -412,86 +283,6 @@ class E2E(nn.Module):
|
||||
return x
|
||||
|
||||
|
||||
from librosa.filters import mel
|
||||
|
||||
|
||||
class MelSpectrogram(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
is_half,
|
||||
n_mel_channels,
|
||||
sampling_rate,
|
||||
win_length,
|
||||
hop_length,
|
||||
n_fft=None,
|
||||
mel_fmin=0,
|
||||
mel_fmax=None,
|
||||
clamp=1e-5,
|
||||
):
|
||||
super().__init__()
|
||||
n_fft = win_length if n_fft is None else n_fft
|
||||
self.hann_window = {}
|
||||
mel_basis = mel(
|
||||
sr=sampling_rate,
|
||||
n_fft=n_fft,
|
||||
n_mels=n_mel_channels,
|
||||
fmin=mel_fmin,
|
||||
fmax=mel_fmax,
|
||||
htk=True,
|
||||
)
|
||||
mel_basis = torch.from_numpy(mel_basis).float()
|
||||
self.register_buffer("mel_basis", mel_basis)
|
||||
self.n_fft = win_length if n_fft is None else n_fft
|
||||
self.hop_length = hop_length
|
||||
self.win_length = win_length
|
||||
self.sampling_rate = sampling_rate
|
||||
self.n_mel_channels = n_mel_channels
|
||||
self.clamp = clamp
|
||||
self.is_half = is_half
|
||||
|
||||
def forward(self, audio, keyshift=0, speed=1, center=True):
|
||||
factor = 2 ** (keyshift / 12)
|
||||
n_fft_new = int(np.round(self.n_fft * factor))
|
||||
win_length_new = int(np.round(self.win_length * factor))
|
||||
hop_length_new = int(np.round(self.hop_length * speed))
|
||||
keyshift_key = str(keyshift) + "_" + str(audio.device)
|
||||
if keyshift_key not in self.hann_window:
|
||||
self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(
|
||||
audio.device
|
||||
)
|
||||
if "privateuseone" in str(audio.device):
|
||||
if not hasattr(self, "stft"):
|
||||
self.stft = STFT(
|
||||
filter_length=n_fft_new,
|
||||
hop_length=hop_length_new,
|
||||
win_length=win_length_new,
|
||||
window="hann",
|
||||
).to(audio.device)
|
||||
magnitude = self.stft.transform(audio)
|
||||
else:
|
||||
fft = torch.stft(
|
||||
audio,
|
||||
n_fft=n_fft_new,
|
||||
hop_length=hop_length_new,
|
||||
win_length=win_length_new,
|
||||
window=self.hann_window[keyshift_key],
|
||||
center=center,
|
||||
return_complex=True,
|
||||
)
|
||||
magnitude = torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
|
||||
if keyshift != 0:
|
||||
size = self.n_fft // 2 + 1
|
||||
resize = magnitude.size(1)
|
||||
if resize < size:
|
||||
magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
|
||||
magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
|
||||
mel_output = torch.matmul(self.mel_basis, magnitude)
|
||||
if self.is_half == True:
|
||||
mel_output = mel_output.half()
|
||||
log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
|
||||
return log_mel_spec
|
||||
|
||||
|
||||
class RMVPE:
|
||||
def __init__(self, model_path: str, is_half, device=None, use_jit=False):
|
||||
self.resample_kernel = {}
|
||||
@@ -501,7 +292,14 @@ class RMVPE:
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
self.device = device
|
||||
self.mel_extractor = MelSpectrogram(
|
||||
is_half, 128, 16000, 1024, 160, None, 30, 8000
|
||||
is_half=is_half,
|
||||
n_mel_channels=128,
|
||||
sampling_rate=16000,
|
||||
win_length=1024,
|
||||
hop_length=160,
|
||||
mel_fmin=30,
|
||||
mel_fmax=8000,
|
||||
device=device,
|
||||
).to(device)
|
||||
if "privateuseone" in str(device):
|
||||
import onnxruntime as ort
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import torch
|
||||
from infer.lib.rmvpe import STFT
|
||||
from rvc.f0.stft import STFT
|
||||
from torch.nn.functional import conv1d, conv2d
|
||||
from typing import Union, Optional
|
||||
from .utils import linspace, temperature_sigmoid, amp_to_db
|
||||
|
||||
Reference in New Issue
Block a user