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

# 1. Konfigurasi
MODEL_PATH = "zai-org/GLM-OCR"

# Paksa deteksi device yang aman
if torch.cuda.is_available():
    device = "cuda"
    dtype = torch.float16
else:
    device = "cpu"
    dtype = torch.float32

print(f"⚡ Loading model ke: {device} dengan tipe data: {dtype}")

# 2. Load Model & Processor
# Kita load model dulu baru processor biar manajemen memori lebih rapi di Space gratisan
try:
    model = AutoModelForImageTextToText.from_pretrained(
        MODEL_PATH,
        torch_dtype=dtype,
        device_map="auto",
        trust_remote_code=True,
        low_cpu_mem_usage=True # IQ move buat hemat RAM
    )
    
    # Kalo device cpu, pastikan model di float32
    if device == "cpu":
        model = model.float()

    processor = AutoProcessor.from_pretrained(
        MODEL_PATH, 
        trust_remote_code=True
    )
except Exception as e:
    print(f"❌ FATAL ERROR LOADING MODEL: {e}")
    raise e

# 3. Logic Inferensi
def run_ocr(image):
    if image is None:
        return "⚠️ Tolong upload gambar dokumennya dulu, Bang."

    # Format Prompt
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": "Text Recognition:"}
            ],
        }
    ]

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

        # Generate (Batasi token biar gak timeout di CPU)
        with torch.no_grad():
            generated_ids = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
            
        output_text = processor.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        return output_text

    except Exception as e:
        return f"Error saat proses OCR: {str(e)}"

# 4. Interface (Gradio)
with gr.Blocks(title="GLM-OCR by IQ 10M") as demo:
    gr.Markdown("# 👁️ GLM-OCR Demo (CPU/GPU Friendly)")
    gr.Markdown("Upload gambar dokumen. Sabar ya kalau pake CPU (Gratis), prosesnya agak lama.")
    
    with gr.Row():
        input_img = gr.Image(type="pil", label="Upload Disini", sources=["upload", "clipboard"])
        output_txt = gr.TextArea(label="Hasil OCR", interactive=False)
            
    btn = gr.Button("🔍 Scan Dokumen", variant="primary")
    btn.click(fn=run_ocr, inputs=input_img, outputs=output_txt)

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