import sys, os, multiprocessing from threading import Thread import numpy as np, os, traceback from .lib.slicer2 import Slicer import traceback from scipy.io import wavfile from .pitch_extraction import FeatureExtractor from .lib.audio import hz_to_mel, load_input_audio, remix_audio, AudioProcessor from .lib.helper import gc_collect from .config import config import torch class Preprocess: def __init__(self, sr, exp_dir, preprocessor: "AudioProcessor"=None, noparallel=True, period=3.0, overlap=.3, max_volume=.95): self.slicer = Slicer( sr=sr, threshold=-50, min_length=1500, min_interval=400, hop_size=15, max_sil_kept=500 ) self.sr = sr self.per = period self.overlap = overlap self.tail = self.per + self.overlap self.max_volume = max_volume self.exp_dir = exp_dir self.gt_wavs_dir = os.path.join(exp_dir,"0_gt_wavs") self.wavs16k_dir = os.path.join(exp_dir,"1_16k_wavs") self.noparallel = noparallel self.preprocessor = preprocessor 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): if len(tmp_audio) > self.overlap*self.sr*2: wavfile.write(os.path.join(self.gt_wavs_dir, f"{idx0}_{idx1}.wav"),self.sr,tmp_audio.astype(np.float32)) remixed_audio = remix_audio((tmp_audio, self.sr), target_sr=16000, max_volume=self.max_volume) wavfile.write(os.path.join(self.wavs16k_dir, f"{idx0}_{idx1}.wav"),16000,remixed_audio[0].astype(np.float32)) else: print(f"skipped short audio clip: {idx0}_{idx1}.wav ({len(tmp_audio)=})") def pipeline(self, path, idx0): try: input_audio = load_input_audio(path, self.sr) if self.preprocessor is not None: input_audio = self.preprocessor(input_audio) idx1 = 0 for audio in self.slicer.slice(input_audio[0]): 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, model, 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 = hz_to_mel(self.f0_min) self.f0_mel_max = hz_to_mel(self.f0_max) self.model = model 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 } feats = self.model.extract_features(version=self.version,**inputs) 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, preprocessor=None, period=3.0, overlap=.3, max_volume=1.): try: pp = Preprocess(sr, exp_dir, preprocessor=preprocessor, period=period, overlap=overlap, max_volume=max_volume) 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() print("Successfully preprocessed data") return True except Exception as e: print(f"Failed to preprocess data: {e}") return False def extract_features_trainset(hubert_model,exp_dir,n_p,f0method,device,version,if_f0,crepe_hop_length): try: featureInput = FeatureInput(f0_method=f0method,exp_dir=exp_dir,device=device,version=version,if_f0=if_f0,model=hubert_model,hop_size=crepe_hop_length) 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())) print(f"Successfully extracted features using {f0method}") return True except Exception as e: print(f"Failed to extract features: {e}") return False