File size: 930 Bytes
12da6cc 725d577 2d6b01a 725d577 12da6cc 725d577 12da6cc 7c3c326 12da6cc 557399e |
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 pickle
import haiku as hk
import jax
import jax.numpy as jnp
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=FLAGS.ckpt_dir / "config.json",
ckpt_file=FLAGS.ckpt_dir / "hk_hifi.pickle",
):
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(ckpt_file, "rb") as f:
params = pickle.load(f)
aux = {}
wav, aux = forward.apply(params, aux, rng, mel)
wav = jnp.squeeze(wav)
jax.config.update('jax_platform_name', 'cpu')
audio = jax.device_get(wav)
return audio |