import torch
import gradio as gr
from model import ECAPA_gender

model = ECAPA_gender.from_pretrained("Beijuka/voice-gender-classifier")
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def predict_gender_confidence(audio_file):
    if audio_file is None:
        return "No audio provided"
    
    try:
        # Load audio
        audio = model.load_audio(audio_file.name if hasattr(audio_file, "name") else audio_file)
        audio = audio.to(device)
        
        # Forward pass
        with torch.no_grad():
            logits = model.forward(audio)
            probs = torch.softmax(logits, dim=1).cpu().numpy()[0] 
            pred_idx = logits.argmax(dim=1).item()
            gender_pred = model.pred2gender[pred_idx].capitalize()
            confidence = probs[pred_idx] * 100 
        
        return f"{gender_pred} — {confidence:.1f}% confidence"
    
    except Exception as e:
        return f"Error: {e}"

iface = gr.Interface(
    fn=predict_gender_confidence,
    inputs=gr.Audio(type="filepath", label="Upload audio file", sources=["upload"]),
    outputs=gr.Textbox(label="Predicted Gender with Confidence"),
    title="Voice Gender Classifier",
    description="Upload an audio file and the model predicts speaker gender with confidence.",
    allow_flagging="never"
)

iface.launch(share=True)