import fn
import gradio as gr

fn.load_model()

with gr.Blocks() as demo:
    audio = gr.Audio(sources="upload", type="filepath")
    model = gr.Dropdown(value='large-v3', choices=["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"])
    run_button = gr.Button(value='Run')
    prompt = gr.Textbox(label='prompt')
    set_button = gr.Button(value='Set Prompt')
    text_only = gr.Textbox(label='output')
    text_with_timestamps = gr.Textbox(label='timestamps')

    run_button.click(
        fn=fn.speech_to_text,
        inputs=[audio, model],
        outputs=[text_only, text_with_timestamps],
    )

    set_button.click(
        fn=fn.set_prompt,
        inputs=[prompt],
        outputs=[],
    )

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