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

# Workaround for a known Gradio/gradio_client bug: when a component's JSON
# schema is a bare boolean (True/False) instead of a dict, both get_type()
# and _json_schema_to_python_type() crash instead of handling it. This patches
# both to treat a boolean schema as a generic type. See gradio-app/gradio #11084.
import gradio_client.utils as _gc_utils

_original_get_type = _gc_utils.get_type
_original_json_schema_to_python_type = _gc_utils._json_schema_to_python_type


def _patched_get_type(schema):
    if isinstance(schema, bool):
        return "bool"
    return _original_get_type(schema)


def _patched_json_schema_to_python_type(schema, defs=None):
    if isinstance(schema, bool):
        return "Any"
    return _original_json_schema_to_python_type(schema, defs)


_gc_utils.get_type = _patched_get_type
_gc_utils._json_schema_to_python_type = _patched_json_schema_to_python_type

# Model weights live in a separate Hugging Face Model repo (not in this Space's
# own storage, which has a much smaller free limit). This downloads them
# automatically the first time the Space starts.
model_path = "122Uswa/urdu-ocr-codesaviours-si26-uswa"

processor = TrOCRProcessor.from_pretrained(model_path)
model = VisionEncoderDecoderModel.from_pretrained(model_path)
model.eval()


@spaces.GPU
def extract_urdu_text(image):
    """Takes an image, returns extracted Urdu text."""
    if image is None:
        return "Please upload an image"
    pixel_values = processor(image, return_tensors="pt").pixel_values
    with torch.no_grad():
        generated_ids = model.generate(pixel_values, max_new_tokens=128)
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return text if text.strip() else "Could not extract text from this image"


interface = gr.Interface(
    fn=extract_urdu_text,
    inputs=gr.Image(type="pil", label="Upload Urdu Image"),
    outputs=gr.Textbox(label="Extracted Urdu Text"),
    title="Urdu OCR -- Code Saviours SI-26",
    description="Upload an image containing Urdu text and get the extracted text.",
    examples=[],
)

interface.queue().launch()