Update app.py

app.py

@@ -1,36 +1,31 @@
-import gradio as gr
 from transformers import AutoTokenizer, AutoModelForImageClassification
-from PIL import Image
-import torch
 
-# Load the model and tokenizer with trust_remote_code set to True
-tokenizer = AutoTokenizer.from_pretrained('stepfun-ai/GOT-OCR2_0', trust_remote_code=True)
-model = AutoModelForImageClassification.from_pretrained('stepfun-ai/GOT-OCR2_0', trust_remote_code=True)
+def load_model_and_tokenizer():
+    try:
+        # Load the tokenizer with the specific revision
+        tokenizer = AutoTokenizer.from_pretrained(
+            'stepfun-ai/GOT-OCR2_0',
+            revision='cf6b7386bc89a54f09785612ba74cb12de6fa17c',  # Pin the specific revision
+            trust_remote_code=True
+        )
 
-def perform_ocr(image):
-    # Ensure the image is in the right format
-    if isinstance(image, Image.Image):
-        image = image.convert("RGB")
-    else:
-        raise ValueError("Input must be a PIL Image")
+        # Load the model with the specific revision
+        model = AutoModelForImageClassification.from_pretrained(
+            'stepfun-ai/GOT-OCR2_0',
+            revision='cf6b7386bc89a54f09785612ba74cb12de6fa17c',  # Pin the specific revision
+            trust_remote_code=True
+        )
 
-    # Use the model to perform OCR
-    inputs = tokenizer(image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Get the predictions
-    predictions = outputs.logits.argmax(dim=1).item()
-    return predictions
+        return model, tokenizer
 
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=perform_ocr,
-    inputs=gr.inputs.Image(type="pil"),
-    outputs="text",
-    title="OCR with GOT-OCR2.0",
-    description="Upload an image for Optical Character Recognition."
-)
+    except Exception as e:
+        print(f"An error occurred while loading the model and tokenizer: {e}")
+        return None, None
 
-# Launch the interface
-iface.launch()
+# Example usage
+if __name__ == "__main__":
+    model, tokenizer = load_model_and_tokenizer()
+    if model and tokenizer:
+        print("Model and tokenizer loaded successfully!")
+    else:
+        print("Failed to load model and tokenizer.")