Update app.py

app.py

@@ -701,23 +701,21 @@
 
 # app.py — TRUST OCR DEMO (Streamlit) — personal-recognition-only
 
+# app.py — TRUST OCR DEMO (Streamlit) with personal recognition model, safe dirs, eager attention, lazy order
+
 import os
 import io
 import tempfile
 from typing import List
 
-import numpy as np
-import cv2
-from PIL import Image
-import pypdfium2
-import pytesseract
-
-# -------------------- Safe, writable dirs & config (BEFORE importing streamlit) --------------------
-# Put everything under /tmp (world-writable on HF Spaces)
+# -------------------- Safe, writable dirs & envs BEFORE any ML imports --------------------
+# همه‌ی مسیرهای نوشتنی را به /tmp ببریم (در HF Spaces world-writable)
 os.environ.setdefault("HOME", "/tmp")
 os.environ.setdefault("STREAMLIT_CONFIG_DIR", "/tmp/.streamlit")
 os.environ.setdefault("STREAMLIT_RUNTIME_DIR", "/tmp/.streamlit")
 os.environ.setdefault("HF_HOME", "/tmp/hf_home")
+# جلوگیری از sdpa برای surya_order
+os.environ.setdefault("TRANSFORMERS_ATTENTION_IMPLEMENTATION", "eager")
 
 for d in (os.environ["STREAMLIT_CONFIG_DIR"], os.environ["STREAMLIT_RUNTIME_DIR"], os.environ["HF_HOME"]):
     os.makedirs(d, exist_ok=True)
@@ -736,16 +734,22 @@ if not os.path.exists(conf_path):
             "gatherUsageStats = false\n"
         )
 
-# HF auth (for private repos) — optional, picked up by transformers/surya automatically
+# HF auth (برای ریپوی خصوصی) — اختیاری؛ transformers/surya خودش از env می‌خواند
 HF_TOKEN = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_TOKEN")
 if HF_TOKEN:
     os.environ["HUGGINGFACE_HUB_TOKEN"] = HF_TOKEN
     try:
         from huggingface_hub import login
-        login(token=HF_TOKEN, add_to_git_credential=True)
+        login(token=HF_TOKEN, add_to_git_credential=False)
     except Exception:
         pass
 
+# -------------------- Light deps (safe to import now) --------------------
+import numpy as np
+import cv2
+from PIL import Image
+import pypdfium2
+import pytesseract
 import streamlit as st
 
 # -------------------- Surya imports (gated) --------------------
@@ -755,7 +759,7 @@ try:
 except Exception:
     DET_AVAILABLE = False
 
-from surya.layout import batch_layout_detection  # will be gated via DET_AVAILABLE
+from surya.layout import batch_layout_detection  # we'll gate usage via DET_AVAILABLE
 
 # Detection model loaders: try newer segformer, fall back to older
 try:
@@ -887,6 +891,7 @@ def load_layout_cached():
 
 @st.cache_resource(show_spinner=True)
 def load_order_cached():
+    # Order model گاهی با attention=sdpa کرش می‌کند؛ ما قبلاً eager را با env ست کرده‌ایم.
     return load_order_model(checkpoint="vikp/surya_order"), load_order_processor(checkpoint="vikp/surya_order")
 
 
@@ -894,7 +899,7 @@ def load_order_cached():
 # Choose ONE source:
 #  (A) Local folder path via env TRUSTOCR_PATH (e.g., /app/models/TrustOCR)
 #  (B) Private HF repo via env TRUSTOCR_REPO (e.g., MohammadReza-Halakoo/TrustOCR) + HUGGINGFACE_HUB_TOKEN
-PERSONAL_MODEL_PATH = os.environ.get("TRUSTOCR_PATH")  # local directory containing save_pretrained files
+PERSONAL_MODEL_PATH = os.environ.get("TRUSTOCR_PATH")  # local directory with save_pretrained files
 PERSONAL_HF_REPO = os.environ.get("TRUSTOCR_REPO")     # private repo id on HF Hub
 
 @st.cache_resource(show_spinner=True)
@@ -910,20 +915,20 @@ def load_rec_personal():
         return m, p
     else:
         raise RuntimeError(
-            "No personal recognition model configured. "
-            "Set TRUSTOCR_PATH to a local folder OR TRUSTOCR_REPO to a private HF repo."
+            "مدل شخصی تنظیم نشده است. یکی را ست کن: TRUSTOCR_PATH (پوشه محلی) یا TRUSTOCR_REPO (ریپوی خصوصی HF)."
         )
 
-# Load everything
+# Load detection/layout immediately (سبک‌ترند). Order را lazy-load می‌کنیم.
 if DET_AVAILABLE:
     det_model, det_processor = load_det_cached()
     layout_model, layout_processor = load_layout_cached()
-    order_model, order_processor = load_order_cached()
 else:
-    det_model = det_processor = layout_model = layout_processor = order_model = order_processor = None
+    det_model = det_processor = layout_model = layout_processor = None
 
 # Always require personal recognition model (no fallback)
 rec_model, rec_processor = load_rec_personal()
+order_model = None
+order_processor = None  # will lazy-load on button click
 
 
 # ===================== High-level Ops =====================
@@ -960,6 +965,11 @@ def layout_detection(pil_img: Image.Image):
 
 
 def order_detection(pil_img: Image.Image):
+    # Lazy-load order model on demand
+    global order_model, order_processor
+    if order_model is None or order_processor is None:
+        order_model, order_processor = load_order_cached()
+
     _, layout_pred = layout_detection(pil_img)
     bboxes = [l.bbox for l in layout_pred.bboxes]
     pred: OrderResult = batch_ordering([pil_img], [bboxes], order_model, order_processor)[0]
@@ -974,10 +984,10 @@ def ocr_page(pil_img: Image.Image, langs: List[str]):
     langs = list(langs) if langs else ["Persian"]
     replace_lang_with_code(langs)  # in-place
 
-    if det_model and det_processor and rec_model and rec_processor:
+    if DET_AVAILABLE and det_model and det_processor and rec_model and rec_processor:
         img_pred: OCRResult = run_ocr([pil_img], [langs], det_model, det_processor, rec_model, rec_processor)[0]
     else:
-        # detection may be missing in some environments; use recognition-only flow
+        # detection may be missing in some environments; use recognition-only path
         img_pred: OCRResult = run_ocr([pil_img], [langs], rec_model=rec_model, rec_processor=rec_processor)[0]
 
     bboxes = [l.bbox for l in img_pred.text_lines]
@@ -1039,6 +1049,8 @@ with col1:
         try:
             order_img, order_pred = order_detection(pil_image)
             st.image(order_img, caption="ترتیب خوانش (Reading Order)", use_column_width=True)
+        except KeyError:
+            st.error("Reading Order فعلاً با نسخه‌های فعلی سازگار نیست (attention=sdpa). نسخه‌ها را طبق requirements پین کن.")
         except Exception as e:
             st.error(f"خطا در ترتیب خوانش: {e}")