Update app.py

app.py

@@ -1,52 +1,44 @@
 from fastapi import FastAPI, UploadFile, File
-from pydantic import BaseModel
 from transformers import AutoProcessor, AutoModel
 from PIL import Image
 import torch
 import io
-import os
-
-# Set higher timeout for model downloading
-os.environ["HF_HUB_READ_TIMEOUT"] = "60" 
 
 app = FastAPI()
-
 model_id = "google/siglip2-so400m-patch14-384"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# FIX 1: Use torch_dtype directly (deprecation fix)
-# FIX 2: use low_cpu_mem_usage to prevent RAM spikes on 16GB
+# Optimized loading based on Feb 2025 Docs
 model = AutoModel.from_pretrained(
     model_id, 
     torch_dtype=torch.float32, 
-    low_cpu_mem_usage=True
+    low_cpu_mem_usage=True,
+    attn_implementation="sdpa" # Scaled Dot Product Attention for speed
 ).to(device).eval()
 
-# FIX 3: Explicitly set use_fast=True to avoid the processor warning
-processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
-
-# OPTIMIZATION: Faster inference
-try:
-    model = torch.compile(model)
-except:
-    pass
-
-class TextRequest(BaseModel):
-    text: str
+processor = AutoProcessor.from_pretrained(model_id)
 
 @app.post("/embed-text")
-def embed_text(request: TextRequest):
-    inputs = processor(text=[request.text], padding="max_length", return_tensors="pt").to(device)
+def embed_text(text: str):
+    # Docs specify max_length=64 for the Gemma-based tokenizer in SigLIP 2
+    inputs = processor(
+        text=[text], 
+        padding="max_length", 
+        max_length=64, 
+        return_tensors="pt"
+    ).to(device)
+    
     with torch.inference_mode():
-        text_outputs = model.get_text_features(**inputs)
-    return {"vector": text_outputs[0].cpu().tolist(), "dim": 1152}
+        outputs = model.get_text_features(**inputs)
+    return {"vector": outputs[0].cpu().tolist()}
 
 @app.post("/embed-image")
 def embed_image(file: UploadFile = File(...)):
-    image_data = file.file.read() 
-    image = Image.open(io.BytesIO(image_data)).convert("RGB")
+    image = Image.open(io.BytesIO(file.file.read())).convert("RGB")
     
+    # NaFlex logic is handled automatically by the processor
     inputs = processor(images=image, return_tensors="pt").to(device)
+    
     with torch.inference_mode():
-        image_outputs = model.get_image_features(**inputs)
-    return {"vector": image_outputs[0].cpu().tolist(), "dim": 1152}
+        outputs = model.get_image_features(**inputs)
+    return {"vector": outputs[0].cpu().tolist()}