initial commit

.gitignore

@@ -0,0 +1,11 @@
+# Python cache
+__pycache__/
+*.pyc
+.ipynb_checkpoints
+
+notebooks/
+
+secrets.toml
+
+cache_sound/*.wav
+assets/*/*.onnx

.streamlit/config.toml

@@ -0,0 +1,6 @@
+[theme]
+
+base="dark"
+primaryColor="#4551b2"
+backgroundColor="#292a57"
+textColor="#faf9fc"

README.md

@@ -1,13 +0,0 @@
----
-title: Nix Tts
-emoji: 📚
-colorFrom: purple
-colorTo: red
-sdk: streamlit
-sdk_version: 1.2.0
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -0,0 +1,56 @@
+# Utils
+import os
+import soundfile as sf
+
+# Streamlit
+import streamlit as st
+
+# Custom elements
+from elements.component import (
+    centered_text,
+)
+from elements.session_states import (
+    init_session_state,
+    update_session_state,
+    update_model,
+)
+from elements.tts import (
+    generate_voice,
+)
+
+st.set_page_config(
+    page_title = "Nix-TTS Interactive Demo",
+    page_icon = "🐤",
+)
+
+# Initiate stuffs
+init_session_state()
+
+# ---------------------------------------------------------------------------------
+
+# Description
+centered_text("🐤 Nix-TTS Interactive Demo")
+centered_text("An incredibly lightweight end-to-end text-to-speech model via knowledge distillation", "h5")
+st.write(" ")
+st.caption("🗒️ This is a demo from our latest paper, **Nix-TTS**. <br>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;You can access the paper and the released models [here](https://github.com/rendchevi/nix-tts). <br>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Authors: Rendi Chevi, Radityo Eko Prasojo, Alham Fikri Aji.<br>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;**Corresponding Author**: Rendi Chevi | rendi.chevi@{kata.ai, gmail.com}.", True)
+
+# Model demo
+st.write(" ")
+st.write(" ")
+col1, col2 = st.columns(2)
+with col1:
+    input_text = st.text_input(
+        "Input Text",
+        value = "Born to multiply, born to gaze into night skies.",
+    )
+with col2:
+    model_variant = st.selectbox("Choose Model Variant", options = ["Deterministic", "Stochastic"], index = 1)
+    if model_variant != st.session_state.model_variant:
+        # Update variant choice
+        update_session_state("model_variant", model_variant)
+        # Re-load model
+        update_model()
+
+button_gen = st.button("Generate Voice")
+if button_gen == True:
+    generate_voice(input_text)

elements/component.py

@@ -0,0 +1,11 @@
+# Streamlit
+import streamlit as st
+
+def centered_text(
+    input_text,
+    mode = "h1",
+):
+    st.markdown(
+        f"<{mode} style='text-align: center;'>{input_text}</{mode}>",
+        unsafe_allow_html = True
+    )

elements/session_states.py

@@ -0,0 +1,30 @@
+# Utils
+import uuid
+
+# Streamlit
+import streamlit as st
+
+# Nix
+from nix.models.TTS import NixTTSInference
+
+# --------------------- SESSION STATE MANAGEMENT -------------------------
+
+def init_session_state():
+    # Model
+    if "init_model" not in st.session_state:
+        st.session_state.init_model = True
+        st.session_state.random_str = uuid.uuid1().hex
+        st.session_state.model_variant = "Stochastic"
+        st.session_state.TTS = NixTTSInference("assets/nix-ljspeech-sdp-v0.1")
+
+def update_model():
+    if st.session_state.model_variant == "Deterministic":
+        st.session_state.TTS = NixTTSInference("assets/nix-ljspeech-v0.1")
+    elif st.session_state.model_variant == "Stochastic":
+        st.session_state.TTS = NixTTSInference("assets/nix-ljspeech-sdp-v0.1")
+
+def update_session_state(
+    state_id,
+    state_value,
+):
+    st.session_state[f"{state_id}"] = state_value

elements/tts.py

@@ -0,0 +1,43 @@
+# Utils
+import timeit
+import soundfile as sf
+
+# Streamlit
+import streamlit as st
+
+# Custom elements
+from elements.component import (
+    centered_text,
+)
+
+def generate_voice(
+    input_text,
+):
+    # TTS Inference
+    start_time = timeit.default_timer()
+    c, c_length, phoneme = st.session_state.TTS.tokenize(input_text)
+    tok_time = timeit.default_timer() - start_time
+
+    start_time = timeit.default_timer()
+    voice = st.session_state.TTS.vocalize(c, c_length)
+    tts_time = timeit.default_timer() - start_time
+
+    # Time stats
+    total_infer_time = tts_time + tok_time
+    audio_time = voice.shape[-1] / 22050
+    rtf = total_infer_time / audio_time
+    rt_ratio = 1 / rtf
+
+    # Save audio (bug in Streamlit, can't play numpy array directly)
+    sf.write(f"cache_sound/{st.session_state.random_str}.wav", voice[0,0], 22050)
+
+    # Play audio
+    st.audio(f"cache_sound/{st.session_state.random_str}.wav", format = "audio/wav")
+    st.caption("Generated Voice")
+
+    st.code(
+        f"💬 Output Audio: {str(audio_time)[:6]} sec.\n\n⏳ Elapsed time for:\n   => Tokenization: {str(tok_time)[:6]} sec.\n   => Model Inference: {str(tts_time)[:6]} sec.\n\n⏰ Real-time Factor (RTF): {str(rtf)[:6]}\n\n🏃 The model runs {str(rt_ratio)[:6]} x faster than real-time \
+        ",
+        language = "bash",
+    )
+    st.caption("Elapsed Time Stats")

