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

# Utilizamos los tres modelos entrenados
pipe_model_1 = pipeline("automatic-speech-recognition", model="IABDs8a/AfinandoElEntrenamiento")
pipe_model_2 = pipeline("automatic-speech-recognition", model="IABDs8a/whisper-base-full")
pipe_model_3 = pipeline("automatic-speech-recognition", model="IABDs8a/whisper-medium-asfe")
pipe_model_4 = pipeline("automatic-speech-recognition", model="IABDs8a/whisper-tiny-top3")

def transcribe(audio, model_choice):
    inicio = time.time()

    if model_choice == "AfinandoElEntrenamiento":
        pipe = pipe_model_1
    elif model_choice == "Whisper Base Full":
        pipe = pipe_model_2
    elif model_choice == "Whisper Medium Asfe":
        pipe = pipe_model_3
    else:
        pipe = pipe_model_4

    # Leer el archivo de audio
    y, sr = librosa.load(audio, sr=16000)
    
    # Convertir a mono si es necesario
    if y.ndim > 1:
        y = librosa.to_mono(y)
    
    # Pasamos el array de muestras a tipo NumPy de 32 bits
    y = y.astype(np.float32)
    
    # Normalizar el audio
    y /= np.max(np.abs(y))
    
    # Realizar la transcripción
    result = pipe({"sampling_rate": sr, "raw": y})
    fin = time.time()

    return result["text"], fin - inicio

# Interfaz de Gradio
demo = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(type="filepath", label="Sube un archivo de audio o graba desde el micrófono"),
        gr.Dropdown(choices=["AfinandoElEntrenamiento", "Whisper Base Full","Whisper Medium Asfe","Whisper Tiny Top 3"], label="Selecciona el modelo", value="Whisper Base Full")
    ],
    outputs=[
        gr.Text(label="Salida"),
        gr.Number(label="Tiempo")
    ],
    title="Transcripción de Audio con LARA",
    description="Sube un archivo de audio o graba desde el micrófono para obtener su transcripción utilizando los modelos Whisper entrenados.",
)

demo.launch()