import gradio as gr
import soundfile as sf
import numpy as np
import tempfile
import torchaudio
from transformers import AutoModel

# Load ASR Model
def load_model():
    return AutoModel.from_pretrained("ai4bharat/indic-conformer-600m-multilingual", trust_remote_code=True)

model = load_model()

def process_audio(audio, language, decoding_method):
    if isinstance(audio, tuple):  # Recorded audio
        sample_rate, data = audio
        temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
        sf.write(temp_wav.name, data, sample_rate)
        audio_path = temp_wav.name
    else:  # Uploaded file
        audio_path = audio
    
    # Load and resample audio
    wav, sr = torchaudio.load(audio_path)
    target_sample_rate = 16000
    if sr != target_sample_rate:
        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sample_rate)
        wav = resampler(wav)
    
    # Perform ASR with selected decoding method
    transcription = model(wav, language, decoding_method)
    
    return transcription

iface = gr.Interface(
    fn=process_audio,
    inputs=[
        gr.Audio(source="microphone", type="numpy"),
        gr.Audio(source="upload"),
        gr.Dropdown(["hi", "ta", "bn", "mr", "te", "gu", "kn", "ml", "pa", "ur"], label="Select Language"),
        gr.Radio(["ctc", "rnnt"], label="Decoding Method")
    ],
    outputs="text",
    title="Multilingual ASR with Indic-Conformer",
    description="Record or upload an audio file, select a language and decoding method, and transcribe it using the AI4Bharat Indic-Conformer model."
)

if __name__ == "__main__":
    iface.launch()