import os
import torch
import gradio as gr
from transformers import pipeline

# Model configuration
MODEL_NAME = "chuuhtetnaing/whisper-large-v3-myanmar"

# Automatically use GPU if available on HF Space, fallback to CPU
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

print(f"Loading pipeline on {device}...")

# Load model pipeline ONCE when the Space starts up
pipe = pipeline(
    "automatic-speech-recognition",
    model=MODEL_NAME,
    torch_dtype=torch_dtype,
    device=device,
    chunk_length_s=30,  # Splice long audio into 30s chunks automatically
)

def transcribe_audio(audio_file):
    if audio_file is None:
        return "Please upload or record an audio snippet."

    # Run transcription
    outputs = pipe(
        audio_file,
        generate_kwargs={
            "language": "burmese", 
            "task": "transcribe"
        },
        return_timestamps=True,
    )
    
    return outputs["text"]

# Define Gradio Interface
demo = gr.Interface(
    fn=transcribe_audio,
    inputs=gr.Audio(
        type="filepath", 
        label="Record or Upload Audio",
        sources=["microphone", "upload"]
    ),
    outputs=gr.Textbox(
        label="Myanmar Transcription Result", 
        lines=6, 
        show_copy_button=True
    ),
    title="🇲🇲 Myanmar Speech-to-Text (Whisper Large v3)",
    description="Upload an audio file (`.wav`, `.mp3`, `.m4a`) or record directly from your microphone to get Myanmar transcriptions.",
)

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