import gradio as gr
from transformers import pipeline
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
import os



# Disponibilidad de GPU
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# ID del modelo 
model_id = "openai/whisper-large-v3-turbo"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)


pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    chunk_length_s=30,
    batch_size=16,  # Para ajustar dependiendo de la capacidad del sistema
    torch_dtype=torch_dtype,
    device=device,
)


def transcribe(audio):
    if audio is None or not os.path.exists(audio):
        return "Input de audio inválido. Por favor, sube un archivo de audio válido."


    try:
        result = pipe(audio)
        return result["text"]
    except Exception as e:
        return f"Ocurrió un error durante la transcripción: {e}"


iface = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(type="filepath", label="Subir archivo de audio"),
    outputs=gr.Textbox(label="Resultado de la transcripción"),
    title="Transcripción por reconocimiento de audio",
    description="Sube tu audio para transcribir a texto.\n\nEV 2024",
    examples=[
        ["example_audio_1.wav"],  # Ojo añadir el camino de los archivo
        ["example_audio_2.wav"],
    ],
    live=True,  
    theme="default",  
)

iface.launch()