Update app.py

app.py

@@ -5,22 +5,21 @@ from PIL import Image
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
 model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
-model = model.eval()  # Remove .cuda() to run on CPU
+model = model.eval()  # Ensure model is set to evaluation mode
 
-# Define the OCR function
+# Define the OCR function with error handling
 def perform_ocr(image):
-    # Convert PIL image to RGB format (if necessary)
-    if image.mode != "RGB":
-        image = image.convert("RGB")
+    try:
+        # Convert PIL image to RGB format (if necessary)
+        if image.mode != "RGB":
+            image = image.convert("RGB")
 
-    # Save the image to a temporary path
-    image_file_path = 'temp_image.jpg'
-    image.save(image_file_path)
+        # Perform OCR using the model
+        res = model.chat(tokenizer, image, ocr_type='ocr')
 
-    # Perform OCR using the model
-    res = model.chat(tokenizer, image_file_path, ocr_type='ocr')
-
-    return res
+        return res
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
 
 # Define the Gradio interface
 interface = gr.Interface(