Update app.py

app.py

@@ -5,24 +5,36 @@ from PIL import Image
 from pdf2image import convert_from_path
 import os
 
-# 1. Hardware & Model Setup
-model_id = "numind/NuMarkdown-8B-Thinking"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# --- CONFIGURATION ---
+MODEL_ID = "numind/NuMarkdown-8B-Thinking"
 
-# 8B models perform best in bfloat16 on Ampere+ GPUs, or float16 on older ones
-dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+# 1. Hardware Detection
+# We explicitly check for CUDA. If not found, we default to CPU with float32.
+if torch.cuda.is_available():
+    device = "cuda"
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    print(f"🚀 Running on GPU ({torch.cuda.get_device_name(0)}) with {dtype}")
+else:
+    device = "cpu"
+    dtype = torch.float32  # CPU must use float32 to avoid "Layer not implemented" errors
+    print("⚠️ Running on CPU. This will be SLOW (2-5 mins per page). Requires ~32GB RAM.")
 
-print(f"Loading {model_id} on {device} with {dtype}...")
-
-# 2. Loading Model & Processor
-# We use device_map="auto" to handle VRAM distribution automatically (Crucial for 8B models)
-processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModelForImageTextToText.from_pretrained(
-    model_id, 
-    torch_dtype=dtype,
-    trust_remote_code=True,
-    device_map="auto" 
-)
+# 2. Load Model & Processor
+print("Loading model... (This may take a while)")
+try:
+    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model = AutoModelForImageTextToText.from_pretrained(
+        MODEL_ID, 
+        torch_dtype=dtype,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True # optimized loading
+    ).to(device)
+    model.eval()
+    print("✅ Model loaded successfully.")
+except ValueError as e:
+    print("\n🛑 CRITICAL ERROR: Transformers version is too old.")
+    print("You MUST install from source: pip install git+https://github.com/huggingface/transformers.git\n")
+    raise e
 
 def process_textbook(pdf_file, start_page, end_page):
     if pdf_file is None:
@@ -34,16 +46,16 @@ def process_textbook(pdf_file, start_page, end_page):
             pdf_file.name, 
             first_page=int(start_page), 
             last_page=int(end_page), 
-            dpi=150 # 150 DPI is usually sufficient for 8B models
+            dpi=150
         )
         
         extracted_text = []
 
         for i, page_image in enumerate(images):
             page_num = int(start_page) + i
+            print(f"Processing page {page_num}...")
             
-            # 3. CONSTRUCT PROMPT FOR NUMARKDOWN
-            # NuMarkdown expects a standard user/assistant chat format
+            # 3. Construct Qwen2.5-VL / NuMarkdown Prompt
             conversation = [
                 {
                     "role": "user",
@@ -54,27 +66,31 @@ def process_textbook(pdf_file, start_page, end_page):
                 }
             ]
 
-            # Apply the chat template to format the text prompt
+            # Apply chat template
             text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
 
-            # Process inputs (Image + Text) -> Tensors
+            # Prepare Inputs
             inputs = processor(
                 text=[text_prompt], 
                 images=[page_image], 
-                padding=True, 
                 return_tensors="pt"
-            ).to(model.device)
+            ).to(device)
+            
+            # Cast inputs to correct dtype if on CPU (processor usually returns float32, but good to be safe)
+            if device == "cpu":
+                inputs["pixel_values"] = inputs["pixel_values"].to(torch.float32)
+            else:
+                inputs["pixel_values"] = inputs["pixel_values"].to(dtype)
 
-            # 4. GENERATION
+            # 4. Generate
             with torch.no_grad():
                 generated_ids = model.generate(
                     **inputs,
-                    max_new_tokens=4096,  # Increased limit for dense textbook pages
-                    do_sample=False       # Deterministic output usually better for OCR
+                    max_new_tokens=4096,
+                    do_sample=False
                 )
-
-            # 5. DECODING
-            # We slice the output to remove the input prompts from the result
+            
+            # 5. Decode (Slice off the prompt)
             input_len = inputs.input_ids.shape[1]
             generated_ids_trimmed = generated_ids[:, input_len:]
             
@@ -89,12 +105,17 @@ def process_textbook(pdf_file, start_page, end_page):
         return "\n\n---\n\n".join(extracted_text)
     
     except Exception as e:
-        # Catch OOM or other processing errors
-        return f"Processing Error: {str(e)}\n\n*Note: An 8B model requires ~16GB+ VRAM. If you have less, consider 4-bit quantization.*"
+        import traceback
+        traceback.print_exc()
+        return f"Error processing file: {str(e)}\n\n(If on CPU, check if you ran out of RAM. You need ~32GB for this model.)"
 
-# 4. UI Layout
+# 3. UI Layout
 with gr.Blocks() as demo:
-    gr.Markdown("## 📚 Textbook OCR Analysis (NuMarkdown-8B)")
+    gr.Markdown("## 📚 NuMarkdown-8B OCR (Qwen2.5-VL Architecture)")
+    
+    if device == "cpu":
+        gr.Markdown("⚠️ **WARNING: Running on CPU.** Expect very slow performance (minutes per page).")
+    
     with gr.Row():
         with gr.Column():
             file_input = gr.File(label="Upload Textbook PDF")