|
import pickle |
|
from argparse import ArgumentParser |
|
from pathlib import Path |
|
|
|
import haiku as hk |
|
import jax |
|
import jax.numpy as jnp |
|
import numpy as np |
|
from tqdm.auto import tqdm |
|
from vietTTS.nat.config import AcousticInput |
|
|
|
from .config import FLAGS, AcousticInput |
|
from .data_loader import load_textgrid_wav |
|
from .dsp import MelFilter |
|
from .model import AcousticModel |
|
|
|
|
|
@hk.transform_with_state |
|
def net(x): |
|
return AcousticModel(is_training=True)(x) |
|
|
|
|
|
@hk.transform_with_state |
|
def val_net(x): |
|
return AcousticModel(is_training=False)(x) |
|
|
|
|
|
def forward_fn_(params, aux, rng, inputs: AcousticInput): |
|
melfilter = MelFilter( |
|
FLAGS.sample_rate, FLAGS.n_fft, FLAGS.mel_dim, FLAGS.fmin, FLAGS.fmax |
|
) |
|
mels = melfilter(inputs.wavs.astype(jnp.float32) / (2**15)) |
|
B, L, D = mels.shape |
|
inp_mels = jnp.concatenate( |
|
(jnp.zeros((B, 1, D), dtype=jnp.float32), mels[:, :-1, :]), axis=1 |
|
) |
|
n_frames = inputs.durations * FLAGS.sample_rate / (FLAGS.n_fft // 4) |
|
inputs = inputs._replace(mels=inp_mels, durations=n_frames) |
|
(mel1_hat, mel2_hat), new_aux = val_net.apply(params, aux, rng, inputs) |
|
return mel2_hat |
|
|
|
|
|
forward_fn = jax.jit(forward_fn_) |
|
|
|
|
|
def generate_gta(out_dir: Path): |
|
out_dir.mkdir(parents=True, exist_ok=True) |
|
data_iter = load_textgrid_wav( |
|
FLAGS.data_dir, |
|
FLAGS.max_phoneme_seq_len, |
|
FLAGS.batch_size, |
|
FLAGS.max_wave_len, |
|
"gta", |
|
) |
|
ckpt_fn = FLAGS.ckpt_dir / "acoustic_latest_ckpt.pickle" |
|
print("Resuming from latest checkpoint at", ckpt_fn) |
|
with open(ckpt_fn, "rb") as f: |
|
dic = pickle.load(f) |
|
_, params, aux, rng, _ = ( |
|
dic["step"], |
|
dic["params"], |
|
dic["aux"], |
|
dic["rng"], |
|
dic["optim_state"], |
|
) |
|
|
|
tr = tqdm(data_iter) |
|
for names, batch in tr: |
|
lengths = batch.wav_lengths |
|
predicted_mel = forward_fn(params, aux, rng, batch) |
|
mel = jax.device_get(predicted_mel) |
|
for idx, fn in enumerate(names): |
|
file = out_dir / f"{fn}.npy" |
|
tr.write(f"saving to file {file}") |
|
l = lengths[idx] // (FLAGS.n_fft // 4) |
|
np.save(file, mel[idx, :l].T) |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = ArgumentParser() |
|
parser.add_argument("-o", "--output-dir", type=Path, default="gta") |
|
generate_gta(parser.parse_args().output_dir) |
|
|
