Update main.py

main.py

@@ -1,21 +1,24 @@
 from fastapi import FastAPI, HTTPException
 from sentence_transformers import SentenceTransformer
+from transformers import AutoImageProcessor, AutoModel
+import torch
 from PIL import Image
 import requests
 from io import BytesIO
 import uvicorn
 
-app = FastAPI(title="Image Embedding API (CLIP)")
+app = FastAPI(title="Movie Linker AI API")
 
 # Load Models
 print("Loading Models... please wait.")
 
-# 1. Image Model: DINOv2 (768 dim)
-img_model_name = 'facebook/dinov2-base'
-img_model = SentenceTransformer(img_model_name)
+# 1. Image Model: DINOv2 (using transformers directly for stability)
+img_model_id = 'facebook/dinov2-base'
+img_processor = AutoImageProcessor.from_pretrained(img_model_id)
+img_model = AutoModel.from_pretrained(img_model_id)
+img_model.eval() # Set to evaluation mode
 
-# 2. Text Model: Qwen (Choice: 1.5B or 0.6B for speed/memory)
-# Much stronger than E5, works great on CPU
+# 2. Text Model: Qwen (Choice: 1.5B or 0.6B)
 text_model_name = 'Alibaba-NLP/gte-Qwen2-1.5b-instruct'
 text_model = SentenceTransformer(text_model_name, trust_remote_code=True)
 
@@ -26,7 +29,7 @@ def home():
     return {
         "status": "online", 
         "models": {
-            "image": img_model_name, 
+            "image": img_model_id, 
             "text": text_model_name
         }
     }
@@ -36,7 +39,16 @@ async def embed_image(image_url: str):
     try:
         response = requests.get(image_url, timeout=10)
         img = Image.open(BytesIO(response.content)).convert("RGB")
-        embedding = img_model.encode(img).tolist()
+        
+        # Process image for DINOv2
+        inputs = img_processor(images=img, return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = img_model(**inputs)
+            # DINOv2 uses the CLS token (first token) for the global representation
+            # This is available in last_hidden_state[:, 0, :]
+            embedding = outputs.last_hidden_state[:, 0, :].squeeze().tolist()
+            
         return {"success": True, "dimension": len(embedding), "embedding": embedding}
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
@@ -44,7 +56,7 @@ async def embed_image(image_url: str):
 @app.post("/embed/text")
 async def embed_text(text: str):
     try:
-        # E5 model requires 'query: ' prefix for similarity tasks
+        # Instruction-tuned models like Qwen work best with prompts
         processed_text = f"query: {text}"
         embedding = text_model.encode(processed_text).tolist()
         return {"success": True, "dimension": len(embedding), "embedding": embedding}
@@ -52,4 +64,4 @@ async def embed_text(text: str):
         raise HTTPException(status_code=400, detail=str(e))
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=7860)