import random
import numpy as np
import torch
from chatterbox.src.orator.tts import OratorTTS
import gradio as gr

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"



model = OratorTTS.from_pretrained(DEVICE)

def generate(text, audio_prompt_path, exaggeration, pace, temperature, seed_num):
    with torch.inference_mode():
        wav = model.generate(
            text,
            audio_prompt_path=audio_prompt_path,
            emotion_adv=exaggeration,
        )
    return model.sr, wav.squeeze(0).numpy()


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(value="I know what you're thinking. \"Did he fire six shots, or only five?\" Well, to tell you the truth, in all this excitement, I kind of lost track myself.", label="Text to synthesize")
            ref_wav = gr.Audio(sources="upload", type="filepath", label="Reference Audio File")
            exaggeration = gr.Slider(0.25, 2, step=.05, label="Exaggeration (Neutral = 0.5, extreme values can be unstable)", value=.5)

            run_btn = gr.Button("Generate", variant="primary")

        with gr.Column():
            audio_output = gr.Audio(label="Output Audio")

    run_btn.click(
        fn=generate,
        inputs=[
            text,
            ref_wav,
            exaggeration,
        ],
        outputs=audio_output,
    )

if __name__ == "__main__":
    demo.launch()