Update app.py

app.py

@@ -1,22 +1,21 @@
 from fastapi import FastAPI, UploadFile, File
-from transformers import AutoProcessor, AutoModel
+from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
 import torch
 import io
 
 app = FastAPI()
-model_id = "google/siglip2-base-patch16-224"
+model_id = "openai/clip-vit-large-patch14"
 
 # Check for GPU, but default to optimized CPU path
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 # 1. Load with memory-efficient settings
-model = AutoModel.from_pretrained(
+model = CLIPModel.from_pretrained(
     model_id, 
     torch_dtype=dtype, 
-    low_cpu_mem_usage=True,
-    attn_implementation="sdpa" # Use Scaled Dot Product Attention
+    low_cpu_mem_usage=True
 ).to(device).eval()
 
 # 2. COMPILE THE MODEL (The huge speed boost)
@@ -26,22 +25,24 @@ try:
 except Exception:
     print("Torch compile not supported on this environment, skipping...")
 
-processor = AutoProcessor.from_pretrained(model_id)
+processor = CLIPProcessor.from_pretrained(model_id)
 
 # 3. USE 'def' (Not 'async def') for CPU-heavy tasks
 # This allows FastAPI to run searches in parallel on different CPU cores
 @app.post("/embed-text")
 def embed_text(text: str):
-    # GEMMA FIX: max_length=64 is required for SigLIP 2
+    # CLIP uses max 77 tokens for text
     inputs = processor(
         text=[text], 
-        padding="max_length", 
-        max_length=64, 
+        padding=True, 
+        truncation=True,
         return_tensors="pt"
     ).to(device)
     
     with torch.inference_mode(): # Faster than no_grad()
         outputs = model.get_text_features(**inputs)
+        # Normalize embeddings for cosine similarity
+        outputs = outputs / outputs.norm(dim=-1, keepdim=True)
     
     return {"vector": outputs[0].cpu().tolist()}
 
@@ -54,5 +55,7 @@ def embed_image(file: UploadFile = File(...)):
     
     with torch.inference_mode():
         outputs = model.get_image_features(**inputs)
+        # Normalize embeddings for cosine similarity
+        outputs = outputs / outputs.norm(dim=-1, keepdim=True)
         
-    return {"vector": outputs[0].cpu().tolist()}
+    return {"vector": outputs[0].cpu().tolist()}