nix/models/TTS.py

@@ -0,0 +1,55 @@
+import os
+import pickle
+import timeit
+
+import numpy as np
+import onnxruntime as ort
+
+from nix.tokenizers.tokenizer_en import NixTokenizerEN
+
+class NixTTSInference:
+
+    def __init__(
+        self,
+        model_dir,
+    ):
+        # Load tokenizer
+        self.tokenizer = NixTokenizerEN(pickle.load(open(os.path.join(model_dir, "tokenizer_state.pkl"), "rb")))
+        # Load TTS model
+        self.encoder = ort.InferenceSession(os.path.join(model_dir, "encoder.onnx"))
+        self.decoder = ort.InferenceSession(os.path.join(model_dir, "decoder.onnx"))
+
+    def tokenize(
+        self,
+        text,
+    ):
+        # Tokenize input text
+        c, c_lengths, phonemes = self.tokenizer([text])
+
+        return np.array(c, dtype = np.int64), np.array(c_lengths, dtype = np.int64), phonemes
+
+    def vocalize(
+        self,
+        c,
+        c_lengths,
+    ):
+        """
+        Single-batch TTS inference
+        """
+        # Infer latent samples from encoder
+        z = self.encoder.run(
+            None,
+            {
+                "c": c,
+                "c_lengths": c_lengths,
+            }
+        )[2]
+        # Decode raw audio with decoder
+        xw = self.decoder.run(
+            None,
+            {
+                "z": z,
+            }
+        )[0]
+
+        return xw

nix/tokenizers/tokenizer_en.py

@@ -0,0 +1,77 @@
+# Regex
+import re
+
+# Phonemizer
+from phonemizer.backend import EspeakBackend
+phonemizer_backend = EspeakBackend(
+    language = 'en-us',
+    preserve_punctuation = True,
+    with_stress = True
+)
+
+class NixTokenizerEN:
+
+    def __init__(
+        self,
+        tokenizer_state,
+    ):
+        # Vocab and abbreviations dictionary
+        self.vocab_dict = tokenizer_state["vocab_dict"]
+        self.abbreviations_dict = tokenizer_state["abbreviations_dict"]
+
+        # Regex recipe
+        self.whitespace_regex = tokenizer_state["whitespace_regex"]
+        self.abbreviations_regex = tokenizer_state["abbreviations_regex"]
+
+    def __call__(
+        self,
+        texts,
+    ):
+        # 1. Phonemize input texts
+        phonemes = [ self._collapse_whitespace(
+            phonemizer_backend.phonemize(
+                self._expand_abbreviations(text.lower()),
+                strip = True,
+            )
+        ) for text in texts ]
+
+        # 2. Tokenize phonemes
+        tokens = [ self._intersperse([self.vocab_dict[p] for p in phoneme], 0) for phoneme in phonemes ]
+
+        # 3. Pad tokens
+        tokens, tokens_lengths = self._pad_tokens(tokens)
+
+        return tokens, tokens_lengths, phonemes
+
+    def _expand_abbreviations(
+        self,
+        text
+    ):
+        for regex, replacement in self.abbreviations_regex:
+            text = re.sub(regex, replacement, text)
+
+        return text
+
+    def _collapse_whitespace(
+        self,
+        text
+    ):
+        return re.sub(self.whitespace_regex, ' ', text)
+
+    def _intersperse(
+        self,
+        lst,
+        item,
+    ):
+        result = [item] * (len(lst) * 2 + 1)
+        result[1::2] = lst
+        return result
+
+    def _pad_tokens(
+        self,
+        tokens,
+    ):
+        tokens_lengths = [len(token) for token in tokens]
+        max_len = max(tokens_lengths)
+        tokens = [token + [0 for _ in range(max_len - len(token))] for token in tokens]
+        return tokens, tokens_lengths

packages.txt

@@ -0,0 +1,2 @@
+libsndfile1-dev
+espeak

requirements.txt

@@ -0,0 +1,5 @@
+streamlit
+numpy
+onnxruntime
+phonemizer
+SoundFile