import gradio as gr
from PIL import Image
import torch
from transformers import AutoProcessor, AutoModel

# Load model and processor once on startup
model_name = "Graf-J/captcha-conv-transformer-finetuned"
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
model.eval()

def predict_captcha(image):
    if image is None:
        return "No image uploaded"
    # Convert to PIL Image
    img = Image.fromarray(image).convert("RGB")
    inputs = processor(img)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    decoded = processor.batch_decode(logits)
    return decoded[0] if decoded else "Failed to decode"

# Create Gradio interface
demo = gr.Interface(
    fn=predict_captcha,
    inputs=gr.Image(type="numpy", label="Upload CAPTCHA"),
    outputs=gr.Textbox(label="Decoded CAPTCHA"),
    title="CAPTCHA Solver API",
    description="Upload a CAPTCHA image to solve it using Graf-J/captcha-conv-transformer-finetuned."
)

demo.launch()