import utils
from models import SynthesizerTrn
import torch
from torch import no_grad, LongTensor
from text import text_to_sequence
import gradio as gr
import commons

model_path = "./OUTPUT_MODEL/G_Amitaro.pth"
config_path = "./OUTPUT_MODEL/config.json"

length = 1.0
device = "cuda:0" if torch.cuda.is_available() else "cpu"

def get_text(text, hps, is_symbol):
    text_norm = text_to_sequence(text, hps.symbols, [] if is_symbol else hps.data.text_cleaners)
    if hps.data.add_blank:
        text_norm = commons.intersperse(text_norm, 0)
    text_norm = LongTensor(text_norm)
    return text_norm

def get_vits_array(text):
    hps = utils.get_hparams_from_file(config_path)
    net_g = SynthesizerTrn(
        len(hps.symbols),
        hps.data.filter_length // 2 + 1,
        hps.train.segment_size // hps.data.hop_length,
        n_speakers=hps.data.n_speakers,
        **hps.model).to(device)
    _ = net_g.eval()
    _ = utils.load_checkpoint(model_path, net_g, None)
    
    speaker_ids = hps.speakers
    
    #text = "[JA]" + text + "[JA]"
    speaker_id = 0
    stn_tst = get_text(text, hps, False)
    with no_grad():
        x_tst = stn_tst.unsqueeze(0).to(device)
        x_tst_lengths = LongTensor([stn_tst.size(0)]).to(device)
        sid = LongTensor([speaker_id]).to(device)
        audio = net_g.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=.667, noise_scale_w=0.6,
                            length_scale=1.0 / length)[0][0, 0].data.cpu().float().numpy()
    del stn_tst, x_tst, x_tst_lengths, sid
    
    return (hps.data.sampling_rate, audio)

app = gr.Blocks()
with app:
    gr.Markdown("# VITS-TTS-Japanese-Only-Amitaro\n\n"
        		"Sample usage of Finetune model [Lycoris53/Vits-Japanese-Only-Amitaro](https://huggingface.co/Lycoris53/Vits-Japanese-Only-Amitaro) \n"
                "Base finetuning code is from [Plachtaa/VITS-fast-fine-tuning](https://github.com/Plachtaa/VITS-fast-fine-tuning)"
                )
    with gr.Row():
        with gr.Column():
            textbox = gr.TextArea(label="Text",
                                  placeholder="Type your sentence here (Maximum 150 words)",
                                  value="おはようございます。")
        with gr.Column():
            audio_output = gr.Audio(label="Output Audio")
            btn = gr.Button("Generate　Voice!")
            btn.click(get_vits_array,
                      inputs=[textbox],
                      outputs=[audio_output])

app.queue(concurrency_count=3).launch()