handler.py

"""Hugging Face Inference Endpoints entry point — deploy this repo as a CPU/GPU API.

Request bodies:
    {"inputs": <image>}                                   -> {person: probability} (all 42, best first)
    {"inputs": <image>, "parameters": {"task": "embed"}}  -> {"embedding": [1024 floats]}
    {"inputs": {"image_a": <image>, "image_b": <image>}, "parameters": {"threshold": 0.5}}
                                                          -> {"same_person", "distance", "threshold", ...}
<image> is a base64 string (or data URL); the endpoint also passes raw image uploads as PIL images.
"""
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import model as M  # noqa: E402


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.predictor = M.load(path or HERE, "cuda" if M.cuda_available() else "cpu")

    def __call__(self, data: dict):
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", None) or {}
        if isinstance(inputs, dict) and {"image_a", "image_b"} <= inputs.keys():
            return self.predictor.verify(inputs["image_a"], inputs["image_b"], parameters.get("threshold"))
        if parameters.get("task") == "embed":
            return {"embedding": self.predictor.embed(inputs)}
        return self.predictor.predict(inputs)