import gradio as gr
import torch
from transformers import pipeline
import numpy as np
import time

pipe_vol_g = pipeline("automatic-speech-recognition", model="IABDs8a/whisper-equipo3-g")
pipe_small = pipeline("automatic-speech-recognition", model="aitor-medrano/whisper-small-lara")
pipe_hombres = pipeline("automatic-speech-recognition", model="IABDs8a/lara-hombres-equipo-3")

def greet(grabacion):
    inicio = time.time()

    sr, y = grabacion
    # Pasamos el array de muestras a tipo NumPy de 32 bits
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))

    inicio_g = time.time()
    result_g = "g:" + pipe_vol_g({"sampling_rate": sr, "raw": y})["text"]
    fin_g = time.time()

    inicio_small = time.time()
    result_small = "small:" + pipe_small({"sampling_rate": sr, "raw": y})["text"]
    fin_small = time.time()

    inicio_hombres = time.time()
    result_hombres = "hombres:" + pipe_hombres({"sampling_rate": sr, "raw": y})["text"]
    fin_hombres = time.time()

    fin = time.time()

    return result_g, fin_g - inicio_g, result_small, fin_small - inicio_small, result_hombres, fin_hombres - inicio_hombres, fin - inicio
    #return result_base, result_small, fin - inicio

demo = gr.Interface(fn=greet,
        inputs=[
                gr.Audio(),
        ],
        outputs=[
            gr.Text(label="Salida (Voluntaria G)"),
            gr.Number(label="Tiempo (Voluntaria G)"),
            gr.Text(label="Salida (Small)"),
            gr.Number(label="Tiempo (Small)"),
            gr.Text(label="Salida (Hombres)"),
            gr.Number(label="Tiempo (Hombres)"),
            gr.Number(label="Tiempo total")
        ])
demo.launch()