app.py

from transformers import pipeline
import gradio as gr
from PIL import Image
import requests
from io import BytesIO

# 修正：使用正確的 pipeline 類型
pipe = pipeline("image-text-to-text", model="google/medgemma-4b-it")

# 檢查輸入是否為 URL，如果是，則下載並轉換為 PIL 圖片
def load_image_from_input(image_input):
    if isinstance(image_input, str) and (image_input.startswith("http://") or image_input.startswith("https://")):
        try:
            response = requests.get(image_input)
            img = Image.open(BytesIO(response.content))
            return img
        except Exception as e:
            raise gr.Error(f"無法從 URL 下載圖片: {e}")
    else:
        # 如果不是 URL，就假設是檔案路徑，由 Gradio 自動處理
        return Image.open(image_input)

# 包裝成 API 函數
def predict(image_input, question):
    image = load_image_from_input(image_input)
    
    # 將輸入轉換為模型所需的 messages 格式
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},  # 使用 "image" 而不是 "url"
                {"type": "text", "text": question}
            ]
        },
    ]
    
    # 修正：使用 text 參數而不是直接傳遞 messages
    result = pipe(text=messages)
    
    # 修正輸出格式：處理新的結果格式
    # 檢查是否是新格式：包含 generated_text 的 dict
    if isinstance(result, list) and len(result) > 0 and isinstance(result[0], dict):
        first_result = result[0]
        if "generated_text" in first_result:
            generated_text = first_result["generated_text"]
            # generated_text 是一個 list，找到 assistant 的回應
            if isinstance(generated_text, list):
                for message in generated_text:
                    if isinstance(message, dict) and message.get("role") == "assistant":
                        return message.get("content", "無法獲取回應")
    
    # 備用：舊格式處理
    if isinstance(result, list):
        for message in result:
            if isinstance(message, dict) and message.get("role") == "assistant":
                return message.get("content", "無法獲取回應")
    
    # 最後備用方案
    return str(result)

# Gradio 介面（同時支援 UI 和 API）
iface = gr.Interface(
    fn=predict,
    inputs=[gr.Image(type="filepath"), "text"],
    outputs="text",
    title="MedGemma API + Demo",
    description="上傳圖片或輸入圖片 URL，以 API 或 UI 測試 MedGemma。"
)

# 啟動應用程式
if __name__ == "__main__":  # 修正語法錯誤：__name__ 不是 **name**
    iface.launch()