Update main.py

main.py

@@ -1,5 +1,4 @@
 from fastapi import FastAPI, HTTPException
-from sentence_transformers import SentenceTransformer
 from transformers import AutoImageProcessor, AutoModel
 import torch
 from PIL import Image
@@ -7,31 +6,25 @@ import requests
 from io import BytesIO
 import uvicorn
 
-app = FastAPI(title="Movie Linker AI API")
+app = FastAPI(title="Movie Linker - Image Embedding API (DINOv2)")
 
-# Load Models
-print("Loading Models... please wait.")
+# Load Model
+print("Loading DINOv2 Model... please wait.")
 
-# 1. Image Model: DINOv2 (using transformers directly for stability)
+# Image Model: DINOv2
 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
+img_model.eval()
 
-# 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)
-
-print("All models loaded successfully.")
+print("DINOv2 loaded successfully.")
 
 @app.get("/")
 def home():
     return {
         "status": "online", 
-        "models": {
-            "image": img_model_id, 
-            "text": text_model_name
-        }
+        "model": img_model_id,
+        "endpoint": "/embed/image"
     }
 
 @app.post("/embed/image")
@@ -40,26 +33,20 @@ async def embed_image(image_url: str):
         response = requests.get(image_url, timeout=10)
         img = Image.open(BytesIO(response.content)).convert("RGB")
         
-        # Process image for DINOv2
+        # Process image
         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, :]
+            # Use CLS token for global representation (768 dimensions)
             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))
-
-@app.post("/embed/text")
-async def embed_text(text: str):
-    try:
-        # 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}
+        return {
+            "success": True, 
+            "model": img_model_id,
+            "dimension": len(embedding), 
+            "embedding": embedding
+        }
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))