kwau/sovits-isla

1import json2 3import torch4 5import utils6from onnxexport.model_onnx_speaker_mix import SynthesizerTrn7 8 9def main(path, config, model):10    #path = "crs"11 12    device = torch.device("cpu")13    hps = utils.get_hparams_from_file(f"checkpoints/{path}/{config}")14    SVCVITS = SynthesizerTrn(15        hps.data.filter_length // 2 + 1,16        hps.train.segment_size // hps.data.hop_length,17        **hps.model)18    _ = utils.load_checkpoint(f"checkpoints/{path}/{model}", SVCVITS, None)19    _ = SVCVITS.eval().to(device)20    for i in SVCVITS.parameters():21        i.requires_grad = False22    23    num_frames = 20024 25    test_hidden_unit = torch.rand(1, num_frames, SVCVITS.gin_channels)26    test_pitch = torch.rand(1, num_frames)27    test_vol = torch.rand(1, num_frames)28    test_mel2ph = torch.LongTensor(torch.arange(0, num_frames)).unsqueeze(0)29    test_uv = torch.ones(1, num_frames, dtype=torch.float32)30    test_noise = torch.randn(1, 192, num_frames)31    test_sid = torch.LongTensor([0])32    export_mix = True33    if len(hps.spk) < 2:34        export_mix = False35    36    if export_mix:37        spk_mix = []38        n_spk = len(hps.spk)39        for i in range(n_spk):40            spk_mix.append(1.0/float(n_spk))41        test_sid = torch.tensor(spk_mix)42        SVCVITS.export_chara_mix(hps.spk)43        test_sid = test_sid.unsqueeze(0)44        test_sid = test_sid.repeat(num_frames, 1)45    46    SVCVITS.eval()47 48    if export_mix:49        daxes = {50            "c": [0, 1],51            "f0": [1],52            "mel2ph": [1],53            "uv": [1],54            "noise": [2],55            "sid":[0]56        }57    else:58        daxes = {59            "c": [0, 1],60            "f0": [1],61            "mel2ph": [1],62            "uv": [1],63            "noise": [2]64        }65    66    input_names = ["c", "f0", "mel2ph", "uv", "noise", "sid"]67    output_names = ["audio", ]68 69    if SVCVITS.vol_embedding:70        input_names.append("vol")71        vol_dadict = {"vol" : [1]}72        daxes.update(vol_dadict)73        test_inputs = (74            test_hidden_unit.to(device),75            test_pitch.to(device),76            test_mel2ph.to(device),77            test_uv.to(device),78            test_noise.to(device),79            test_sid.to(device),80            test_vol.to(device)81        )82    else:83        test_inputs = (84            test_hidden_unit.to(device),85            test_pitch.to(device),86            test_mel2ph.to(device),87            test_uv.to(device),88            test_noise.to(device),89            test_sid.to(device)90        )91 92    # SVCVITS = torch.jit.script(SVCVITS)93    SVCVITS(test_hidden_unit.to(device),94            test_pitch.to(device),95            test_mel2ph.to(device),96            test_uv.to(device),97            test_noise.to(device),98            test_sid.to(device),99            test_vol.to(device))100 101    SVCVITS.dec.OnnxExport()102 103    torch.onnx.export(104        SVCVITS,105        test_inputs,106        f"checkpoints/{path}/{path}_SoVits.onnx",107        dynamic_axes=daxes,108        do_constant_folding=False,109        opset_version=16,110        verbose=False,111        input_names=input_names,112        output_names=output_names113    )114 115    vec_lay = "layer-12" if SVCVITS.gin_channels == 768 else "layer-9"116    spklist = []117    for key in hps.spk.keys():118        spklist.append(key)119 120    MoeVSConf = {121        "Folder" : f"{path}",122        "Name" : f"{path}",123        "Type" : "SoVits",124        "Rate" : hps.data.sampling_rate,125        "Hop" : hps.data.hop_length,126        "Hubert": f"vec-{SVCVITS.gin_channels}-{vec_lay}",127        "SoVits4": True,128        "SoVits3": False,129        "CharaMix": export_mix,130        "Volume": SVCVITS.vol_embedding,131        "HiddenSize": SVCVITS.gin_channels,132        "Characters": spklist133    }134 135    with open(f"checkpoints/{path}/{model}_MoeVS.json", 'w') as MoeVsConfFile:136        json.dump(MoeVSConf, MoeVsConfFile, indent = 4)137 138 139if __name__ == '__main__':140    import argparse141    parser = argparse.ArgumentParser()142    parser.add_argument('-p', '--path', type=str, default="crs")143    parser.add_argument('-c', '--config', type=str, default='config.json')144    parser.add_argument('-m', '--model', type=str, default='model.pth')145    args = parser.parse_args()146    147    path = args.path148    config = args.config149    model = args.model150    main(path, config, model)151