import gradio as gr
import torch
from PIL import Image
import requests
from transformers import AutoProcessor, SiglipModel

# 1. 載入模型 (只會在啟動時載入一次)
model_id = "google/siglip2-base-patch16-224"
print("正在載入 SigLIP 2 模型...")
processor = AutoProcessor.from_pretrained(model_id)
model = SiglipModel.from_pretrained(model_id)
print("模型載入完成！")

# 2. 定義產生 Embedding 的核心運算
def get_embedding(image_url):
    try:
        # 下載圖片
        image = Image.open(requests.get(image_url, stream=True).raw)
        inputs = processor(images=image, return_tensors="pt")
        
        with torch.no_grad():
            # 提取 768 維度圖片向量並歸一化
            image_features = model.get_image_features(**inputs)
            image_features = image_features / image_features.norm(p=2, dim=-1, keepdim=True)
            
        embedding_list = image_features.squeeze().tolist()
        return {"status": "success", "dimension": len(embedding_list), "embedding": embedding_list}
    except Exception as e:
        return {"status": "error", "message": str(e)}

# 3. 建立 Gradio API 介面 (這會自動產生 n8n 可用的 API)
demo = gr.Interface(
    fn=get_embedding,
    inputs=gr.Textbox(label="輸入圖片網址 (Image URL)"),
    outputs=gr.JSON(label="輸出的 768 維 Embedding 向量"),
    title="SigLIP 2 Embedding Generator"
)

# 啟動服務
if __name__ == "__main__":
    demo.launch()