import gradio as gr
import torch
import numpy as np
import librosa
from transformers import AutoProcessor, Gemma4ForConditionalGeneration

try:
    import spaces
    HAS_ZEROGPU = True
except ImportError:
    HAS_ZEROGPU = False

PROCESSOR_ID = "google/gemma-4-E2B-it"
MODEL_ID = "himalaya-ai/himalaya-gemma-4-e2b-it"
LANGUAGE = "Nepali"
TARGET_SR = 16000

print(f"Loading processor from {PROCESSOR_ID} ...")
processor = AutoProcessor.from_pretrained(PROCESSOR_ID, trust_remote_code=True)

print(f"Loading model weights from {MODEL_ID} ...")
model = Gemma4ForConditionalGeneration.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    trust_remote_code=True,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()


def _transcribe_core(audio_array, sr):
    if audio_array.ndim > 1:
        audio_array = np.mean(audio_array, axis=1)
    audio_array = audio_array.astype(np.float32)
    if audio_array.max() > 1.0 or audio_array.min() < -1.0:
        audio_array = audio_array / 32768.0
    if sr != TARGET_SR:
        audio_array = librosa.resample(audio_array, orig_sr=sr, target_sr=TARGET_SR)

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": (
                        f"Transcribe the following speech segment in {LANGUAGE} "
                        f"into {LANGUAGE} text. "
                        "Follow these specific instructions for formatting the answer: "
                        "* Only output the transcription, with no newlines. "
                        "* When transcribing numbers, write the digits."
                    ),
                },
                {"type": "audio", "audio": audio_array},
            ],
        }
    ]

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=False,
        )

    input_len = inputs["input_ids"].shape[1]
    decoded = processor.decode(
        outputs[0][input_len:], skip_special_tokens=True
    ).strip()
    return decoded


if HAS_ZEROGPU:
    @spaces.GPU(duration=60)
    def transcribe_audio(audio):
        if audio is None:
            return "⚠️ No audio provided."
        sr, data = audio
        return _transcribe_core(data, sr)
else:
    def transcribe_audio(audio):
        if audio is None:
            return "⚠️ No audio provided."
        sr, data = audio
        return _transcribe_core(data, sr)


with gr.Blocks(title="Himalaya Gemma — Nepali ASR") as demo:
    gr.Markdown(
        "# 🏔️ Himalaya Gemma — Nepali ASR Demo\n"
        f"Model: `{MODEL_ID}`\n\n"
        "Record your voice or upload a `.wav`, then click **Transcribe**."
    )
    audio_in = gr.Audio(
        sources=["microphone", "upload"],
        type="numpy",
        label="Record or upload Nepali audio",
    )
    transcribe_btn = gr.Button("📝 Transcribe", variant="primary")
    output_text = gr.Textbox(label="Transcription", lines=4)

    transcribe_btn.click(
        fn=transcribe_audio,
        inputs=audio_in,
        outputs=output_text,
    )

    gr.Examples(
        examples=[],
        inputs=audio_in,
        label="Sample audio (add your own .wav files here if you want)",
    )

if __name__ == "__main__":
    demo.queue().launch()