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

model_name = "ucaslcl/GOT-OCR2_0"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, device_map=device)
model = model.eval().to(device)


@spaces.GPU()
def ocr_process(image):
    if image is None:
        return "错误:未提供图片"

    try:
        res = model.chat(tokenizer, image, ocr_type="ocr")
        return res
    except Exception as e:
        return f"错误: {str(e)}"


with gr.Blocks() as demo:
    gr.Markdown("# OCR 图像识别")

    with gr.Row():
        image_input = gr.Image(type="filepath", label="上传图片")

    submit_button = gr.Button("开始OCR识别")

    output_text = gr.Textbox(label="识别结果")

    submit_button.click(ocr_process, inputs=[image_input], outputs=[output_text])

if __name__ == "__main__":
    demo.launch()