Update main.py

main.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoImageProcessor, AutoModel
 import torch
 from PIL import Image
@@ -8,29 +9,19 @@ import uvicorn
 
 app = FastAPI(title="Movie Linker - Image Embedding API (DINOv2)")
 
-# Load Model
-print("Loading DINOv2 Model... please wait.")
+class ImageRequest(BaseModel):
+    image_url: str
 
-# Image Model: DINOv2
+# Load Model
 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()
 
-print("DINOv2 loaded successfully.")
-
-@app.get("/")
-def home():
-    return {
-        "status": "online", 
-        "model": img_model_id,
-        "endpoint": "/embed/image"
-    }
-
 @app.post("/embed/image")
-async def embed_image(image_url: str):
+async def embed_image(request: ImageRequest):
     try:
-        response = requests.get(image_url, timeout=10)
+        response = requests.get(request.image_url, timeout=10)
         img = Image.open(BytesIO(response.content)).convert("RGB")
         
         # Process image