import spaces
import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import torch

print("Loading TrOCR model...")
device = "cuda" if torch.cuda.is_available() else "cpu"

# use_fast=False completely bypasses the transformers tokenizer conversion bug!
processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten', use_fast=False)
model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten').to(device)
print("Model loaded successfully!")

@spaces.GPU
def predict(image):
    if image is None:
        return ""
    
    pixel_values = processor(image, return_tensors="pt").pixel_values.to(model.device)
    
    generated_ids = model.generate(pixel_values)
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return generated_text

demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs="text",
    title="TrOCR Hub Backend"
)

demo.launch()