Update app.py

app.py

@@ -5,7 +5,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from io import BytesIO
 from PIL import Image
-from pdf2image import convert_from_bytes
+import fitz  # PyMuPDF
+import numpy as np
 from transformers import NougatProcessor, VisionEncoderDecoderModel
 
 # Set environment variables
@@ -92,11 +93,18 @@ def extract_text_from_pdf(pdf_bytes):
         # Load Nougat model
         processor, model = load_nougat_model()
         
-        # Convert PDF to images
-        images = convert_from_bytes(pdf_bytes)
+        # Convert PDF to images using PyMuPDF
+        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
         full_text = ""
         
-        for img in images:
+        for page_num in range(len(doc)):
+            page = doc.load_page(page_num)
+            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x zoom for better quality
+            
+            # Convert to PIL Image
+            img_data = pix.samples
+            img = Image.frombytes("RGB", [pix.width, pix.height], img_data)
+            
             # Process with Nougat
             pixel_values = processor(img, return_tensors="pt").pixel_values.to(model.device)
             
@@ -104,7 +112,7 @@ def extract_text_from_pdf(pdf_bytes):
             outputs = model.generate(
                 pixel_values,
                 min_length=1,
-                max_new_tokens=1024,  # Adjust based on expected page content length
+                max_new_tokens=1024,
                 bad_words_ids=[[processor.tokenizer.unk_token_id]],
             )
             
@@ -113,6 +121,9 @@ def extract_text_from_pdf(pdf_bytes):
             page_text = processor.post_process_generation(page_text, fix_markdown=True)
             
             full_text += page_text + "\n\n"
+            
+            # Print progress
+            print(f"Processed page {page_num+1}/{len(doc)}")
         
         # Clear GPU memory
         del pixel_values, outputs
@@ -120,7 +131,9 @@ def extract_text_from_pdf(pdf_bytes):
         
         return full_text
     except Exception as e:
-        print(f"PDF extraction error: {str(e)}")
+        import traceback
+        error_details = traceback.format_exc()
+        print(f"PDF extraction error: {str(e)}\n{error_details}")
         return default_paper_content
     finally:
         # Clear GPU memory