ChazzyG/Retrieval-based-Voice-Conversion-WebUI

1import os, sys, torch, warnings, pdb2 3warnings.filterwarnings("ignore")4import librosa5import importlib6import numpy as np7import hashlib, math8from tqdm import tqdm9from uvr5_pack.lib_v5 import spec_utils10from uvr5_pack.utils import _get_name_params, inference11from uvr5_pack.lib_v5.model_param_init import ModelParameters12from scipy.io import wavfile13 14 15class _audio_pre_:16    def __init__(self, agg, model_path, device, is_half):17        self.model_path = model_path18        self.device = device19        self.data = {20            # Processing Options21            "postprocess": False,22            "tta": False,23            # Constants24            "window_size": 512,25            "agg": agg,26            "high_end_process": "mirroring",27        }28        nn_arch_sizes = [29            31191,  # default30            33966,31            61968,32            123821,33            123812,34            537238,  # custom35        ]36        self.nn_architecture = list("{}KB".format(s) for s in nn_arch_sizes)37        model_size = math.ceil(os.stat(model_path).st_size / 1024)38        nn_architecture = "{}KB".format(39            min(nn_arch_sizes, key=lambda x: abs(x - model_size))40        )41        nets = importlib.import_module(42            "uvr5_pack.lib_v5.nets"43            + f"_{nn_architecture}".replace("_{}KB".format(nn_arch_sizes[0]), ""),44            package=None,45        )46        model_hash = hashlib.md5(open(model_path, "rb").read()).hexdigest()47        param_name, model_params_d = _get_name_params(model_path, model_hash)48 49        mp = ModelParameters(model_params_d)50        model = nets.CascadedASPPNet(mp.param["bins"] * 2)51        cpk = torch.load(model_path, map_location="cpu")52        model.load_state_dict(cpk)53        model.eval()54        if is_half:55            model = model.half().to(device)56        else:57            model = model.to(device)58 59        self.mp = mp60        self.model = model61 62    def _path_audio_(self, music_file, ins_root=None, vocal_root=None):63        if ins_root is None and vocal_root is None:64            return "No save root."65        name = os.path.basename(music_file)66        if ins_root is not None:67            os.makedirs(ins_root, exist_ok=True)68        if vocal_root is not None:69            os.makedirs(vocal_root, exist_ok=True)70        X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}71        bands_n = len(self.mp.param["band"])72        # print(bands_n)73        for d in range(bands_n, 0, -1):74            bp = self.mp.param["band"][d]75            if d == bands_n:  # high-end band76                (77                    X_wave[d],78                    _,79                ) = librosa.core.load(  # 理论上librosa读取可能对某些音频有bug，应该上ffmpeg读取，但是太麻烦了弃坑80                    music_file,81                    bp["sr"],82                    False,83                    dtype=np.float32,84                    res_type=bp["res_type"],85                )86                if X_wave[d].ndim == 1:87                    X_wave[d] = np.asfortranarray([X_wave[d], X_wave[d]])88            else:  # lower bands89                X_wave[d] = librosa.core.resample(90                    X_wave[d + 1],91                    self.mp.param["band"][d + 1]["sr"],92                    bp["sr"],93                    res_type=bp["res_type"],94                )95            # Stft of wave source96            X_spec_s[d] = spec_utils.wave_to_spectrogram_mt(97                X_wave[d],98                bp["hl"],99                bp["n_fft"],100                self.mp.param["mid_side"],101                self.mp.param["mid_side_b2"],102                self.mp.param["reverse"],103            )104            # pdb.set_trace()105            if d == bands_n and self.data["high_end_process"] != "none":106                input_high_end_h = (bp["n_fft"] // 2 - bp["crop_stop"]) + (107                    self.mp.param["pre_filter_stop"] - self.mp.param["pre_filter_start"]108                )109                input_high_end = X_spec_s[d][110                    :, bp["n_fft"] // 2 - input_high_end_h : bp["n_fft"] // 2, :111                ]112 113        X_spec_m = spec_utils.combine_spectrograms(X_spec_s, self.mp)114        aggresive_set = float(self.data["agg"] / 100)115        aggressiveness = {116            "value": aggresive_set,117            "split_bin": self.mp.param["band"][1]["crop_stop"],118        }119        with torch.no_grad():120            pred, X_mag, X_phase = inference(121                X_spec_m, self.device, self.model, aggressiveness, self.data122            )123        # Postprocess124        if self.data["postprocess"]:125            pred_inv = np.clip(X_mag - pred, 0, np.inf)126            pred = spec_utils.mask_silence(pred, pred_inv)127        y_spec_m = pred * X_phase128        v_spec_m = X_spec_m - y_spec_m129 130        if ins_root is not None:131            if self.data["high_end_process"].startswith("mirroring"):132                input_high_end_ = spec_utils.mirroring(133                    self.data["high_end_process"], y_spec_m, input_high_end, self.mp134                )135                wav_instrument = spec_utils.cmb_spectrogram_to_wave(136                    y_spec_m, self.mp, input_high_end_h, input_high_end_137                )138            else:139                wav_instrument = spec_utils.cmb_spectrogram_to_wave(y_spec_m, self.mp)140            print("%s instruments done" % name)141            wavfile.write(142                os.path.join(143                    ins_root, "instrument_{}_{}.wav".format(name, self.data["agg"])144                ),145                self.mp.param["sr"],146                (np.array(wav_instrument) * 32768).astype("int16"),147            )  #148        if vocal_root is not None:149            if self.data["high_end_process"].startswith("mirroring"):150                input_high_end_ = spec_utils.mirroring(151                    self.data["high_end_process"], v_spec_m, input_high_end, self.mp152                )153                wav_vocals = spec_utils.cmb_spectrogram_to_wave(154                    v_spec_m, self.mp, input_high_end_h, input_high_end_155                )156            else:157                wav_vocals = spec_utils.cmb_spectrogram_to_wave(v_spec_m, self.mp)158            print("%s vocals done" % name)159            wavfile.write(160                os.path.join(161                    vocal_root, "vocal_{}_{}.wav".format(name, self.data["agg"])162                ),163                self.mp.param["sr"],164                (np.array(wav_vocals) * 32768).astype("int16"),165            )166 167 168if __name__ == "__main__":169    device = "cuda"170    is_half = True171    model_path = "uvr5_weights/2_HP-UVR.pth"172    pre_fun = _audio_pre_(model_path=model_path, device=device, is_half=True)173    audio_path = "神女劈观.aac"174    save_path = "opt"175    pre_fun._path_audio_(audio_path, save_path, save_path)176