import gradio as gr
import torch
from transformers import AutoModel, AutoTokenizer

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map=device, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
model = model.eval().to(device)

def ocr_pipeline(image):
    """OCR processing function for Gradio."""
    res = model.chat(tokenizer, image, ocr_type='ocr')
    return res

# Gradio Interface
iface = gr.Interface(
    fn=ocr_pipeline,
    inputs=gr.Image(type="filepath"),  # Allows users to upload an image
    outputs="text",
    title="OCR Model",
    description="Upload an image to extract text using GOT-OCR2_0."
)

# Launch the app
iface.launch()