FloatyPotato/snap-nexus-embed

1from fastapi import FastAPI, UploadFile, File, HTTPException2from PIL import Image3import io4import torch5from transformers import CLIPProcessor, CLIPModel6 7app = FastAPI(title="Snap Nexus Embedding Server")8 9model_name = "openai/clip-vit-base-patch16"10print(f"Loading CLIP model {model_name}...")11model = CLIPModel.from_pretrained(model_name)12processor = CLIPProcessor.from_pretrained(model_name)13device = "cuda" if torch.cuda.is_available() else "cpu"14model.to(device)15print(f"Model loaded successfully on {device}!")16 17@app.post("/embed")18async def get_embedding(file: UploadFile = File(...)):19    try:20        # Read image21        image_bytes = await file.read()22        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")23        24        # Preprocess and embed25        inputs = processor(images=image, return_tensors="pt")26        inputs = {k: v.to(device) for k, v in inputs.items()}27        28        with torch.no_grad():29            image_features = model.get_image_features(**inputs)30            31        # Convert to list and return32        embedding = image_features[0].cpu().numpy().tolist()33        return {"embedding": embedding}34    except Exception as e:35        raise HTTPException(status_code=500, detail=str(e))36 37@app.get("/")38def read_root():39    return {"status": "ok", "model": model_name, "device": device}40 41