import gradio as gr
import torch
from TTS.api import TTS
import os

# Get device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize TTS model
tts = TTS(model_name="voice_conversion_models/multilingual/vctk/freevc24", progress_bar=False).to(device)

# Get examples from Examples folder
examples_folder = "Examples/"
example_files = [f for f in os.listdir(examples_folder) if f.endswith(".wav")]

def voice_conversion(input_audio, target_voice):
    output_path = "output.wav"
    # Perform voice conversion
    target_voice = f"{examples_folder}{target_voice}"
    print(f"Target voice is: {target_voice}")
    tts.voice_conversion_to_file(source_wav=input_audio, target_wav=target_voice, file_path=output_path)
    return output_path

# Define Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Voice Conversion using Coqui TTS")
    
    with gr.Row():
        input_audio = gr.Audio(label="Record or Upload Your Voice", type="filepath")
        target_voice = gr.Dropdown(choices=example_files, label="Select Target Voice from Examples", 
                                   value=example_files[0], info="Located in Examples/ folder")

    convert_button = gr.Button("Convert Voice")
    output_audio = gr.Audio(label="Converted Voice", type="filepath")

    convert_button.click(voice_conversion, inputs=[input_audio, target_voice], outputs=output_audio)

# Launch with public=True for public URL access and share link
demo.launch(share=True)