import os

import gradio as gr
import spaces
import torch
from qwen_asr import Qwen3ASRModel


MODEL_ID = os.getenv(
    "MODEL_ID",
    "tryorato/orato-asr-hindi-v1",
)

print(f"Loading {MODEL_ID}...", flush=True)

# ZeroGPU provides CUDA emulation during startup.
# The real GPU is allocated when the decorated function runs.
model = Qwen3ASRModel.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    device_map="cuda:0",
    max_inference_batch_size=1,
    max_new_tokens=512,
)

print("Model loaded successfully.", flush=True)


@spaces.GPU(duration=120)
def transcribe(audio_path):
    if not audio_path:
        return "Please record or upload an audio file."

    try:
        with torch.inference_mode():
            results = model.transcribe(
                audio=audio_path,
                language="Hindi",
            )

        if not results:
            return "No transcription was generated."

        text = (results[0].text or "").strip()
        return text or "No speech detected."

    except Exception as exc:
        print(f"Transcription error: {exc}", flush=True)
        return f"Transcription error: {exc}"


demo = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(
        sources=["microphone", "upload"],
        type="filepath",
        label="Record or upload Hindi/Hinglish audio",
    ),
    outputs=gr.Textbox(
        label="ORATO transcription",
        lines=6,
        show_copy_button=True,
    ),
    title="ORATO Hindi ASR",
    description=(
        "Hindi and Hinglish speech recognition using "
        "tryorato/orato-asr-hindi-v1."
    ),
    flagging_mode="never",
)

demo.queue(max_size=10)

if __name__ == "__main__":
    demo.launch(ssr_mode=False)