initial refactor of rvc studio
This commit is contained in:
SayanoAI
@@ -0,0 +1,274 @@
|
||||
import sys, os, multiprocessing
|
||||
from threading import Thread
|
||||
from scipy import signal
|
||||
import numpy as np, os, traceback
|
||||
from .lib.model_utils import load_hubert
|
||||
from .lib.slicer2 import Slicer
|
||||
import librosa, traceback
|
||||
from scipy.io import wavfile
|
||||
from .lib.audio import load_audio
|
||||
from .pitch_extraction import FeatureExtractor
|
||||
from .lib.audio import load_input_audio
|
||||
from .lib.utils import gc_collect
|
||||
from .config import config
|
||||
import torch
|
||||
|
||||
class Preprocess:
|
||||
def __init__(self, sr, exp_dir, noparallel=True, period=3.0, overlap=.3, max_volume=.95):
|
||||
self.slicer = Slicer(
|
||||
sr=sr,
|
||||
threshold=-42,
|
||||
min_length=1500,
|
||||
min_interval=400,
|
||||
hop_size=15,
|
||||
max_sil_kept=500
|
||||
)
|
||||
self.sr = sr
|
||||
self.bh, self.ah = signal.butter(N=5, Wn=48, btype="high", fs=self.sr)
|
||||
self.per = period
|
||||
self.overlap = overlap
|
||||
self.tail = self.per + self.overlap
|
||||
self.max = max_volume
|
||||
self.alpha = 0.75
|
||||
self.exp_dir = exp_dir
|
||||
self.gt_wavs_dir = "%s/0_gt_wavs" % exp_dir
|
||||
self.wavs16k_dir = "%s/1_16k_wavs" % exp_dir
|
||||
self.noparallel = noparallel
|
||||
os.makedirs(self.exp_dir, exist_ok=True)
|
||||
os.makedirs(self.gt_wavs_dir, exist_ok=True)
|
||||
os.makedirs(self.wavs16k_dir, exist_ok=True)
|
||||
|
||||
def println(self,strr):
|
||||
# mutex.acquire()
|
||||
print(strr)
|
||||
with open("%s/preprocess.log" % self.exp_dir, "a+") as f:
|
||||
f.write("%s\n" % strr)
|
||||
f.flush()
|
||||
# mutex.release()
|
||||
|
||||
def norm_write(self, tmp_audio, idx0, idx1):
|
||||
tmp_max = np.abs(tmp_audio).max()
|
||||
if tmp_max > 2.5:
|
||||
print("%s-%s-%s-filtered" % (idx0, idx1, tmp_max))
|
||||
return
|
||||
tmp_audio = (tmp_audio / tmp_max * (self.max * self.alpha)) + (
|
||||
1 - self.alpha
|
||||
) * tmp_audio
|
||||
wavfile.write(
|
||||
"%s/%s_%s.wav" % (self.gt_wavs_dir, idx0, idx1),
|
||||
self.sr,
|
||||
tmp_audio.astype(np.float32),
|
||||
)
|
||||
tmp_audio = librosa.resample(
|
||||
tmp_audio, orig_sr=self.sr, target_sr=16000
|
||||
) # , res_type="soxr_vhq"
|
||||
wavfile.write(
|
||||
"%s/%s_%s.wav" % (self.wavs16k_dir, idx0, idx1),
|
||||
16000,
|
||||
tmp_audio.astype(np.float32),
|
||||
)
|
||||
|
||||
def pipeline(self, path, idx0):
|
||||
try:
|
||||
audio = load_audio(path, self.sr)
|
||||
# zero phased digital filter cause pre-ringing noise...
|
||||
# audio = signal.filtfilt(self.bh, self.ah, audio)
|
||||
# audio = signal.lfilter(self.bh, self.ah, audio)
|
||||
|
||||
idx1 = 0
|
||||
for audio in self.slicer.slice(audio):
|
||||
i = 0
|
||||
while 1:
|
||||
start = int(self.sr * (self.per - self.overlap) * i)
|
||||
i += 1
|
||||
if len(audio[start:]) > self.tail * self.sr:
|
||||
tmp_audio = audio[start : start + int(self.per * self.sr)]
|
||||
self.norm_write(tmp_audio, idx0, idx1)
|
||||
idx1 += 1
|
||||
else:
|
||||
tmp_audio = audio[start:]
|
||||
idx1 += 1
|
||||
break
|
||||
self.norm_write(tmp_audio, idx0, idx1)
|
||||
self.println("%s->Suc." % path)
|
||||
except:
|
||||
self.println("%s->%s" % (path, traceback.format_exc()))
|
||||
|
||||
def pipeline_mp(self, infos):
|
||||
for path, idx0 in infos:
|
||||
self.pipeline(path, idx0)
|
||||
|
||||
def pipeline_mp_inp_dir(self, inp_root, n_p):
|
||||
try:
|
||||
infos = [
|
||||
("%s/%s" % (inp_root, name), idx)
|
||||
for idx, name in enumerate(sorted(list(os.listdir(inp_root))))
|
||||
]
|
||||
if self.noparallel:
|
||||
for i in range(n_p):
|
||||
self.pipeline_mp(infos[i::n_p])
|
||||
else:
|
||||
ps = []
|
||||
for i in range(n_p):
|
||||
p = multiprocessing.Process(
|
||||
target=self.pipeline_mp, args=(infos[i::n_p],)
|
||||
)
|
||||
ps.append(p)
|
||||
p.start()
|
||||
for i in range(n_p):
|
||||
ps[i].join()
|
||||
except:
|
||||
self.println("Fail. %s" % traceback.format_exc())
|
||||
|
||||
class FeatureInput(FeatureExtractor):
|
||||
def __init__(self, f0_method, exp_dir, samplerate=16000, hop_size=160, device="cpu", version="v2", if_f0=False):
|
||||
self.sr = samplerate
|
||||
self.hop = hop_size
|
||||
self.f0_method = f0_method
|
||||
self.exp_dir = exp_dir
|
||||
self.device = device
|
||||
self.version = version
|
||||
self.if_f0 = if_f0
|
||||
|
||||
self.f0_bin = 256
|
||||
self.f0_max = 1100.0
|
||||
self.f0_min = 50.0
|
||||
self.f0_mel_min = 1127 * np.log(1 + self.f0_min / 700)
|
||||
self.f0_mel_max = 1127 * np.log(1 + self.f0_max / 700)
|
||||
|
||||
self.model = load_hubert(config)
|
||||
|
||||
super().__init__(samplerate, config, onnx=False)
|
||||
|
||||
def printt(self,strr):
|
||||
print(strr)
|
||||
with open("%s/extract_f0_feature.log" % self.exp_dir, "a+") as f:
|
||||
f.write("%s\n" % strr)
|
||||
f.flush()
|
||||
|
||||
def compute_feats(self,x):
|
||||
feats = torch.from_numpy(x).float()
|
||||
if feats.dim() == 2: # double channels
|
||||
feats = feats.mean(-1)
|
||||
assert feats.dim() == 1, feats.dim()
|
||||
feats = feats.view(1, -1)
|
||||
padding_mask = torch.BoolTensor(feats.shape).fill_(False)
|
||||
|
||||
inputs = {
|
||||
"source": feats.half().to(self.device)
|
||||
if self.device not in ["mps", "cpu"]
|
||||
else feats.to(self.device),
|
||||
"padding_mask": padding_mask.to(self.device),
|
||||
"output_layer": 9 if self.version == "v1" else 12, # layer 9
|
||||
}
|
||||
with torch.no_grad():
|
||||
logits = self.model.extract_features(**inputs)
|
||||
feats = (
|
||||
self.model.final_proj(logits[0]) if self.version == "v1" else logits[0]
|
||||
)
|
||||
|
||||
feats = feats.squeeze(0).float().cpu().numpy()
|
||||
if np.isnan(feats).sum() == 0:
|
||||
return feats
|
||||
else:
|
||||
return self.printt("==contains nan==")
|
||||
|
||||
def compute_f0(self,x):
|
||||
return self.get_f0(x,0,self.f0_method,crepe_hop_length=self.hop)
|
||||
|
||||
def go(self, paths):
|
||||
if len(paths) == 0:
|
||||
self.printt("no-f0-todo")
|
||||
else:
|
||||
self.printt("todo-f0-%s" % len(paths))
|
||||
# n = max(len(paths) // 5, 1) # 每个进程最多打印5条
|
||||
for idx, (inp_path, opt_path1, opt_path2, opt_path3) in enumerate(paths):
|
||||
try:
|
||||
# if idx % n == 0:
|
||||
# self.printt("f0ing,now-%s,all-%s,-%s" % (idx, len(paths), inp_path))
|
||||
if (
|
||||
os.path.exists(opt_path1 + ".npy") == True
|
||||
and os.path.exists(opt_path2 + ".npy") == True
|
||||
and os.path.exists(opt_path3 + ".npy") == True
|
||||
):
|
||||
continue
|
||||
x,_ = load_input_audio(inp_path,self.sr)
|
||||
if self.model:
|
||||
feats = self.compute_feats(x)
|
||||
if feats is not None:
|
||||
np.save(
|
||||
opt_path3,
|
||||
feats,
|
||||
allow_pickle=False,
|
||||
) # features
|
||||
if self.if_f0: # uses pitch
|
||||
coarse_pit, featur_pit = self.compute_f0(x)
|
||||
np.save(
|
||||
opt_path2,
|
||||
featur_pit,
|
||||
allow_pickle=False,
|
||||
) # nsf
|
||||
np.save(
|
||||
opt_path1,
|
||||
coarse_pit,
|
||||
allow_pickle=False,
|
||||
) # ori
|
||||
except:
|
||||
self.printt("f0fail-%s-%s-%s" % (idx, inp_path, traceback.format_exc()))
|
||||
|
||||
def preprocess_trainset(inp_root, sr, n_p, exp_dir, period=3.0, overlap=.3):
|
||||
try:
|
||||
pp = Preprocess(sr, exp_dir, period=period, overlap=overlap)
|
||||
pp.println("start preprocess")
|
||||
pp.println(sys.argv)
|
||||
pp.pipeline_mp_inp_dir(inp_root, n_p)
|
||||
pp.println("end preprocess")
|
||||
del pp
|
||||
gc_collect()
|
||||
return "Successfully preprocessed data"
|
||||
except Exception as e:
|
||||
return f"Failed to preprocess data: {e}"
|
||||
|
||||
def extract_features_trainset(exp_dir,n_p,f0method,device,version,if_f0):
|
||||
try:
|
||||
featureInput = FeatureInput(f0_method=f0method,exp_dir=exp_dir,device=device,version=version,if_f0=if_f0)
|
||||
paths = []
|
||||
inp_root = os.path.join(exp_dir,"1_16k_wavs")
|
||||
opt_root1 = os.path.join(exp_dir,"2a_f0")
|
||||
opt_root2 = os.path.join(exp_dir,"2b-f0nsf")
|
||||
opt_root3 = os.path.join(exp_dir,"3_feature256" if version == "v1" else "3_feature768")
|
||||
|
||||
os.makedirs(opt_root1, exist_ok=True)
|
||||
os.makedirs(opt_root2, exist_ok=True)
|
||||
os.makedirs(opt_root3, exist_ok=True)
|
||||
|
||||
for name in sorted(list(os.listdir(inp_root))):
|
||||
inp_path = os.path.join(inp_root, name)
|
||||
if "spec" in inp_path:
|
||||
continue
|
||||
|
||||
opt_path1 = os.path.join(opt_root1, ",".join([str(f0method),name]))
|
||||
opt_path2 = os.path.join(opt_root2, ",".join([str(f0method),name]))
|
||||
opt_path3 = os.path.join(opt_root3, ",".join([str(f0method),name]))
|
||||
paths.append([inp_path, opt_path1, opt_path2, opt_path3])
|
||||
|
||||
ps = []
|
||||
n_p = max(n_p,1)
|
||||
for i in range(n_p):
|
||||
if device=="cuda":
|
||||
featureInput.go(paths[i::n_p])
|
||||
else:
|
||||
p = Thread(target=featureInput.go,args=(paths[i::n_p],),daemon=True)
|
||||
ps.append(p)
|
||||
p.start()
|
||||
|
||||
if device != "cuda":
|
||||
for p in ps:
|
||||
try:
|
||||
p.join()
|
||||
except:
|
||||
featureInput.printt("f0_all_fail-%s" % (traceback.format_exc()))
|
||||
|
||||
return f"Successfully extracted features using {f0method}"
|
||||
except Exception as e:
|
||||
return f"Failed to extract features: {e}"
|
||||
Reference in New Issue
Block a user