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

# --- KONFIGURASI ---
MODEL_PATH = "zai-org/GLM-OCR"

# 1. Deteksi Hardware (Biar gak maksain CPU nangis)
if torch.cuda.is_available():
    device = "cuda"
    dtype = torch.float16
else:
    device = "cpu"
    # Pake float32 aja buat CPU biar aman, walau lambat
    dtype = torch.float32

print(f"🚀 Mulai System: Device={device} | Dtype={dtype}")

# 2. LOAD MODEL (LOGIKA BARU)
# GLM-OCR butuh arsitektur khusus. Kita pake 'AutoModel' bukan 'AutoModelForImageTextToText' 
# karena 'AutoModel' lebih fleksibel buat nerima arsitektur custom via remote code.
try:
    print("⏳ Loading Processor...")
    processor = AutoProcessor.from_pretrained(
        MODEL_PATH, 
        trust_remote_code=True
    )

    print("⏳ Loading Model...")
    # Pake AutoModel biasa, dia bakal baca config.json dan narik class GLMOCRModel otomatis
    model = AutoModel.from_pretrained(
        MODEL_PATH,
        torch_dtype=dtype,
        trust_remote_code=True,
        # Opsi memori buat CPU (PENTING BUAT SPACE GRATISAN)
        device_map="auto" if device == "cuda" else "cpu", 
        low_cpu_mem_usage=True 
    )
    
    # Kalau CPU, model di eval mode
    model = model.eval()

except Exception as e:
    print(f"❌ KEGAGALAN SISTEM: {e}")
    # Pesan ini biar muncul di log kalau error
    raise ValueError(f"Gagal Load Model GLM-OCR. Detail: {e}")

# 3. Fungsi Kerja (Inferensi)
def proses_gambar(image):
    if image is None:
        return "⚠️ Woi bro, upload gambarnya dulu dong."

    # Prompt standar GLM-OCR
    pesan_user = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": "Text Recognition:"}
            ]
        }
    ]

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

        # Generasi Teks
        with torch.no_grad():
            output_ids = model.generate(
                **inputs, 
                max_new_tokens=1024, # Biar ga kepanjangan render di CPU
                do_sample=False      # Greedy search biar stabil
            )

        # Ambil hasilnya doang (potong input dari output)
        hasil_asli = output_ids[0][len(inputs["input_ids"][0]):]
        teks_final = processor.decode(hasil_asli, skip_special_tokens=True)
        
        return teks_final

    except Exception as e:
        return f"😭 Error Pas Baca Gambar: {str(e)}"

# 4. Tampilan Web (UI)
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🔍 GLM-OCR Detector (Versi IQ Tinggi)")
    gr.Markdown("Model OCR canggih buat baca dokumen rumit. Upload aja langsung scan.")
    
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="Gambar Dokumen/Teks", sources=["upload", "clipboard"])
            scan_btn = gr.Button("🚀 MULAI SCAN", variant="primary")
        
        with gr.Column():
            txt_output = gr.TextArea(label="Hasil Bacaan", show_copy_button=True, lines=20)
            
    scan_btn.click(fn=proses_gambar, inputs=img_input, outputs=txt_output)

# 5. Eksekusi
if __name__ == "__main__":
    print("✅ App siap dijalankan...")
    app.launch()