#@title Install software

print("Installing...")
!git clone -q https://github.com/rhasspy/piper
%cd /content/piper/src/python
!pip install -q cython>=0.29.0 espeak-phonemizer>=1.1.0 librosa>=0.9.2 numpy>=1.19.0 pytorch-lightning~=1.7.0 torch~=1.11.0
!pip install -q onnx onnxruntime-gpu
!bash build_monotonic_align.sh
!apt-get install espeak-ng
!pip install -q torchtext==0.12.0
# fixing recent compativility isswes:
!pip install -q torchaudio==0.11.0 torchmetrics==0.11.4
print("Done!")

#@title Voice package generation section
%cd /content/piper/src/python
import os
import ipywidgets as widgets
from IPython.display import display
import json
from google.colab import output
guideurl = "https://github.com/rmcpantoja/piper/blob/master/notebooks/wav/en"
#@markdown #### Download:
#@markdown **Drive ID or direct download link of the model in another cloud:**
model_id = "" #@param {type:"string"}
#@markdown **Drive ID or direct download link of the config.json file:**
config_id = "" #@param {type:"string"}
#@markdown ---

#@markdown #### Creation process:
#@markdown **Choose the language code (iso639-1 format):**

#@markdown You can see a list of language codes and names [here](https://www.loc.gov/standards/iso639-2/php/English_list.php)

language = "en-us" #@param ["ca", "da", "de", "en", "en-us", "es", "fi", "fr", "grc", "is", "it", "k", "nb", "ne", "nl", "pl", "pt-br", "ru", "sv", "uk", "vi-vn-x-central", "yue"]
voice_name = "Myvoice" #@param {type:"string"}
voice_name = voice_name.lower()
quality = "medium" #@param ["high", "low", "medium", "x-low"]
def start_process():
    if not os.path.exists("/content/project/model.ckpt"):
        raise Exception("Could not download model! make sure the file is shareable to everyone")
    output.eval_js(f'new Audio("{guideurl}/starting.wav?raw=true").play()')
    !python -m piper_train.export_onnx "/content/project/model.ckpt" "{export_voice_path}/{export_voice_name}.onnx"
    print("compressing...")
    !tar -czvf "{packages_path}/voice-{export_voice_name}.tar.gz" -C "{export_voice_path}" .
    output.eval_js(f'new Audio("{guideurl}/success.wav?raw=true").play()')
    print("Done!")

export_voice_name = f"{language}-{voice_name}-{quality}"
export_voice_path = "/content/project/voice-"+export_voice_name
packages_path = "/content/project/packages"
if not os.path.exists(export_voice_path):
    os.makedirs(export_voice_path)
if not os.path.exists(packages_path):
    os.makedirs(packages_path)
print("Downloading model and his config...")
if model_id.startswith("1"):
    !gdown -q "{model_id}" -O /content/project/model.ckpt
elif model_id.startswith("https://drive.google.com/file/d/"):
    !gdown -q "{model_id}" -O "/content/project/model.ckpt" --fuzzy
else:
    !wget "{model_id}" -O "/content/project/model.ckpt"
if config_id.startswith("1"):
    !gdown -q "{config_id}" -O "{export_voice_path}/{export_voice_name}.onnx.json"
elif config_id.startswith("https://drive.google.com/file/d/"):
    !gdown -q "{config_id}" -O "{export_voice_path}/{export_voice_name}.onnx.json" --fuzzy
else:
    !wget "{config_id}" -O "{export_voice_path}/{export_voice_name}.onnx.json"
#@markdown **Do you want to write a model card?**
write_model_card = False #@param {type:"boolean"}
if write_model_card:
    with open(f"{export_voice_path}/{export_voice_name}.onnx.json", "r") as file:
        config = json.load(file)
    sample_rate = config["audio"]["sample_rate"]
    num_speakers = config["num_speakers"]
    output.eval_js(f'new Audio("{guideurl}/waiting.wav?raw=true").play()')
    text_area = widgets.Textarea(
        description = "fill in this following template and press start to generate the voice package",
        value=f'# Model card for {voice_name} ({quality})\n\n* Language: {language} (normaliced)\n* Speakers: {num_speakers}\n* Quality: {quality}\n* Samplerate: {sample_rate}Hz\n\n## Dataset\n\n* URL: \n* License: \n\n## Training\n\nTrained from scratch.\nOr finetuned from: ',
        layout=widgets.Layout(width='500px', height='200px')
    )
    button = widgets.Button(description='Start')

    def create_model_card(button):
        model_card_text = text_area.value.strip()
        with open(f'{export_voice_path}/MODEL_CARD', 'w') as file:
            file.write(model_card_text)
        text_area.close()
        button.close()
        output.clear()
        start_process()

    button.on_click(create_model_card)

    display(text_area, button)
else:
    start_process()

#@title Download/export your generated voice package

#@markdown #### How do you want to export your model?
export_mode = "upload it to my Google Drive" #@param ["Download the voice package on my device (may take some time)", "upload it to my Google Drive"]
print("Exporting package...")
if export_mode == "Download the voice package on my device (may take some time)":
    from google.colab import files
    files.download(f"{packages_path}/voice-{export_voice_name}.tar.gz")
    msg = "Please wait a moment while the package is being downloaded."
else:
    voicepacks_folder = "/content/drive/MyDrive/piper voice packages"
    from google.colab import drive
    drive.mount('/content/drive')
    if not os.path.exists(voicepacks_folder):
        os.makedirs(voicepacks_folder)
    !cp "{packages_path}/voice-{export_voice_name}.tar.gz" "{voicepacks_folder}"
    msg = f"You can find the generated voice package at: {voicepacks_folder}."
print(f"Done! {msg}")