Update app.py

app.py

@@ -1,39 +1,32 @@
 import gradio as gr
 from transformers import AutoModel, AutoTokenizer
 from PIL import Image
-import tempfile
+import torch
 
-# Load the tokenizer and model without GPU settings
+# Load the OCR model and tokenizer
 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()  # Set to evaluation mode
+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()
 
-# Define the OCR function with error handling
-def perform_ocr(image):
-    try:
-        # Convert PIL image to RGB format (if necessary)
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        # Save the image to a temporary file
-        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
-            image.save(temp_file.name)
+# Check if GPU is available and use it, else use CPU
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device)
 
-            # Perform OCR using the model
-            res = model.chat(tokenizer, temp_file.name, ocr_type='ocr')  # Pass the file path
-
-        return res
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
+# Function to perform OCR on the image
+def perform_ocr(image):
+    # Perform OCR using the model
+    result = model.chat(tokenizer, image, ocr_type='ocr')
+    return result
 
-# Define the Gradio interface
+# Create the Gradio interface
 interface = gr.Interface(
-    fn=perform_ocr,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Textbox(label="Extracted Text"),
-    title="OCR and Document Search Web Application",
-    description="Upload an image to extract text using the GOT-OCR2_0 model."
+    fn=perform_ocr, 
+    inputs=gr.inputs.Image(type="pil"),  # Accepts an image input
+    outputs="text",  # Outputs extracted text
+    title="OCR Web App",
+    description="Upload an image to extract text using the GOT-OCR2.0 model."
 )
 
-# Launch the Gradio app
+# Launch the app
 interface.launch()