275 lines
10 KiB
Python
275 lines
10 KiB
Python
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}"
|