Add custom Inference Endpoints handler for BioCLIP 2 (open_clip ViT-L/14)

handler.py

@@ -0,0 +1,107 @@
+# Custom handler for Hugging Face Inference Endpoints.
+# Serves imageomics/bioclip-2 (open_clip ViT-L/14, BioCLIP 2, 768-dim embeddings).
+#
+# API contract (POST JSON to the endpoint):
+#   {"inputs": {
+#       "image": "<base64-encoded image bytes>",   # or "https://...jpg" URL
+#       "candidate_labels": ["Canis familiaris (dog)", ...],   # required for classify
+#       "mode": "classify"   # default: zero-shot classification over candidate_labels
+#   }}
+#   -> [{"label": "...", "score": 0.93}, ...]  sorted by score desc
+#
+#   {"inputs": {"image": "<base64>", "mode": "embed"}}
+#   -> {"embedding": [768 floats]}  (L2-normalized image embedding)
+#
+#   {"inputs": {"text": "...", "mode": "embed_text"}}
+#   -> {"embedding": [768 floats]}  (L2-normalized text embedding)
+
+from typing import Any, Dict, List, Union
+
+import base64
+import io
+import urllib.request
+
+import open_clip
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # The endpoint mounts the repo contents at `path`; open_clip's
+        # "local-dir:" loader reads open_clip_config.json + weights + tokenizer
+        # from there. Fall back to the Hub if the mounted path is empty.
+        model_id = f"local-dir:{path}" if path else "hf-hub:imageomics/bioclip-2"
+
+        load_kwargs = {}
+        if self.device == "cuda":
+            load_kwargs["precision"] = "fp16"  # ~0.9GB VRAM instead of ~1.7GB
+        try:
+            self.model, _, self.preprocess = open_clip.create_model_and_transforms(
+                model_id, **load_kwargs
+            )
+        except TypeError:
+            # open_clip versions without the `precision` factory kwarg
+            self.model, _, self.preprocess = open_clip.create_model_and_transforms(model_id)
+
+        self.model = self.model.to(self.device).eval()
+        self.tokenizer = open_clip.get_tokenizer(model_id)
+        print(f"[bioclip-2 handler] loaded {model_id} on {self.device}")
+
+    def _load_image(self, image_field: str) -> Image.Image:
+        if image_field.startswith("http://") or image_field.startswith("https://"):
+            with urllib.request.urlopen(image_field, timeout=15) as r:
+                raw = r.read()
+        else:
+            if "," in image_field[:64] and image_field[:5].lower() == "data:":
+                image_field = image_field.split(",", 1)[1]
+            raw = base64.b64decode(image_field)
+        return Image.open(io.BytesIO(raw)).convert("RGB")
+
+    def _autocast(self):
+        # fp16 weights on CUDA: align input/weight dtypes via autocast —
+        # casting the image tensor manually breaks the first conv in open_clip.
+        return torch.autocast(
+            device_type=self.device, dtype=torch.float16, enabled=self.device == "cuda"
+        )
+
+    @torch.no_grad()
+    def __call__(self, data: Union[Dict[str, Any], List[Dict[str, Any]]]) -> Any:
+        # Some handler generations wrap payloads in a list; accept both.
+        if isinstance(data, list):
+            data = data[0]
+
+        inputs = data.pop("inputs", data)
+        mode = inputs.pop("mode", "classify")
+
+        if mode == "embed_text":
+            with self._autocast():
+                tok = self.tokenizer([inputs["text"]]).to(self.device)
+                emb = F.normalize(self.model.encode_text(tok), dim=-1)[0].float()
+            return {"embedding": emb.tolist()}
+
+        image = self._load_image(inputs["image"])
+        img = self.preprocess(image).unsqueeze(0).to(self.device)
+
+        with self._autocast():
+            img_emb = F.normalize(self.model.encode_image(img), dim=-1)[0].float()
+
+        if mode == "embed":
+            return {"embedding": img_emb.tolist()}
+
+        labels: List[str] = inputs["candidate_labels"]
+
+        with self._autocast():
+            tok = self.tokenizer(labels).to(self.device)
+            txt_emb = F.normalize(self.model.encode_text(tok), dim=-1).float()
+            logits = (self.model.logit_scale.exp() * (img_emb @ txt_emb.T)).float()
+
+        # Standard CLIP zero-shot scoring: logit_scale * (img . txt^T), softmax
+        # over candidates — the recipe from the BioCLIP model card.
+        probs = F.softmax(logits.squeeze(0), dim=0)
+
+        ranked = sorted(zip(labels, probs.tolist()), key=lambda x: -x[1])
+        return [{"label": label, "score": float(score)} for label, score in ranked]