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

# 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-tiny-top3")

def transcribe(audio, model_choice):
    if model_choice == "AfinandoElEntrenamiento":
        pipe = pipe_model_1
    elif model_choice == "Whisper Base Full":
        pipe = pipe_model_2
    else:
        pipe = pipe_model_3

    # 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})
    return result["text"]

# Ruta de la imagen del logo
logo_path = "img/logoLara.png"

# Estilos CSS personalizados
custom_css = """
<style>
body {
    background-color: #c5e6ed; /* Color de fondo principal */
    color: #040404; /* Color del texto principal */
}
.gradio-container {
    border-radius: 10px;
    border: 2px solid #477b84; /* Color de borde */
    padding: 20px;
    background-color: #8cd3dc; /* Color de fondo del contenedor */
}
input[type="file"] {
    border-radius: 10px;
    border: 2px solid #56646c; /* Color del borde del input */
}
button {
    background-color: #477b84; /* Color de fondo de los botones */
    color: #ffffff; /* Color del texto de los botones */
    border-radius: 5px;
    border: 2px solid #040404; /* Color del borde de los botones */
    padding: 10px 20px;
}
button:hover {
    background-color: #6dc6d6; /* Color de fondo de los botones al pasar el ratón */
}
#logo {
    text-align: center;
    margin-bottom: 20px;
}
</style>
"""

# HTML para el logo
logo_html = f"<div id='logo'><img src='{logo_path}' alt='Logo' width='300'></div>"

# 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 Tiny Top 3"], label="Selecciona el modelo", value="Whisper Base Full")
    ],
    outputs=gr.Textbox(label="Texto Generado"),

    title="Transcripci&oacute;n de Audio con LARA",
    description="Sube un archivo de audio o graba desde el micr&oacute;fono para obtener su transcripci&oacute;n utilizando los modelos Whisper entrenados.",
    article=custom_css,
    
)

demo.launch(share=True)