File size: 905 Bytes
12da6cc |
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 |
import json
import os
import pickle
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from .config import FLAGS
from .model import Generator
class AttrDict(dict):
def __init__(self, *args, **kwargs):
super(AttrDict, self).__init__(*args, **kwargs)
self.__dict__ = self
def mel2wave(mel):
config_file = "assets/hifigan/config.json"
MAX_WAV_VALUE = 32768.0
with open(config_file) as f:
data = f.read()
json_config = json.loads(data)
h = AttrDict(json_config)
@hk.transform_with_state
def forward(x):
net = Generator(h)
return net(x)
rng = next(hk.PRNGSequence(42))
with open(FLAGS.ckpt_dir / "hk_hifi.pickle", "rb") as f:
params = pickle.load(f)
aux = {}
wav, aux = forward.apply(params, aux, rng, mel)
wav = jnp.squeeze(wav)
audio = jax.device_get(wav)
return audio
|