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

# Initialize model on CPU with float16
model = pipeline(
    "automatic-speech-recognition",
    model="Aekanun/whisper-small-hi",
    device="cuda",  # เปลี่ยนจาก cpu เป็น cuda
    torch_dtype=torch.float16
)

@spaces.GPU
def transcribe_speech(audio):
    """Speech transcription with GPU support"""
    try:
        if audio is None:
            return "กรุณาบันทึกเสียงก่อน"
            
        # ไม่ต้องย้ายโมเดลเพราะอยู่บน GPU อยู่แล้ว
        with torch.amp.autocast('cuda'):
            # Process audio with chunk_length_s
            result = model(
                audio, 
                batch_size=1,
                chunk_length_s=30
            )
            
            # Get text result
            text = result["text"] if isinstance(result, dict) else result
            
        return text
        
    except Exception as e:
        return f"เกิดข้อผิดพลาด: {str(e)}"

# Create Gradio interface
demo = gr.Interface(
    fn=transcribe_speech,
    inputs=gr.Audio(type="filepath"),
    outputs=gr.Textbox(label="ข้อความ"),
    title="Thai Speech Transcription",
    description="บันทึกเสียงเพื่อแปลงเป็นข้อความภาษาไทย",
)

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