update

app.py

@@ -7,16 +7,17 @@
 import os
 import re
 import unicodedata
+
 import regex
 
 if not os.path.isfile("./wavegru_mod.so"):
     os.system("bash ./build_ext.sh")
 
 import gradio as gr
+
 from inference import load_tacotron_model, load_wavegru_net, mel_to_wav, text_to_mel
 from wavegru_cpp import extract_weight_mask, load_wavegru_cpp
 
-
 alphabet, tacotron_net, tacotron_config = load_tacotron_model(
     "./alphabet.txt", "./tacotron.toml", "./mono_tts_cbhg_small_0700000.ckpt"
 )
@@ -26,9 +27,7 @@ wavegru_config, wavegru_net = load_wavegru_net(
 )
 
 wave_cpp_weight_mask = extract_weight_mask(wavegru_net)
-wavecpp = load_wavegru_cpp(
-    wave_cpp_weight_mask, wavegru_config["upsample_factors"][-1]
-)
+wavecpp = load_wavegru_cpp(wave_cpp_weight_mask, wavegru_config["upsample_factors"][-1])
 
 
 space_re = regex.compile(r"\s+")

extract_tacotrons_model.py

@@ -1,5 +1,8 @@
 import pickle
 
+import jax
+
 dic = pickle.load(open("./mono_tts_cbhg_small_0700000.ckpt", "rb"))
 del dic["optim_state_dict"]
+dic = jax.device_get(dic)
 pickle.dump(dic, open("./mono_tts_cbhg_small_0700000.ckpt", "wb"))

extract_wavegru_model.py

@@ -1,5 +1,12 @@
 import pickle
 
-dic = pickle.load(open("./wavegru_vocoder_tpu_gta_preemphasis_pruning_0800000.ckpt", "rb"))
+import jax
+
+dic = pickle.load(
+    open("./wavegru_vocoder_tpu_gta_preemphasis_pruning_0800000.ckpt", "rb")
+)
+dic = jax.device_get(dic)
 del dic["optim_state_dict"]
-pickle.dump(dic, open("./wavegru_vocoder_tpu_gta_preemphasis_pruning_0800000.ckpt", "wb"))
+pickle.dump(
+    dic, open("./wavegru_vocoder_tpu_gta_preemphasis_pruning_0800000.ckpt", "wb")
+)

mynumbers.py

@@ -1,8 +1,8 @@
 """ from https://github.com/keithito/tacotron """
 
-import inflect
 import re
 
+import inflect
 
 _inflect = inflect.engine()
 _comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")

wavegru_vocoder_tpu_gta_preemphasis_pruning_0800000.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:17d7801c9aaa72992671be5066953ae9598ee91ecfeb502a88e2fd27aac19cc6
-size 69718512
+oid sha256:1225b31bfc3a79ae4f0884c559b4698b3ea6b33a70b589883bf0bc76e23fb467
+size 69717656