#!/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

original: 🍵 https://github.com/shivammehta25/Matcha-TTS

my personal checkpoints: a speech-synthesis model of Vietnamese M.C. Nguyễn Ngọc Ngạn

the checkpoints is released under CC-BY-NC-SA-4.0

In accordance with the terms of the CC-BY-NC-SA-4.0 license, the use of my checkpoints and any audio output generated by them for commercial purposes is strictly prohibited. This includes, but is not limited to:
- online and offline voice cloning as a service
- online and offline text-to-speech as a service
- content creation for monetization on social media platforms

Căn cứ vào các điều khoản của giấp phép CC-BY-NC-SA-4.0, việc sử dụng các checkpoints này và bất kỳ đầu ra âm thanh nào được tạo bởi chúng đều bị nghiêm cấm sử dụng cho mục đích thương mại. Điều này bao gồm, nhưng không giới hạn ở:
- các dịch vụ nhân bản giọng nói trực tuyến và ngoại tuyến
- các dịch vụ chuyển văn bản thành giọng nói trực tuyến và ngoại tuyến
- tạo nội dung để kiếm tiền trên các nền tảng mạng xã hội"""


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(default_concurrency_limit=5).launch()