File size: 2,713 Bytes
b51190f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | #!/usr/bin/env python3
"""Export Piper checkpoint using JIT tracing"""
import sys
import torch
from pathlib import Path
sys.path.insert(0, str(Path("/root/piper_msa/piper1-gpl/src")))
from piper.train.vits.lightning import VitsModel
def main():
checkpoint_path = "/root/piper_msa/piper1-gpl/lightning_logs/version_1/checkpoints/epoch=455-step=1189248.ckpt"
output_path = "/root/piper_msa/output/saudi_msa_epoch455.onnx"
print(f"Loading checkpoint: {checkpoint_path}")
model = VitsModel.load_from_checkpoint(checkpoint_path, map_location="cpu")
model_g = model.model_g
# Inference only
model_g.eval()
with torch.no_grad():
model_g.dec.remove_weight_norm()
def infer_forward(text, text_lengths, scales, sid=None):
noise_scale = scales[0]
length_scale = scales[1]
noise_scale_w = scales[2]
audio = model_g.infer(
text,
text_lengths,
noise_scale=noise_scale,
length_scale=length_scale,
noise_scale_w=noise_scale_w,
sid=sid,
)[0].unsqueeze(1)
return audio
model_g.forward = infer_forward
num_symbols = model_g.n_vocab
num_speakers = model_g.n_speakers
dummy_input_length = 50
sequences = torch.randint(
low=0, high=num_symbols, size=(1, dummy_input_length), dtype=torch.long
)
sequence_lengths = torch.LongTensor([sequences.size(1)])
sid = None
if num_speakers > 1:
sid = torch.LongTensor([0])
scales = torch.FloatTensor([0.667, 1.0, 0.8])
dummy_input = (sequences, sequence_lengths, scales, sid)
print(f"Exporting to ONNX using JIT: {output_path}")
# Use JIT tracing with legacy exporter
with torch.no_grad():
torch.onnx.export(
model=model_g,
args=dummy_input,
f=output_path,
verbose=False,
opset_version=15,
input_names=["input", "input_lengths", "scales", "sid"],
output_names=["output"],
dynamic_axes={
"input": {0: "batch_size", 1: "phonemes"},
"input_lengths": {0: "batch_size"},
"output": {0: "batch_size", 2: "time"},
},
export_params=True,
do_constant_folding=True,
# Use legacy JIT-based exporter
dynamo=False,
)
print(f"✓ Model exported successfully to: {output_path}")
print(f"\nTo test the model:")
print(f" echo 'مرحبا بك' | piper --model {output_path} --config /root/piper_msa/output/config.json --output_file test.wav")
if __name__ == "__main__":
main()
|