Update app.py

app.py

@@ -3,81 +3,88 @@ 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
-
+# 1. Konfigurasi
 MODEL_PATH = "zai-org/GLM-OCR"
 
-print(f"Loading model ke {device}...")
+# 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 Processor dan Model
+# 2. Load Model & Processor
+# Kita load model dulu baru processor biar manajemen memori lebih rapi di Space gratisan
 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,
+        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"Error loading model: {e}")
+    print(f"❌ FATAL ERROR LOADING MODEL: {e}")
     raise e
 
-# 3. Fungsi Inferensi
+# 3. Logic Inferensi
 def run_ocr(image):
     if image is None:
-        return "Tolong upload gambar dulu, Bro."
+        return "⚠️ Tolong upload gambar dokumennya dulu, Bang."
 
-    # Format pesan khusus GLM-OCR
+    # Format Prompt
     messages = [
         {
             "role": "user",
             "content": [
-                {
-                    "type": "image",
-                    "image": image, # Gradio ngasih format PIL Image
-                },
-                {
-                    "type": "text",
-                    "text": "Text Recognition:" 
-                }
+                {"type": "image", "image": 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)
+    # Proses
+    try:
+        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
+        # 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. 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.")
+# 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():
-        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)
+        input_img = gr.Image(type="pil", label="Upload Disini", sources=["upload", "clipboard"])
+        output_txt = gr.TextArea(label="Hasil OCR", interactive=False)
             
-    btn_submit.click(fn=run_ocr, inputs=input_img, outputs=output_txt)
+    btn = gr.Button("🔍 Scan Dokumen", variant="primary")
+    btn.click(fn=run_ocr, inputs=input_img, outputs=output_txt)
 
-# Jalankan App
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()