import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from PIL import Image
import torch

# Load model dan tokenizer
model_id = "vikhyatk/moondream2"
# Menggunakan revision terbaru agar lebih stabil
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)

def answer_question(image, question):
    if image is None:
        return "No image provided"
    
    # Proses gambar ke format yang dikenali model
    enc_image = model.encode_image(image)
    answer = model.answer_question(enc_image, question, tokenizer)
    return answer

# Interface Gradio sederhana
interface = gr.Interface(
    fn=answer_question,
    inputs=[gr.Image(type="pil"), gr.Textbox(label="Question")],
    outputs=gr.Text(label="Answer"),
    title="Moondream Captcha Solver"
)

interface.launch()