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

# 載入模型與前處理器
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# 定義圖片描述生成函數
def generate_caption(image):
    if image is None:
        return "請上傳一張圖片。"
    inputs = processor(images=image, return_tensors="pt")
    out = model.generate(**inputs, max_new_tokens=50)
    caption = processor.decode(out[0], skip_special_tokens=True)
    return caption

# Gradio 介面設定
app = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil", label="上傳圖片"),
    outputs=gr.Textbox(label="AI 生成的圖片描述"),
    title="🖼️ AI 影像描述生成器",
    description="上傳圖片後，AI 將自動生成自然語言描述。使用模型：Salesforce / BLIP Image Captioning Base"
)

# 啟動應用
if __name__ == "__main__":
    app.launch()