import gradio as gr

from fbeeper_hubert import HubertBeeper

BEEPER = HubertBeeper()


def greet(audio):
    global BEEPER
    sample_rate, wav = audio

    # Gradio loads sound files as ints, the model expects normalized floats
    wav = wav / 32768

    # TODO: output sample_rate doesn't have to match input sample rate
    text, result_wav = BEEPER.f_beep_waveform(wav, sample_rate)

    result_wav = (result_wav * 32768).astype("int16")
    return (text, (sample_rate, result_wav))


if __name__ == '__main__':
    iface = gr.Interface(
        fn=greet,
        inputs="audio",
        outputs=[
            gr.outputs.Textbox(label="ASR output"),
            gr.outputs.Audio(label="Censored speech"),
        ],
        examples=[["data/speech-short-16k.wav"], ["data/sample.wav"]],
        server_port=7860,
    )
    iface.launch()