Update app.py

app.py

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