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

MODEL_ID = "Salesforce/blip-image-captioning-base"

processor = BlipProcessor.from_pretrained(MODEL_ID)
model = BlipForConditionalGeneration.from_pretrained(MODEL_ID)
model.eval()


@spaces.GPU
def caption(image: Image.Image) -> str:
    if image is None:
        return ""
    # GPU is only available inside this function on ZeroGPU
    model.to("cuda")
    # BLIP expects RGB input
    inputs = processor(images=image.convert("RGB"), return_tensors="pt").to("cuda")
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=50)
    return processor.decode(out[0], skip_special_tokens=True)


demo = gr.Interface(
    fn=caption,
    inputs=gr.Image(type="pil", label="Image"),
    outputs=gr.Textbox(label="Caption"),
    title="BLIP Image Captioning",
)

demo.launch()