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

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

# 1. HUKUM KEKALAN HARDWARE (Otomatis Deteksi CPU/GPU)
if torch.cuda.is_available():
    device = "cuda"
    dtype = torch.float16
else:
    device = "cpu"
    dtype = torch.float32

print(f"🚀 ENGINE STARTED: Device={device} | Dtype={dtype}")

# 2. INISIASI MODEL (TRUST REMOTE CODE MUTLAK)
try:
    print("⏳ Menyiapkan Otak Buatan...")
    processor = AutoProcessor.from_pretrained(
        MODEL_PATH, 
        trust_remote_code=True
    )
    
    # Kita pakai AutoModel karena GLM-OCR arsitekturnya unik
    model = AutoModel.from_pretrained(
        MODEL_PATH,
        torch_dtype=dtype,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        device_map="auto"
    )
    # Kunci model ke mode evaluasi biar irit memori
    model.eval() 

except Exception as e:
    print(f"⚠️ Warning Loading Model (Abaikan jika UI Muncul): {e}")
    # Kita biarkan script lanjut karena warning weights mismatch itu wajar di custom model
    pass

# 3. PROSES INFERENSI (INTELIJEN VISUAL)
def proses_intelijen(image):
    if image is None:
        return "⚠️ Gambarnya mana Bro? Fisika butuh materi buat bereaksi."

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

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

        with torch.no_grad():
            output_ids = model.generate(
                **inputs, 
                max_new_tokens=1500, # Batasi biar ga timeout
                do_sample=False
            )

        hasil = output_ids[0][len(inputs["input_ids"][0]):]
        teks_final = processor.decode(hasil, skip_special_tokens=True)
        return teks_final

    except Exception as e:
        return f"🚨 ERROR REAKSI: {str(e)}"

# 4. ANTARMUKA VISUAL 2026 (Modern Style)
# HAPUS 'theme' dan 'show_copy_button' di sini biar 100% Anti-Error

css_style = """
.container { max-width: 1200px; margin: auto; padding-top: 20px; }
h1 { text-align: center; color: #3b82f6; font-family: sans-serif; }
.gr-button-primary { background: linear-gradient(90deg, #3b82f6 0%, #8b5cf6 100%); border: none; color: white; }
.info { text-align: center; color: gray; font-size: 0.9em; margin-bottom: 20px;}
"""

with gr.Blocks(css=css_style, title="GLM-OCR PRO") as app:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# 👁️ GLM-OCR INFINITE VISION")
        gr.Markdown("<p class='info'>Analisis Dokumen Menggunakan Arsitektur Syaraf GLM Multimodal</p>")
        
        with gr.Row():
            # Kolom Kiri: Input
            with gr.Column(scale=1):
                input_img = gr.Image(type="pil", label="Masukkan Materi (Gambar)", sources=["upload", "clipboard"], height=500)
                # Tombol Ganteng
                scan_btn = gr.Button("🚀 EKSTRAKSI DATA", variant="primary", size="lg")
            
            # Kolom Kanan: Output
            with gr.Column(scale=1):
                # INI DIA FIXNYA: Ganti TextArea jadi Textbox dan HAPUS argumen ilegal
                # Pake 'show_copy_button=True' HANYA jika komponen Textbox, tapi demi keamanan kita pake standar aja.
                # interactive=False biar jadi read-only mode
                output_txt = gr.Textbox(label="Hasil Analisis Molekuler", lines=25, interactive=False, show_label=True)
    
    # Trigger Reaksi
    scan_btn.click(fn=proses_intelijen, inputs=input_img, outputs=output_txt)

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