preprocessing_utils.py

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 lib 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}"