import gradio as gr
from transformers import pipeline

# Load model
pipe = pipeline(
    "image-text-to-text",
    model="ibm-granite/granite-4.0-3b-vision",
    trust_remote_code=True
)

def analyze_image(image, question):
    if image is None or question.strip() == "":
        return "Please upload an image and enter a question."

    # Format messages like your original code
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": question}
            ]
        }
    ]

    result = pipe(text=messages)

    # Extract output safely
    try:
        return result[0]["generated_text"]
    except:
        return str(result)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ Image Question Answering (Granite Vision Model)")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")
        question_input = gr.Textbox(label="Ask a question about the image")

    output = gr.Textbox(label="Answer")

    submit_btn = gr.Button("Submit")
    submit_btn.click(analyze_image, inputs=[image_input, question_input], outputs=output)

# Launch
demo.launch()