import gradio as gr
from transformers import AutoProcessor, AutoModelForImageTextToText
import torch
from PIL import Image

# 1. Konfigurasi Model
# Pake CPU kalau gak ada GPU, pake cuda kalau ada
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32

MODEL_PATH = "zai-org/GLM-OCR"

print(f"Loading model ke {device}...")

# 2. Load Processor dan Model
try:
    processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        MODEL_PATH,
        torch_dtype=dtype,
        device_map="auto",
        trust_remote_code=True
    )
except Exception as e:
    print(f"Error loading model: {e}")
    raise e

# 3. Fungsi Inferensi
def run_ocr(image):
    if image is None:
        return "Tolong upload gambar dulu, Bro."

    # Format pesan khusus GLM-OCR
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image, # Gradio ngasih format PIL Image
                },
                {
                    "type": "text",
                    "text": "Text Recognition:" 
                }
            ],
        }
    ]

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

    # Generate (OCR)
    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=2048) # Token bisa dinaikin kalo dokumen panjang
        
    # Decode hasil jadi teks
    output_text = processor.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    
    return output_text

# 4. Bikin Tampilan UI (Gradio)
with gr.Blocks() as demo:
    gr.Markdown("# 👁️ GLM-OCR Demo")
    gr.Markdown("Upload gambar dokumen, surat, atau teks tulisan tangan. Model ini jago baca layout.")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Upload Gambar")
            btn_submit = gr.Button("Baca Teks (OCR)", variant="primary")
        
        with gr.Column():
            output_txt = gr.Textbox(label="Hasil Text / Markdown", lines=20)
            
    btn_submit.click(fn=run_ocr, inputs=input_img, outputs=output_txt)

# Jalankan App
demo.launch()