Bindkushal/IndicVoice-82M

1import argparse2import os3import torch4import onnx5import onnxruntime as ort6import sounddevice as sd7 8from indicvoice import IndicModel, IndicPipeline9from indicvoice.model import IndicModelForONNX10 11def export_onnx(model, output):12    onnx_file = output + "/" + "indicvoice.onnx"13 14    input_ids = torch.randint(1, 100, (48,)).numpy()15    input_ids = torch.LongTensor([[0, *input_ids, 0]])16    style = torch.randn(1, 256)17    speed = torch.randint(1, 10, (1,)).int()18 19    torch.onnx.export(20        model, 21        args = (input_ids, style, speed), 22        f = onnx_file, 23        export_params = True, 24        verbose = True, 25        input_names = [ 'input_ids', 'style', 'speed' ], 26        output_names = [ 'waveform', 'duration' ],27        opset_version = 17, 28        dynamic_axes = {29            'input_ids': {0: "batch_size", 1: 'input_ids_len' }, 30            'style': {0: "batch_size"}, 31            "speed": {0: "batch_size"}32        }, 33        do_constant_folding = True, 34    )35 36    print('export indicvoice.onnx ok!')37 38    onnx_model = onnx.load(onnx_file)39    onnx.checker.check_model(onnx_model)40    print('onnx check ok!')41 42def load_input_ids(pipeline, text):43    if pipeline.lang_code in 'ab':44        _, tokens = pipeline.g2p(text)45        for gs, ps, tks in pipeline.en_tokenize(tokens):46            if not ps:47                continue48    else:49        ps, _ = pipeline.g2p(text)50 51    if len(ps) > 510:52        ps = ps[:510]53 54    input_ids = list(filter(lambda i: i is not None, map(lambda p: pipeline.model.vocab.get(p), ps)))55    print(f"text: {text} -> phonemes: {ps} -> input_ids: {input_ids}")56    input_ids = torch.LongTensor([[0, *input_ids, 0]]).to(pipeline.model.device)57    return ps, input_ids58 59def load_voice(pipeline, voice, phonemes):60    pack = pipeline.load_voice(voice).to('cpu')61    return pack[len(phonemes) - 1]62 63def load_sample(model):64    pipeline = IndicPipeline(lang_code='a', model=model.kmodel, device='cpu')65    text = '''66    In today's fast-paced tech world, building software applications has never been easier — thanks to AI-powered coding assistants.'67    '''68    text = '''69    The sky above the port was the color of television, tuned to a dead channel.70    '''71    voice = 'checkpoints/voices/af_heart.pt'72 73    pipeline = IndicPipeline(lang_code='z', model=model.kmodel, device='cpu')74    text = '''75    2月15日晚，猫眼专业版数据显示，截至发稿，《哪吒之魔童闹海》（或称《哪吒2》）今日票房已达7.8亿元，累计票房（含预售）超过114亿元。76    '''77    voice = 'checkpoints/voices/zf_xiaoxiao.pt'78 79    phonemes, input_ids = load_input_ids(pipeline, text)80    style = load_voice(pipeline, voice, phonemes)81    speed = torch.IntTensor([1])82 83    return input_ids, style, speed84 85def inference_onnx(model, output):86    onnx_file = output + "/" + "indicvoice.onnx"87    session = ort.InferenceSession(onnx_file)88 89    input_ids, style, speed = load_sample(model)90 91    outputs = session.run(None, {92        'input_ids': input_ids.numpy(), 93        'style': style.numpy(), 94        'speed': speed.numpy(), 95    })96 97    output = torch.from_numpy(outputs[0])98    print(f'output: {output.shape}')99    print(output)100 101    audio = output.numpy()102    sd.play(audio, 24000)103    sd.wait()104 105def check_model(model):106    input_ids, style, speed = load_sample(model)107    output, duration = model(input_ids, style, speed)108 109    print(f'output: {output.shape}')110    print(f'duration: {duration.shape}')111    print(output)112 113    audio = output.numpy()114    sd.play(audio, 24000)115    sd.wait()116 117if __name__ == "__main__":118    parser = argparse.ArgumentParser("Export IndicVoice Model to ONNX", add_help=True)119    parser.add_argument("--inference", "-t", help="test indicvoice.onnx model", action="store_true")120    parser.add_argument("--check", "-m", help="check indicvoice model", action="store_true")121    parser.add_argument(122        "--config_file", "-c", type=str, default="checkpoints/config.json", help="path to config file"123    )124    parser.add_argument(125        "--checkpoint_path", "-p", type=str, default="checkpoints/indicvoice-v1_0.pth", help="path to checkpoint file"126    )127    parser.add_argument(128        "--output_dir", "-o", type=str, default="onnx", help="output directory"129    )130 131    args = parser.parse_args()132 133    # cfg134    config_file = args.config_file  # change the path of the model config file135    checkpoint_path = args.checkpoint_path  # change the path of the model136    output_dir = args.output_dir137    138    # make dir139    os.makedirs(output_dir, exist_ok=True)140 141    kmodel = IndicModel(config=config_file, model=checkpoint_path, disable_complex=True)142    model = IndicModelForONNX(kmodel).eval()143 144    if args.inference:145        inference_onnx(model, output_dir)146    elif args.check:147        check_model(model)148    else:149        export_onnx(model, output_dir)150