app.py

import gradio as gr
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
import torch
from huggingface_hub import login
import os

# Authenticate before any model loading
hf_token = os.getenv("HF_TOKEN")
# if hf_token:
#     login(token=hf_token)
# else:
#     print("No HF_TOKEN found in environment")

# Load model components with explicit token
# try:
#     model = PaliGemmaForConditionalGeneration.from_pretrained(
#         "pyimagesearch/finetuned_paligemma_vqav2_small",
#         token=hf_token
#     )
#     processor = AutoProcessor.from_pretrained(
#         "pyimagesearch/finetuned_paligemma_vqav2_small", 
#         token=hf_token
#     )
#     print("Model loaded successfully!")
# except Exception as e:
#     print(f"Error: {e}")
    
# Load model and processor
model_id = "google/paligemma2-3b-pt-224"
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id)
processor = AutoProcessor.from_pretrained("google/paligemma2-3b-pt-224", token=hf_token)

# 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")
  
   try:
       # 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):]
   except IndexError as e:
       print(f"IndexError: {e}")
       return "An error occurred during processing."
# 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="PaliGemma 2 Model Showcase", description="Upload an image and ask questions to get answers.")

# Launch the app
demo.launch()