Codecooker/rvcapi

1from multiprocessing import cpu_count2from pathlib import Path3 4import torch5from fairseq import checkpoint_utils6from scipy.io import wavfile7 8from infer_pack.models import (9    SynthesizerTrnMs256NSFsid,10    SynthesizerTrnMs256NSFsid_nono,11    SynthesizerTrnMs768NSFsid,12    SynthesizerTrnMs768NSFsid_nono,13)14from my_utils import load_audio15from vc_infer_pipeline import VC16 17BASE_DIR = Path(__file__).resolve().parent.parent18 19 20class Config:21    def __init__(self, device, is_half):22        self.device = device23        self.is_half = is_half24        self.n_cpu = 025        self.gpu_name = None26        self.gpu_mem = None27        self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()28 29    def device_config(self) -> tuple:30        if torch.cuda.is_available():31            i_device = int(self.device.split(":")[-1])32            self.gpu_name = torch.cuda.get_device_name(i_device)33            if (34                    ("16" in self.gpu_name and "V100" not in self.gpu_name.upper())35                    or "P40" in self.gpu_name.upper()36                    or "1060" in self.gpu_name37                    or "1070" in self.gpu_name38                    or "1080" in self.gpu_name39            ):40                print("16 series/10 series P40 forced single precision")41                self.is_half = False42                for config_file in ["32k.json", "40k.json", "48k.json"]:43                    with open(BASE_DIR / "src" / "configs" / config_file, "r") as f:44                        strr = f.read().replace("true", "false")45                    with open(BASE_DIR / "src" / "configs" / config_file, "w") as f:46                        f.write(strr)47                with open(BASE_DIR / "src" / "trainset_preprocess_pipeline_print.py", "r") as f:48                    strr = f.read().replace("3.7", "3.0")49                with open(BASE_DIR / "src" / "trainset_preprocess_pipeline_print.py", "w") as f:50                    f.write(strr)51            else:52                self.gpu_name = None53            self.gpu_mem = int(54                torch.cuda.get_device_properties(i_device).total_memory55                / 102456                / 102457                / 102458                + 0.459            )60            if self.gpu_mem <= 4:61                with open(BASE_DIR / "src" / "trainset_preprocess_pipeline_print.py", "r") as f:62                    strr = f.read().replace("3.7", "3.0")63                with open(BASE_DIR / "src" / "trainset_preprocess_pipeline_print.py", "w") as f:64                    f.write(strr)65        elif torch.backends.mps.is_available():66            print("No supported N-card found, use MPS for inference")67            self.device = "mps"68        else:69            print("No supported N-card found, use CPU for inference")70            self.device = "cpu"71            self.is_half = True72 73        if self.n_cpu == 0:74            self.n_cpu = cpu_count()75 76        if self.is_half:77            # 6G memory config78            x_pad = 379            x_query = 1080            x_center = 6081            x_max = 6582        else:83            # 5G memory config84            x_pad = 185            x_query = 686            x_center = 3887            x_max = 4188 89        if self.gpu_mem != None and self.gpu_mem <= 4:90            x_pad = 191            x_query = 592            x_center = 3093            x_max = 3294 95        return x_pad, x_query, x_center, x_max96 97 98def load_hubert(device, is_half, model_path):99    models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task([model_path], suffix='', )100    hubert = models[0]101    hubert = hubert.to(device)102 103    if is_half:104        hubert = hubert.half()105    else:106        hubert = hubert.float()107 108    hubert.eval()109    return hubert110 111 112def get_vc(device, is_half, config, model_path):113    cpt = torch.load(model_path, map_location='cpu')114    tgt_sr = cpt["config"][-1]115    cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]116    if_f0 = cpt.get("f0", 1)117    version = cpt.get("version", "v1")118 119    if version == "v1":120        if if_f0 == 1:121            net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=is_half)122        else:123            net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"])124    elif version == "v2":125        if if_f0 == 1:126            net_g = SynthesizerTrnMs768NSFsid(*cpt["config"], is_half=is_half)127        else:128            net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])129 130    del net_g.enc_q131    print(net_g.load_state_dict(cpt["weight"], strict=False))132    net_g.eval().to(device)133 134    if is_half:135        net_g = net_g.half()136    else:137        net_g = net_g.float()138 139    vc = VC(tgt_sr, config)140    return cpt, version, net_g, tgt_sr, vc141 142 143def rvc_infer(index_path, index_rate, input_path, output_path, pitch_change, cpt, version, net_g, filter_radius, tgt_sr, rms_mix_rate, protect, vc, hubert_model):144    audio = load_audio(input_path, 16000)145    times = [0, 0, 0]146    if_f0 = cpt.get('f0', 1)147    audio_opt = vc.pipeline(hubert_model, net_g, 0, audio, input_path, times, pitch_change, 'rmvpe', index_path, index_rate, if_f0, filter_radius, tgt_sr, 0, rms_mix_rate, version, protect, None)148    wavfile.write(output_path, tgt_sr, audio_opt)149