File size: 4,000 Bytes
12da6cc 725d577 12da6cc 725d577 12da6cc 725d577 12da6cc 725d577 12da6cc 2a0dc49 725d577 12da6cc 449bd33 725d577 449bd33 12da6cc 725d577 |
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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import pickle
from argparse import ArgumentParser
from pathlib import Path
import haiku as hk
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from .config import FLAGS, DurationInput
from .data_loader import load_phonemes_set
from .model import AcousticModel, DurationModel
def load_lexicon(fn):
lines = open(fn, "r").readlines()
lines = [l.lower().strip().split("\t") for l in lines]
return dict(lines)
def predict_duration(tokens, ckpt_file):
def fwd_(x):
return DurationModel(is_training=False)(x)
forward_fn = jax.jit(hk.transform_with_state(fwd_).apply)
with open(ckpt_file, "rb") as f:
dic = pickle.load(f)
x = DurationInput(
np.array(tokens, dtype=np.int32)[None, :],
np.array([len(tokens)], dtype=np.int32),
None,
)
return forward_fn(dic["params"], dic["aux"], dic["rng"], x)[0]
def text2tokens(text, lexicon_fn):
phonemes = load_phonemes_set()
lexicon = load_lexicon(lexicon_fn)
words = text.strip().lower().split()
tokens = [FLAGS.sil_index]
for word in words:
if word in FLAGS.special_phonemes:
tokens.append(phonemes.index(word))
elif word in lexicon:
p = lexicon[word]
p = p.split()
p = [phonemes.index(pp) for pp in p]
tokens.extend(p)
tokens.append(FLAGS.word_end_index)
else:
for p in word:
if p in phonemes:
tokens.append(phonemes.index(p))
tokens.append(FLAGS.word_end_index)
tokens.append(FLAGS.sil_index) # silence
return tokens
def predict_mel(tokens, durations, ckpt_fn):
with open(ckpt_fn, "rb") as f:
dic = pickle.load(f)
last_step, params, aux, rng, optim_state = (
dic["step"],
dic["params"],
dic["aux"],
dic["rng"],
dic["optim_state"],
)
@hk.transform_with_state
def forward(tokens, durations, n_frames):
net = AcousticModel(is_training=False)
return net.inference(tokens, durations, n_frames)
durations = durations * FLAGS.sample_rate / (FLAGS.n_fft // 4)
n_frames = int(jnp.sum(durations).item())
predict_fn = jax.jit(forward.apply, static_argnums=[5])
tokens = np.array(tokens, dtype=np.int32)[None, :]
return predict_fn(params, aux, rng, tokens, durations, n_frames)[0]
def text2mel(
text: str,
lexicon_fn=FLAGS.data_dir / "lexicon.txt",
silence_duration: float = -1.0,
acoustic_ckpt=FLAGS.ckpt_dir / "acoustic_latest_ckpt.pickle",
duration_ckpt=FLAGS.ckpt_dir / "duration_latest_ckpt.pickle",
):
tokens = text2tokens(text, lexicon_fn)
print("text2mel-token : ")
print(tokens)
durations = predict_duration(tokens, duration_ckpt)
durations = jnp.where(
np.array(tokens)[None, :] == FLAGS.sil_index,
jnp.clip(durations, a_min=silence_duration, a_max=None),
durations,
)
durations = jnp.where(
np.array(tokens)[None, :] == FLAGS.word_end_index, 0.0, durations
)
print("acoustic_ckpt : ")
print(acoustic_ckpt)
print("duration_ckpt : ")
print(duration_ckpt)
mels = predict_mel(tokens, durations, acoustic_ckpt)
print("mels : ")
print(mels)
if tokens[-1] == FLAGS.sil_index:
end_silence = durations[0, -1].item()
silence_frame = int(end_silence * FLAGS.sample_rate / (FLAGS.n_fft // 4))
mels = mels[:, : (mels.shape[1] - silence_frame)]
return mels
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
mel = text2mel(args.text)
plt.figure(figsize=(10, 5))
plt.imshow(mel[0].T, origin="lower", aspect="auto")
plt.savefig(str(args.output))
plt.close()
mel = jax.device_get(mel)
mel.tofile("clip.mel") |