Update app.py

app.py

@@ -1,29 +1,42 @@
 from fastapi import FastAPI, UploadFile, File
 from pydantic import BaseModel
-from sentence_transformers import SentenceTransformer
+from transformers import AutoProcessor, AutoModel
 from PIL import Image
+import torch
 import io
 
 app = FastAPI()
 
-# Load model into memory (do this globally so it only happens once)
-model = SentenceTransformer('google/siglip-so400m-patch14-384')
+# Load SigLIP 2
+model_id = "google/siglip2-so400m-patch14-384"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load Model and Processor
+model = AutoModel.from_pretrained(model_id).to(device).eval()
+processor = AutoProcessor.from_pretrained(model_id)
 
 class TextRequest(BaseModel):
     text: str
 
 @app.post("/embed-text")
 async def embed_text(request: TextRequest):
-    # Convert text to vector
-    vector = model.encode(request.text).tolist()
-    return {"vector": vector}
+    inputs = processor(text=[request.text], padding="max_length", return_tensors="pt").to(device)
+    with torch.no_grad():
+        # Get the text embeddings
+        text_outputs = model.get_text_features(**inputs)
+    
+    vector = text_outputs[0].cpu().tolist()
+    return {"vector": vector, "dim": len(vector)}
 
 @app.post("/embed-image")
 async def embed_image(file: UploadFile = File(...)):
-    # Read uploaded image
     image_data = await file.read()
-    image = Image.open(io.BytesIO(image_data))
+    image = Image.open(io.BytesIO(image_data)).convert("RGB")
     
-    # Convert image to vector
-    vector = model.encode(image).tolist()
-    return {"vector": vector}
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        # Get the image embeddings
+        image_outputs = model.get_image_features(**inputs)
+        
+    vector = image_outputs[0].cpu().tolist()
+    return {"vector": vector, "dim": len(vector)}