Update app.py

app.py

@@ -1,46 +1,38 @@
 import gradio as gr
-import torch
-from transformers import AutoModel, AutoTokenizer
 from PIL import Image
+# Assuming 'model' and 'tokenizer' are defined elsewhere in your code
+# from your_model_file import model, tokenizer
+
+def load_image(image_file):
+    """Load and preprocess the image."""
+    if isinstance(image_file, Image.Image):  # Check if the input is an Image object
+        return image_file.convert("RGB")  # Convert to RGB if necessary
+    elif isinstance(image_file, str) and (image_file.startswith('http') or image_file.startswith('https')):
+        # Handle URL case (you can use an external library to fetch the image if needed)
+        return Image.open(requests.get(image_file, stream=True).raw).convert("RGB")
+    else:
+        # Handle file path case
+        return Image.open(image_file).convert("RGB")
 
-# Check GPU availability
-if torch.cuda.is_available():
-    print("CUDA is available! GPU is present.")
-    print(f"Number of GPUs: {torch.cuda.device_count()}")
-    print(f"GPU Name: {torch.cuda.get_device_name(0)}")
-else:
-    print("CUDA is not available. Running on CPU.")
-
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
-
-# Initialize the model
-if torch.cuda.is_available():
-    model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
-    model = model.eval().cuda()
-else:
-    model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, pad_token_id=tokenizer.eos_token_id)
-    model = model.eval()  # Keep model on CPU
-
-# Define the OCR function
 def perform_ocr(image):
-    # Convert PIL image to RGB format (if necessary)
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-
-    # Perform OCR using the model
-    res = model.chat(tokenizer, image, ocr_type='ocr')
-
-    return res
-
-# Define the Gradio interface
-interface = gr.Interface(
+    """Perform OCR on the uploaded image."""
+    try:
+        # Load and preprocess the image
+        processed_image = load_image(image)
+        # Use the model for OCR
+        res = model.chat(tokenizer, processed_image, ocr_type='ocr')
+        return res
+    except Exception as e:
+        return str(e)  # Return the error message
+
+# Gradio interface setup
+iface = 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."
+    inputs=gr.Image(type="pil"),  # Ensure Gradio accepts images as PIL images
+    outputs="text",
+    title="OCR Application",
+    description="Upload an image to perform Optical Character Recognition (OCR)."
 )
 
-# Launch the Gradio app
-interface.launch()
+if __name__ == "__main__":
+    iface.launch()