109 lines
4.0 KiB
Python
109 lines
4.0 KiB
Python
from typing import Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from modules.F0Predictor.F0Predictor import F0Predictor
|
|
|
|
from .fcpe.model import FCPEInfer
|
|
|
|
|
|
class FCPEF0Predictor(F0Predictor):
|
|
def __init__(self, hop_length=512, f0_min=50, f0_max=1100, dtype=torch.float32, device=None, sampling_rate=44100,
|
|
threshold=0.05):
|
|
self.fcpe = FCPEInfer(model_path="pretrain/fcpe.pt", device=device, dtype=dtype)
|
|
self.hop_length = hop_length
|
|
self.f0_min = f0_min
|
|
self.f0_max = f0_max
|
|
if device is None:
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
else:
|
|
self.device = device
|
|
self.threshold = threshold
|
|
self.sampling_rate = sampling_rate
|
|
self.dtype = dtype
|
|
self.name = "fcpe"
|
|
|
|
def repeat_expand(
|
|
self, content: Union[torch.Tensor, np.ndarray], target_len: int, mode: str = "nearest"
|
|
):
|
|
ndim = content.ndim
|
|
|
|
if content.ndim == 1:
|
|
content = content[None, None]
|
|
elif content.ndim == 2:
|
|
content = content[None]
|
|
|
|
assert content.ndim == 3
|
|
|
|
is_np = isinstance(content, np.ndarray)
|
|
if is_np:
|
|
content = torch.from_numpy(content)
|
|
|
|
results = torch.nn.functional.interpolate(content, size=target_len, mode=mode)
|
|
|
|
if is_np:
|
|
results = results.numpy()
|
|
|
|
if ndim == 1:
|
|
return results[0, 0]
|
|
elif ndim == 2:
|
|
return results[0]
|
|
|
|
def post_process(self, x, sampling_rate, f0, pad_to):
|
|
if isinstance(f0, np.ndarray):
|
|
f0 = torch.from_numpy(f0).float().to(x.device)
|
|
|
|
if pad_to is None:
|
|
return f0
|
|
|
|
f0 = self.repeat_expand(f0, pad_to)
|
|
|
|
vuv_vector = torch.zeros_like(f0)
|
|
vuv_vector[f0 > 0.0] = 1.0
|
|
vuv_vector[f0 <= 0.0] = 0.0
|
|
|
|
# 去掉0频率, 并线性插值
|
|
nzindex = torch.nonzero(f0).squeeze()
|
|
f0 = torch.index_select(f0, dim=0, index=nzindex).cpu().numpy()
|
|
time_org = self.hop_length / sampling_rate * nzindex.cpu().numpy()
|
|
time_frame = np.arange(pad_to) * self.hop_length / sampling_rate
|
|
|
|
vuv_vector = F.interpolate(vuv_vector[None, None, :], size=pad_to)[0][0]
|
|
|
|
if f0.shape[0] <= 0:
|
|
return torch.zeros(pad_to, dtype=torch.float, device=x.device).cpu().numpy(), vuv_vector.cpu().numpy()
|
|
if f0.shape[0] == 1:
|
|
return (torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[
|
|
0]).cpu().numpy(), vuv_vector.cpu().numpy()
|
|
|
|
# 大概可以用 torch 重写?
|
|
f0 = np.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
|
|
# vuv_vector = np.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
|
|
|
|
return f0, vuv_vector.cpu().numpy()
|
|
|
|
def compute_f0(self, wav, p_len=None):
|
|
x = torch.FloatTensor(wav).to(self.dtype).to(self.device)
|
|
if p_len is None:
|
|
p_len = x.shape[0] // self.hop_length
|
|
else:
|
|
assert abs(p_len - x.shape[0] // self.hop_length) < 4, "pad length error"
|
|
f0 = self.fcpe(x, sr=self.sampling_rate, threshold=self.threshold)[0,:,0]
|
|
if torch.all(f0 == 0):
|
|
rtn = f0.cpu().numpy() if p_len is None else np.zeros(p_len)
|
|
return rtn, rtn
|
|
return self.post_process(x, self.sampling_rate, f0, p_len)[0]
|
|
|
|
def compute_f0_uv(self, wav, p_len=None):
|
|
x = torch.FloatTensor(wav).to(self.dtype).to(self.device)
|
|
if p_len is None:
|
|
p_len = x.shape[0] // self.hop_length
|
|
else:
|
|
assert abs(p_len - x.shape[0] // self.hop_length) < 4, "pad length error"
|
|
f0 = self.fcpe(x, sr=self.sampling_rate, threshold=self.threshold)[0,:,0]
|
|
if torch.all(f0 == 0):
|
|
rtn = f0.cpu().numpy() if p_len is None else np.zeros(p_len)
|
|
return rtn, rtn
|
|
return self.post_process(x, self.sampling_rate, f0, p_len) |