166 lines
6.3 KiB
Python
166 lines
6.3 KiB
Python
import os
|
|
import logging
|
|
from .lib import BASE_MODELS_DIR
|
|
from .lib.model_utils import change_rms, load_hubert
|
|
from .config import config
|
|
from .pitch_extraction import FeatureExtractor
|
|
from .lib.utils import gc_collect, get_filenames
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
import fairseq
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
if config.dml == True:
|
|
|
|
def forward_dml(ctx, x, scale):
|
|
ctx.scale = scale
|
|
res = x.clone().detach()
|
|
return res
|
|
|
|
fairseq.modules.grad_multiply.GradMultiply.forward = forward_dml
|
|
|
|
|
|
# config.device=torch.device("cpu")########强制cpu测试
|
|
# config.is_half=False########强制cpu测试
|
|
class RVC(FeatureExtractor):
|
|
def __init__(self, model_path, config, onnx=False, device=None):
|
|
|
|
cpt = torch.load(model_path, map_location="cpu")
|
|
tgt_sr = cpt["config"][-1]
|
|
cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0] # n_spk
|
|
if_f0 = cpt.get("f0", 1)
|
|
version = cpt.get("version", "v1")
|
|
self.sid = 0
|
|
|
|
if version == "v1":
|
|
if if_f0 == 1:
|
|
from .lib.infer_pack.models import SynthesizerTrnMs256NSFsid
|
|
net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=config.is_half)
|
|
else:
|
|
from .lib.infer_pack.models import SynthesizerTrnMs256NSFsid_nono
|
|
net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"])
|
|
elif version == "v2":
|
|
if if_f0 == 1:
|
|
from .lib.infer_pack.models import SynthesizerTrnMs768NSFsid
|
|
net_g = SynthesizerTrnMs768NSFsid(*cpt["config"], is_half=config.is_half)
|
|
else:
|
|
from .lib.infer_pack.models import SynthesizerTrnMs768NSFsid_nono
|
|
net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])
|
|
del net_g.enc_q
|
|
|
|
net_g.load_state_dict(cpt["weight"], strict=False)
|
|
net_g.eval().to(device if device else config.device)
|
|
if config.is_half:
|
|
net_g = net_g.half()
|
|
else:
|
|
net_g = net_g.float()
|
|
hubert_model = load_hubert(config)
|
|
model_name = os.path.basename(model_path).split(".")[0]
|
|
index_files = get_filenames(root=os.path.join(BASE_MODELS_DIR,"RVC"),folder=".index",exts=["index"],name_filters=[model_name])
|
|
file_index = index_files.pop() if len(index_files) else ""
|
|
|
|
self.cpt = cpt
|
|
self.net_g = net_g
|
|
self.hubert_model = hubert_model
|
|
self.model_name = model_name
|
|
self.index, self.big_npy = self.load_index(file_index)
|
|
self.tgt_sr = tgt_sr
|
|
self.if_f0 = if_f0
|
|
self.version = version
|
|
super().__init__(tgt_sr, config, onnx) # initiate Feature Extraction
|
|
|
|
def __del__(self):
|
|
super().__del__()
|
|
del self.cpt, self.net_g, self.hubert_model, self.index, self.big_npy
|
|
gc_collect()
|
|
|
|
# def process_input(self, x: np.ndarray, **kwargs) -> np.ndarray:
|
|
def vc(self, x: np.ndarray, **kwargs) -> np.ndarray:
|
|
index_rate = kwargs.pop("index_rate",.5)
|
|
protect = kwargs.pop("protect",.5)
|
|
rms_mix_rate = kwargs.pop("rms_mix_rate",1.)
|
|
|
|
feats = torch.from_numpy(x.copy())
|
|
feats = feats.view(1, -1)
|
|
if config.is_half:
|
|
feats = feats.half()
|
|
else:
|
|
feats = feats.float()
|
|
feats = feats.to(self.device)
|
|
|
|
with torch.no_grad():
|
|
padding_mask = torch.BoolTensor(feats.shape).to(self.device).fill_(False)
|
|
inputs = {
|
|
"source": feats,
|
|
"padding_mask": padding_mask,
|
|
"output_layer": 9 if self.version == "v1" else 12,
|
|
}
|
|
logits = self.hubert_model.extract_features(**inputs)
|
|
feats = (
|
|
self.hubert_model.final_proj(logits[0]) if self.version == "v1" else logits[0]
|
|
)
|
|
|
|
if protect < 0.5 and self.if_f0:
|
|
feats0 = F.interpolate(feats.permute(0, 2, 1), scale_factor=2).permute(0, 2, 1)
|
|
|
|
if self.index is not None and self.big_npy is not None and index_rate != 0:
|
|
npy = feats[0].cpu().numpy()
|
|
if self.is_half:
|
|
npy = npy.astype("float16")
|
|
|
|
score, ix = self.index.search(npy, k=8)
|
|
weight = np.square(1 / score)
|
|
weight /= weight.sum(axis=1, keepdims=True)
|
|
npy = np.sum(self.big_npy[ix] * np.expand_dims(weight, axis=2), axis=1)
|
|
|
|
if self.is_half:
|
|
npy = npy.astype("float16")
|
|
feats = (
|
|
torch.from_numpy(npy).unsqueeze(0).to(self.device) * index_rate
|
|
+ (1 - index_rate) * feats
|
|
)
|
|
|
|
feats = F.interpolate(feats.permute(0, 2, 1), scale_factor=2).permute(0, 2, 1)
|
|
|
|
if self.if_f0 == 1:
|
|
pitch, pitchf = self.get_f0(x, **kwargs)
|
|
p_len = min(feats.shape[1], pitch.shape[0])
|
|
pitch = pitch[:p_len]
|
|
pitchf = pitchf[:p_len]
|
|
pitch = torch.from_numpy(pitch).to(self.device).unsqueeze(0)
|
|
pitchf = torch.from_numpy(pitchf).to(self.device).unsqueeze(0)
|
|
if protect < 0.5:
|
|
pitchff = pitchf.clone()
|
|
pitchff[pitchf > 0] = 1
|
|
pitchff[pitchf < 1] = protect
|
|
pitchff = pitchff.unsqueeze(-1)
|
|
feats = feats * pitchff + feats0 * (1 - pitchff)
|
|
del pitchff
|
|
else:
|
|
pitch, pitchf = None, None
|
|
p_len = feats.shape[1]
|
|
|
|
p_len = torch.LongTensor([p_len]).to(self.device)
|
|
sid = torch.LongTensor([self.sid]).to(self.device)
|
|
with torch.no_grad():
|
|
if self.is_half: feats = feats.to(torch.half)
|
|
if self.if_f0 == 1:
|
|
# print("process_output",feats,p_len,pitch,pitchf)
|
|
# print(12222222222,feats.dtype,pitch.dtype,pitchf.dtype,sid.dtype,self.is_half)
|
|
infered_audio = (
|
|
self.net_g.infer(feats, p_len, pitch, pitchf, sid)[0][0, 0].data
|
|
)
|
|
else:
|
|
infered_audio = (
|
|
self.net_g.infer(feats, p_len, sid)[0][0, 0].data
|
|
)
|
|
|
|
audio_opt = infered_audio.cpu().float().numpy()
|
|
if rms_mix_rate < 1.:
|
|
audio_opt = change_rms(x, self.sr, audio_opt, self.tgt_sr, rms_mix_rate)
|
|
|
|
del feats, p_len, sid, pitch, pitchf, infered_audio
|
|
return audio_opt |