import gradio as gr
from transformers import (
    BlipProcessor,
    BlipForConditionalGeneration,
    BlipForQuestionAnswering
)

print("Loading models...")

caption_processor = BlipProcessor.from_pretrained(
    "Salesforce/blip-image-captioning-base"
)

caption_model = BlipForConditionalGeneration.from_pretrained(
    "Salesforce/blip-image-captioning-base"
)

vqa_processor = BlipProcessor.from_pretrained(
    "Salesforce/blip-vqa-base"
)

vqa_model = BlipForQuestionAnswering.from_pretrained(
    "Salesforce/blip-vqa-base"
)

print("Models Loaded!")


def analyze_image(image, question):

    if image is None:
        return "", ""

    caption_inputs = caption_processor(
        image,
        return_tensors="pt"
    )

    caption_ids = caption_model.generate(
        **caption_inputs,
        max_new_tokens=30
    )

    caption = caption_processor.decode(
        caption_ids[0],
        skip_special_tokens=True
    )

    answer = ""

    if question and question.strip():

        vqa_inputs = vqa_processor(
            image,
            question,
            return_tensors="pt"
        )

        answer_ids = vqa_model.generate(
            **vqa_inputs,
            max_new_tokens=10
        )

        answer = vqa_processor.decode(
            answer_ids[0],
            skip_special_tokens=True
        )

    return caption, answer


with gr.Blocks(
    title="Visual AI Assistant"
) as demo:

    gr.Markdown(
        """
# Visual AI Assistant

This application uses BLIP models for:

- Image Captioning
- Visual Question Answering (VQA)

Upload an image and optionally ask a question about it.
"""
    )

    image_input = gr.Image(
        type="pil",
        label="Upload Image"
    )

    question_input = gr.Textbox(
        label="Ask a Question",
        placeholder="What is happening in this image?"
    )

    with gr.Row():
        analyze_btn = gr.Button(
            "Analyze Image"
        )

        clear_btn = gr.ClearButton(
            [
                image_input,
                question_input
            ]
        )

    caption_output = gr.Textbox(
        label="Generated Caption"
    )

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

    analyze_btn.click(
        fn=analyze_image,
        inputs=[
            image_input,
            question_input
        ],
        outputs=[
            caption_output,
            answer_output
        ]
    )

demo.launch()