import torch
import torchaudio
import gradio as gr
from scipy.io import wavfile
from scipy.io.wavfile import write

knn_vc = torch.hub.load('bshall/knn-vc', 'knn_vc', prematched=True, trust_repo=True, pretrained=True, device='cpu')

def voice_change(audio_in, audio_ref):
  samplerate1, data1 = wavfile.read(audio_in)
  samplerate2, data2 = wavfile.read(audio_ref)
  write("./audio_in.wav", samplerate1, data1)
  write("./audio_ref.wav", samplerate2, data2)

  query_seq = knn_vc.get_features("./audio_in.wav")
  matching_set = knn_vc.get_matching_set(["./audio_ref.wav"])
  out_wav = knn_vc.match(query_seq, matching_set, topk=4)
  torchaudio.save('output.wav', out_wav[None], 16000)
  return 'output.wav'



app = gr.Blocks()

with app:
    gr.Markdown("# <center>🥳🎶🎡 - KNN-VC AI变声</center>")
    gr.Markdown("### <center>🌟 - 3秒实时AI变声，支持中日英在内的所有语言！无需训练、一键变声！🍻 </center>")
    gr.Markdown("### <center>🌊 - 更多精彩应用，敬请关注[滔滔AI](http://www.talktalkai.com)；滔滔AI，为爱滔滔！💕</center>")

    with gr.Row():
      with gr.Column():
        inp1 = gr.Audio(type="filepath", label="请上传AI变声的原音频（决定变声后的语音内容）")
        inp2 = gr.Audio(type="filepath", label="请上传AI变声的参照音频（决定变声后的语音音色）")
        btn1 = gr.Button("一键开启AI变声吧", variant="primary")
      with gr.Column():
        out1 = gr.Audio(type="filepath", label="AI变声后的专属音频")

      btn1.click(voice_change, [inp1, inp2], out1)

    gr.Markdown("### <center>注意❗：请不要生成会对个人以及组织造成侵害的内容，此程序仅供科研、学习及个人娱乐使用。</center>")
    gr.HTML('''
        <div class="footer">
                    <p>🌊🏞️🎶 - 江水东流急，滔滔无尽声。 明·顾璘
                    </p>
        </div>
    ''')

app.launch(show_error=True)