app.py

import gradio as gr
from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration
from PIL import Image
import torch

# Load model and processor
model_id = "cosmo3769/finetuned_paligemma_vqav2_small"
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id)
processor = PaliGemmaProcessor.from_pretrained("google/paligemma-3b-pt-224")

# Define inference function
def process_image(image, prompt):
    # Process the image and prompt using the processor
    inputs = processor(image.convert("RGB"), prompt, return_tensors="pt")
    
    # Generate output from the model
    output = model.generate(**inputs, max_new_tokens=20)
    
    # Decode and return the output
    decoded_output = processor.decode(output[0], skip_special_tokens=True)
    
    # Return the answer (exclude the prompt part from output)
    return decoded_output[len(prompt):]

# Define the Gradio interface
inputs = [
    gr.Image(type="pil"),
    gr.Textbox(label="Prompt", placeholder="Enter your question")
]
outputs = gr.Textbox(label="Answer")

# Create the Gradio app
demo = gr.Interface(fn=process_image, inputs=inputs, outputs=outputs, title="Finetuned PaliGemma on VQAv2 Small Dataset", 
                    description="Ask a question about an image")

# Launch the app
demo.launch()