app.py

"""
CLIP Embed — image → vecteur CLIP 512-dim (normalisé), pour la recherche
produit par image. Aucun LLM ne lit l'image : embedding géométrique pur.
Contrat : POST /embed { "image_base64": "..." } -> { "embedding": [...512], "dim", "model" }
"""
import base64
import io
import os

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from PIL import Image
from sentence_transformers import SentenceTransformer

MODEL_ID = os.environ.get("MODEL_ID", "sentence-transformers/clip-ViT-B-32")
model = SentenceTransformer(MODEL_ID)

app = FastAPI(title="CLIP Embed")


class EmbedRequest(BaseModel):
    image_base64: str | None = None
    model: str | None = None  # ignoré (info appelant) — le modèle est fixé côté Space


@app.api_route("/", methods=["GET", "HEAD"])
def health():
    return {"ok": True, "model": MODEL_ID}


@app.api_route("/embed", methods=["GET", "HEAD"])
def embed_health():
    return {"ok": True, "hint": "POST image_base64 to this endpoint"}


@app.post("/embed")
def embed(req: EmbedRequest):
    if not req.image_base64:
        raise HTTPException(status_code=400, detail="image_base64 required")
    try:
        raw = base64.b64decode(req.image_base64)
        img = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(status_code=400, detail=f"bad image: {exc}")
    vec = model.encode(img, normalize_embeddings=True).tolist()
    return {"embedding": vec, "dim": len(vec), "model": MODEL_ID}