# app.py
import os
import io
from huggingface_hub import InferenceClient
from PIL import Image
import gradio as gr
from transformers import pipeline

# 模型設定
CAPTION_MODEL = "Salesforce/blip-image-captioning-base"
VQA_MODEL = "Salesforce/blip-vqa-base"

# 初始化 pipelines
captioner = pipeline("image-to-text", model=CAPTION_MODEL)
vqa_pipeline = pipeline("visual-question-answering", model=VQA_MODEL)


def caption_image(image: Image.Image):
    """圖片描述"""
    try:
        result = captioner(image)
        return result[0]["generated_text"] if result else "無法產生描述"
    except Exception as e:
        return f"Caption 錯誤：{e}"


def vqa_image(image: Image.Image, question: str):
    """圖片問答 (VQA)"""
    try:
        result = vqa_pipeline(image=image, question=question)
        if isinstance(result, list) and len(result) > 0:
            return result[0]["answer"]
        return "無法取得答案"
    except Exception as e:
        return f"VQA 錯誤：{e}"


# 建立 Gradio 介面
with gr.Blocks(title="Vision-Language Web UI") as demo:
    gr.Markdown("# 🖼️ Vision-Language Web UI — Caption + VQA")

    with gr.Row():
        with gr.Column(scale=1):
            img_in = gr.Image(type="pil", label="上傳圖片")
            question = gr.Textbox(label="圖片問題", placeholder="這張圖片裡有什麼？")
            btn_caption = gr.Button("生成圖片描述")
            btn_vqa = gr.Button("回答圖片問題")

        with gr.Column(scale=1):
            out_caption = gr.Textbox(label="圖片描述結果")
            out_vqa = gr.Textbox(label="VQA 回答")

    btn_caption.click(caption_image, inputs=img_in, outputs=out_caption)
    btn_vqa.click(vqa_image, inputs=[img_in, question], outputs=out_vqa)

    gr.Markdown("---\n**提示**：第一次啟動會下載模型，請稍等幾十秒。")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))