import gradio as gr
from TTS.api import TTS

# Initialize TTS models
tts = TTS(model_name="tts_models/multilingual/multi-dataset/your_tts", progress_bar=False, gpu=False)

def text_to_speech(text: str, speaker_wav, speaker_wav_file, language: str):
    if speaker_wav_file and not speaker_wav:
        speaker_wav = speaker_wav_file
    file_path = "output.wav"
    if speaker_wav is not None:
        tts.tts_to_file(text, speaker_wav=speaker_wav, language=language, file_path=file_path)
    else:
        tts.tts_to_file(text, speaker=tts.speakers[0], language=language, file_path=file_path)
    return file_path

title = "Prateek's Voice Cloning Demo"

def toggle(choice):
    if choice == "mic":
        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
    else:
        return gr.update(visible=False, value=None), gr.update(visible=True, value=None)

with gr.Blocks() as demo:
    gr.Markdown("# 🎙️ Prateek's Voice Cloning Demo")
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(label="Enter the text you want to convert to speech", value="", max_lines=3)
            lan_input = gr.Radio(label="Language", choices=["en"], value="en", visible=False)
            radio = gr.Radio(["Record your voice", "Upload an audio file"], value="mic",
                             label="How would you like to provide your voice sample?")
            audio_input_mic = gr.Audio(label="Record your voice", source="microphone", type="filepath", visible=True)
            audio_input_file = gr.Audio(label="Upload an audio file", type="filepath", visible=False)

            with gr.Row():
                with gr.Column():
                    btn_clear = gr.Button("Clear")
                with gr.Column():
                    btn = gr.Button("Generate Speech", variant="primary")
        with gr.Column():
            audio_output = gr.Audio(label="Generated Speech")

    btn.click(text_to_speech, inputs=[text_input, audio_input_mic,
              audio_input_file, lan_input], outputs=audio_output)
    radio.change(toggle, radio, [audio_input_mic, audio_input_file])

demo.launch(enable_queue=True)