app.py

import os, io, traceback
import gradio as gr
import fitz  # PyMuPDF
from PIL import Image
import spaces  # ل ZeroGPU

# تعطيل torchvision داخل Transformers (لمنع AutoVideoProcessor)
os.environ["TRANSFORMERS_NO_TORCHVISION"] = "1"

# إعدادات النماذج
BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen2-VL-2B-Instruct")
HOROOF_ADAPTER = os.environ.get("HOROOF_MODEL", "NaserNajeh/Horoof")

_model = None
_processor = None  # AutoProcessor من الـBase Model

def load_model_merged():
    """
    1) نحمل Base Qwen2-VL-2B-Instruct على GPU (fp16).
    2) نركب لورا Horoof (PEFT) على الـBase.
    3) ندمج الأوزان merge_and_unload لتفادي أي تبعية لـ bitsandbytes.
    """
    global _model, _processor
    if _model is not None:
        return
    try:
        import torch
        from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
        from peft import PeftModel

        if not torch.cuda.is_available():
            raise AssertionError("هذه النسخة تتطلب GPU (CUDA) مفعّل على الـSpace.")

        # التحميل
        _processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=False)

        base = Qwen2VLForConditionalGeneration.from_pretrained(
            BASE_MODEL, torch_dtype=torch.float16
        ).to("cuda")

        # ركب لورا Horoof ثم دمجها
        peft_model = PeftModel.from_pretrained(base, HOROOF_ADAPTER)
        _model = peft_model.merge_and_unload()  # نحصل على نموذج مدمج بلا تبعيات إضافية
        _model.to("cuda")

    except Exception as e:
        raise RuntimeError(f"تعذّر تحميل النموذج: {e}")

def pdf_to_images(pdf_bytes: bytes, dpi: int = 220, max_pages: int = 0):
    """تحويل PDF إلى قائمة صور PIL."""
    pages_imgs = []
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    total = doc.page_count
    n_pages = total if (not max_pages or max_pages <= 0) else min(max_pages, total)
    for i in range(n_pages):
        page = doc.load_page(i)
        pix = page.get_pixmap(dpi=dpi, alpha=False)
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        pages_imgs.append((i + 1, img))
    doc.close()
    return pages_imgs

def ocr_page_gpu(pil_img: Image.Image, max_new_tokens: int = 1200) -> str:
    """OCR لصفحة واحدة باستخدام Qwen2-VL المدموج مع لورا Horoof."""
    load_model_merged()
    import torch

    # نبني رسالة محادثة متوافقة مع Qwen2-VL
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": pil_img},
                {"type": "text", "text": "اقرأ النص العربي في الصورة كما هو دون أي تعديل أو تفسير."},
            ],
        }
    ]

    # تحويل الرسائل إلى Prompt داخلي
    prompt = _processor.apply_chat_template(messages, add_generation_prompt=True)

    # تجهيز المدخلات (نمرر النص + الصورة)
    inputs = _processor(text=[prompt], images=[pil_img], return_tensors="pt").to("cuda")

    # توليد
    with torch.inference_mode():
        output_ids = _model.generate(**inputs, max_new_tokens=max_new_tokens)
    text = _processor.batch_decode(output_ids, skip_special_tokens=True)[0]
    return (text or "").strip()

@spaces.GPU  # مهم ل ZeroGPU: يحجز GPU عند الاستدعاء
def ocr_pdf(pdf_file, dpi, limit_pages):
    """الدالة الرئيسة التي يستدعيها Gradio."""
    if pdf_file is None:
        return "لم يتم رفع ملف."
    try:
        pdf_bytes = pdf_file.read() if hasattr(pdf_file, "read") else pdf_file
        limit = int(limit_pages) if limit_pages else 1  # صفحة واحدة افتراضيًا للاختبار
        pages = pdf_to_images(pdf_bytes, dpi=int(dpi), max_pages=limit)
        if not pages:
            return "لا توجد صفحات."
        out = []
        for idx, img in pages:
            txt = ocr_page_gpu(img)
            out.append(f"--- صفحة {idx} ---\n{txt}")
        return "\n\n".join(out)
    except AssertionError as ae:
        return f"⚠️ {ae}"
    except Exception as e:
        traceback.print_exc()
        return f"حدث خطأ: {repr(e)}"

with gr.Blocks(title="Horoof OCR (ZeroGPU)") as demo:
    gr.Markdown("### Horoof OCR على ZeroGPU — Qwen2-VL + LoRA (مُدمج أثناء التشغيل، بدون bitsandbytes).")
    pdf_in = gr.File(label="ارفع ملف PDF", file_types=[".pdf"], type="binary")
    dpi = gr.Slider(150, 300, value=220, step=10, label="دقّة التحويل (DPI)")
    limit_pages = gr.Number(value=1, precision=0, label="عدد الصفحات (للاختبار؛ زِد لاحقًا)")
    run_btn = gr.Button("بدء التحويل")
    out = gr.Textbox(label="النص المستخرج", lines=24)

    demo.queue()
    run_btn.click(fn=ocr_pdf, inputs=[pdf_in, dpi, limit_pages], outputs=out, api_name="ocr_pdf")

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