import gradio as gr
from transformers import AutoProcessor, BlipForConditionalGeneration
from PIL import Image
import numpy as np

# Load BLIP model
processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

def caption_image(input_image: np.ndarray):
    # Convert numpy array to PIL Image
    raw_image = Image.fromarray(input_image).convert('RGB')
    
    # Generate caption
    inputs = processor(images=raw_image, text="a photo of", return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=50)
    caption = processor.decode(outputs[0], skip_special_tokens=True)
    
    return caption

# Gradio interface
iface = gr.Interface(
    fn=caption_image,
    inputs=gr.Image(),
    outputs="text",
    title="Image Captioning",
)

iface.launch()