import os
import gradio as gr
import random

os.system("pip install kantts -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html")
os.system("pip install numpy==1.22.0")

from modelscope.models.audio.tts import SambertHifigan
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

model_dir = os.path.abspath("./pretrain_work_dir")

custom_infer_abs = {
    'voice_name':
    'F7',
    'am_ckpt':
    os.path.join(model_dir, 'tmp_am', 'ckpt'),
    'am_config':
    os.path.join(model_dir, 'tmp_am', 'config.yaml'),
    'voc_ckpt':
    os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan', 'ckpt'),
    'voc_config':
    os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan',
             'config.yaml'),
    'audio_config':
    os.path.join(model_dir, 'data', 'audio_config.yaml'),
    'se_file':
    os.path.join(model_dir, 'data', 'se', 'se.npy')
}
kwargs = {'custom_ckpt': custom_infer_abs}

model_id = SambertHifigan(os.path.join(model_dir, "orig_model"), **kwargs)

inference = pipeline(task=Tasks.text_to_speech, model=model_id)

def infer(text):
    output = inference(input=text)
    filename = str(random.randint(1, 1000000000000))
    
    with open(filename + "myfile.wav", mode='bx') as f:
        f.write(output["output_wav"])
    return filename + "myfile.wav"

app = gr.Blocks()

with app:
    gr.HTML("<center>"
            "<h1>🥳🎶🎡 - KanTTS中文声音克隆</h1>"
            "</center>")
    gr.Markdown("## <center>🌊 - 更多精彩应用，敬请关注[滔滔AI](http://www.talktalkai.com)；滔滔AI，为爱滔滔！💕</center>")

   
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(lines=5, label="请填写您想要转换的中文文本")
            btn = gr.Button("一键生成")
        out = gr.Audio(label="为您生成的专属音频", type="filepath")

        btn.click(fn=infer, inputs=[inp], outputs=[out])

app.launch(show_error=True)