from random import choices

import gradio as gr
from TTS.api import TTS

# voice_conversion_models/multilingual/multi-dataset/knnvc
# voice_conversion_models/multilingual/multi-dataset/openvoice_v1
# voice_conversion_models/multilingual/multi-dataset/openvoice_v2
tts = TTS(model_name="voice_conversion_models/multilingual/vctk/freevc24", gpu=False)
tts2 = TTS(model_name="voice_conversion_models/multilingual/multi-dataset/knnvc", gpu=False)
tts3 = TTS(model_name="voice_conversion_models/multilingual/multi-dataset/openvoice_v1", gpu=False)
tts4 = TTS(model_name="voice_conversion_models/multilingual/multi-dataset/openvoice_v2", gpu=False)

def greet(source, target, choices_s):
    if choices_s == "freevc24":
        tts.voice_conversion_to_file(
            source_wav=source,
            target_wav=target,
            file_path="output.wav"
        )
    elif choices_s == "knnvc":
        tts2.voice_conversion_to_file(
            source_wav=source,
            target_wav=target,
            file_path="output.wav"
        )
    elif choices_s == "openvoice_v1":
        tts3.voice_conversion_to_file(
            source_wav=source,
            target_wav=target,
            file_path="output.wav"
        )
    elif choices_s == "openvoice_v2":
        tts4.voice_conversion_to_file(
            source_wav=source,
            target_wav=target,
            file_path="output.wav"
        )
    return "output.wav"

demo = gr.Interface(fn=greet, inputs=[gr.Audio(label="Source audio", value="kratos.wav", type="filepath"),
                                      gr.Audio(label="Target audio", value="nikole_kidman.wav", type="filepath"), 
                                      gr.Dropdown(label="Model name", choices=["freevc24", "knnvc", "openvoice_v1", "openvoice_v2"])],
                    outputs=gr.Audio(type="filepath"))
demo.launch()