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

model_id = "iGwangsu/my-blip-model"

processor = BlipProcessor.from_pretrained(model_id, use_fast=True)
model = BlipForConditionalGeneration.from_pretrained(
    model_id,
    low_cpu_mem_usage=True
)


def generate_caption(img):
    if img is None: return "이미지를 업로드해주세요."
    inputs = processor(images=img, return_tensors="pt")

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_length=50
        )

    caption = processor.decode(out[0], skip_special_tokens=True)
    return caption

demo = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="BLIP Image Captioning"
)

demo.launch()