import spaces
import torch
import gradio as gr
from chatterbox_vc import VoiceConverter

vc = None


def get_converter():
    global vc

    if vc is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Loading ChatterboxVC on {device}...")
        vc = VoiceConverter(device=device)

    return vc


@spaces.GPU
def convert(source_audio, target_audio):
    if source_audio is None:
        raise gr.Error("Please upload a source audio file.")

    if target_audio is None:
        raise gr.Error("Please upload a reference voice.")

    output = "output.wav"

    converter = get_converter()

    converter.convert(
        source_audio,
        target_audio,
        output,
    )

    return output


with gr.Blocks() as demo:

    gr.Markdown("# Chatterbox Voice Conversion")

    source = gr.Audio(
        type="filepath",
        label="Source Speech",
    )

    target = gr.Audio(
        type="filepath",
        label="Reference Voice",
    )

    button = gr.Button("Convert")

    output = gr.Audio(
        type="filepath",
        label="Converted Speech",
    )

    button.click(
        convert,
        inputs=[source, target],
        outputs=output,
    )


demo.launch()