app.py

from transformers import VisionEncoderDecoderModel, TrOCRProcessor, AutoTokenizer, ViTImageProcessor
import gradio as gr
from PIL import Image

MODEL_ID = "danush99/Model_TrOCR-Sin-Printed-Text"

# Load model and processor once at startup
model = None
processor = None

print("Loading model and processor...")
try:
    model = VisionEncoderDecoderModel.from_pretrained(MODEL_ID)
    tokenizer = AutoTokenizer.from_pretrained("NLPC-UOM/SinBERT-large")
    feature_extractor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
    processor = TrOCRProcessor(image_processor=feature_extractor, tokenizer=tokenizer)
    print("Ready.")
except OSError as e:
    if "pytorch_model.bin" in str(e) or "model.safetensors" in str(e):
        print("Weights not found. Upload model.safetensors to", MODEL_ID)
    else:
        raise


def ocr(image):
    if model is None or processor is None:
        return (
            "Model weights are not yet on the Hub. "
            "Upload model.safetensors to https://huggingface.co/danush99/Model_TrOCR-Sin-Printed-Text "
        )
    pixel_values = processor(image, return_tensors="pt").pixel_values
    generated_ids = model.generate(pixel_values)
    return processor.batch_decode(generated_ids, skip_special_tokens=True)[0]


demo = gr.Interface(
    fn=ocr,
    inputs=gr.Image(show_label=False, type="pil"),
    outputs=gr.Textbox(),
    title="Sinhala OCR (TrOCR-Sin-Printed-Text)",
    description="Uses **danush99/Model_TrOCR-Sin-Printed-Text** for Sinhala normal/printed text OCR. Upload an image to get started.",
)
# share=True is not supported on Hugging Face Spaces
demo.launch(share=False)