import os
import torch
import gradio as gr
from transformers import AutoModel
import numpy as np

MODEL_NAME = "ARTPARK-IISc/SraVaani"

print("Loading model, this may take a few minutes...")
model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True, token=os.getenv("HF_TOKEN"))
model.eval()
print("Model loaded successfully.")


def transcribe(audio_input):
    if audio_input is None:
        return ""

    sr, audio = audio_input

    if audio.ndim == 2:
        audio = np.mean(audio, axis=1)

    audio = audio.astype(np.float32)
    audio = audio / (np.max(np.abs(audio)) + 1e-9)

    import tempfile, soundfile as sf
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        sf.write(f.name, audio, sr)
        results = model.transcribe([f.name])

    return results[0] if results else ""


demo = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(
        sources=["microphone", "upload"],
        type="numpy",
        label="Record or upload WAV audio"
    ),
    outputs=gr.Textbox(label="Transcription"),
    title="SraVaani Multilingual ASR",
    description="Upload a WAV file and get the multilingual ASR transcription."
)

demo.launch(share=True)