You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
piper/src/python/piper_train/infer_generator.py

84 lines
2.3 KiB
Python

#!/usr/bin/env python3
import argparse
import json
import logging
import sys
import time
from pathlib import Path
import torch
from .vits.utils import audio_float_to_int16
from .vits.wavfile import write as write_wav
_LOGGER = logging.getLogger("piper_train.infer_generator")
def main():
"""Main entry point"""
logging.basicConfig(level=logging.DEBUG)
parser = argparse.ArgumentParser(prog="piper_train.infer_generator")
parser.add_argument("--model", required=True, help="Path to generator (.pt)")
parser.add_argument("--output-dir", required=True, help="Path to write WAV files")
parser.add_argument("--sample-rate", type=int, default=22050)
args = parser.parse_args()
args.output_dir = Path(args.output_dir)
args.output_dir.mkdir(parents=True, exist_ok=True)
model = torch.load(args.model)
# Inference only
model.eval()
for i, line in enumerate(sys.stdin):
line = line.strip()
if not line:
continue
utt = json.loads(line)
utt_id = str(i)
phoneme_ids = utt["phoneme_ids"]
speaker_id = utt.get("speaker_id")
text = torch.LongTensor(phoneme_ids).unsqueeze(0)
text_lengths = torch.LongTensor([len(phoneme_ids)])
sid = torch.LongTensor([speaker_id]) if speaker_id is not None else None
start_time = time.perf_counter()
audio = (
model(
text,
text_lengths,
sid,
# torch.FloatTensor([0.667]),
# torch.FloatTensor([1.0]),
# torch.FloatTensor([0.8]),
)[0]
.detach()
.numpy()
)
audio = audio_float_to_int16(audio)
end_time = time.perf_counter()
audio_duration_sec = audio.shape[-1] / args.sample_rate
infer_sec = end_time - start_time
real_time_factor = (
infer_sec / audio_duration_sec if audio_duration_sec > 0 else 0.0
)
_LOGGER.debug(
"Real-time factor for %s: %0.2f (infer=%0.2f sec, audio=%0.2f sec)",
i + 1,
real_time_factor,
infer_sec,
audio_duration_sec,
)
output_path = args.output_dir / f"{utt_id}.wav"
write_wav(str(output_path), args.sample_rate, audio)
if __name__ == "__main__":
main()