Update app.py

app.py

@@ -1,23 +1,30 @@
+import torch
 import gradio as gr
 from transformers import AutoModel, AutoTokenizer
 import os
-import re  # Import regular expressions module
+import re  # For keyword searching and highlighting
+
+# Ensure that the code runs on the CPU
+device = torch.device('cpu')
 
 # Load the OCR model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
+
+# Load the model onto the CPU
 model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', 
                                   trust_remote_code=True, 
                                   low_cpu_mem_usage=True, 
-                                  pad_token_id=tokenizer.eos_token_id).eval()  # Removed device_map='cuda' and .cuda()
+                                  pad_token_id=tokenizer.eos_token_id).eval().to(device)
 
-# Define the function to process images and extract text
+# Function to extract text from an image
 def extract_text_from_image(image):
     # Save the uploaded image temporarily
     image_path = "temp_image.jpg"
     image.save(image_path)
     
-    # Call the model to perform OCR
-    extracted_text = model.chat(tokenizer, image_path, ocr_type='ocr')
+    # Perform OCR using the model, ensuring it runs on CPU
+    with torch.no_grad():
+        extracted_text = model.chat(tokenizer, image_path, ocr_type='ocr')
     
     # Remove the temporary image file
     os.remove(image_path)