#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import gradio as gr
import torch
from huggingface_hub import hf_hub_download

from matcha.cli import get_torch_device, load_matcha, load_vocoder, process_text, to_waveform
from matcha.utils.utils import plot_tensor


TITLE="🍵 Matcha-TTS: A fast TTS architecture with conditional flow matching"
DESCRIPTION = """# chuyển đổi chuyển văn bản thành giọng nói, sử dụng giọng đọc của nhà văn Nguyễn Ngọc Ngạn"""


DEMO_TEXT_0 = "Đây là một chương trình chuyển đổi văn bản thành giọng nói, sử dụng giọng đọc của nhà văn Nguyễn Ngọc Ngạn"
DEMO_TEXT_1 = "Kính thưa quý vị, xin quý vị ghé vào thăm kênh Youtube Nguyễn Ngọc Ngạn, quý vị sẽ gặp lại tất cả các băng đọc truyện của Nguyễn Ngọc Ngạn do Trung tâm Thúy Nga thực hiện và những truyện mới cùng những buổi nói chuyện về nhiều đề tài phổ biến khác nhau. Xin chân thành cảm ơn và chờ đón quý vị."


# model files already pre-downloaded, see README.md >>> preload_from_hub
MODEL_PATH = hf_hub_download(repo_id="doof-ferb/matcha_ngngngan", filename="ckpt/checkpoint_epoch420_slim.pt")
VOCODER_PATH = hf_hub_download(repo_id="doof-ferb/matcha_ngngngan", filename="hifigan/g_02500000")
DEVICE = get_torch_device()
MODEL = load_matcha(MODEL_PATH, DEVICE)
VOCODER, DENOISER = load_vocoder(VOCODER_PATH, DEVICE)


@torch.inference_mode()
def process_text_gradio(text):
	output = process_text(text, DEVICE)
	return output["x_phones"][1::2], output["x"], output["x_lengths"]


@torch.inference_mode()
def synthesise_mel(text, text_length, n_timesteps, temperature, length_scale):
	output = MODEL.synthesise(text, text_length, n_timesteps=n_timesteps,temperature=temperature, spks=None, length_scale=length_scale)
	waveform = to_waveform(output["mel"], VOCODER, DENOISER).numpy()
	return (22050, waveform), plot_tensor(output["mel"].squeeze().cpu().numpy())
	# sample rate 22.05 kHz


def example_cacher(text, n_timesteps, mel_temp, length_scale):
	phones, text, text_lengths = process_text_gradio(text)
	audio, mel_spectrogram = synthesise_mel(text, text_lengths, n_timesteps, mel_temp, length_scale)
	return phones, audio, mel_spectrogram


with gr.Blocks(title=TITLE, theme="soft") as demo:
	processed_text = gr.State(value=None)
	processed_text_len = gr.State(value=None)

	with gr.Row():
		gr.Markdown(DESCRIPTION)

	with gr.Group():
		with gr.Row():
			text = gr.Textbox(label="nhập văn bản tiếng việt", lines=3)
			synth_btn = gr.Button("Chuyển thành giọng nói", variant="primary", scale=0, size="lg")
		with gr.Accordion(label="tuỳ chọn nâng cao", open=False):
			with gr.Row():
				n_timesteps = gr.Slider(label="số bước xoá nhiễu", minimum=1, maximum=100, step=1, value=50, interactive=True)
				length_scale = gr.Slider(label="length scale", info="càng lớn thì đọc càng chậm", minimum=0.5, maximum=1.5,  step=0.05, value=.95, interactive=True)
				mel_temp = gr.Slider(label="sampling temperature", minimum=0.00, maximum=2.001, step=0.16675, value=0.667, interactive=True,)
				# idk min max of denoiser strength so let it default at 2.5e-4

	with gr.Group():
		with gr.Row():
			audio = gr.Audio(label="audio", interactive=False)
		with gr.Accordion(label="thông tin chuyên sâu", open=False):
			with gr.Row():
				phonetised_text = gr.Textbox(label="Văn bản dưới dạng mẫu tự biểu âm quốc tế (IPA)", info="khẩu âm Hà Nội", interactive=False, lines=3)
				mel_spectrogram = gr.Image(label="mel spectrogram", interactive=False)

	with gr.Row():
		examples = gr.Examples(
			label="ví dụ văn bản đầu vào",
			examples=[
				[DEMO_TEXT_0,  2, 0.677, 0.95],
				[DEMO_TEXT_0,  4, 0.677, 0.95],
				[DEMO_TEXT_0, 10, 0.677, 0.95],
				[DEMO_TEXT_0, 50, 0.677, 0.95],
				[DEMO_TEXT_1, 50, 0.677, 0.95],
			],
			fn=example_cacher,
			inputs=[text, n_timesteps, mel_temp, length_scale],
			outputs=[phonetised_text, audio, mel_spectrogram],
			cache_examples=True,
		)

	with gr.Row():
		gr.Markdown("[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/#fileId=https://huggingface.co/spaces/doof-ferb/MatchaTTS_ngngngan/blob/main/colab-notebook.ipynb)")

	synth_btn.click(
		fn=process_text_gradio,
		inputs=[text],
		outputs=[phonetised_text, processed_text, processed_text_len],
		api_name="phonemize",
		queue=True,
	).then(
		fn=synthesise_mel,
		inputs=[processed_text, processed_text_len, n_timesteps, mel_temp, length_scale],
		outputs=[audio, mel_spectrogram],
		api_name="speak",
	)

demo.queue().launch(server_port=8080)