import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import torch

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

def describe_image(image):
    # Convert image to RGB
    image = image.convert("RGB")
    
    # Prepare inputs
    inputs = processor(images=image, return_tensors="pt")
    
    # Generate caption
    with torch.no_grad():
        output = model.generate(**inputs, max_length=50)  # max_length ensures longer captions
    caption = processor.decode(output[0], skip_special_tokens=True)
    
    # Ensure at least 20 words
    if len(caption.split()) < 20:
        caption += " This image appears detailed and contains multiple elements that make it visually interesting and descriptive."
    
    return caption

# Gradio interface
iface = gr.Interface(fn=describe_image, inputs=gr.Image(type="pil"), outputs="text", title="Image Description App",
                     description="Upload an image and get a descriptive caption (at least 20 words).")
iface.launch()