Upload 2 files

handler.py

@@ -0,0 +1,140 @@
+"""
+Custom inference handler for Hugging Face Inference Endpoints.
+
+Serves the Marqo GCL ViT-L-14 open_clip checkpoint as a pure encoder:
+given text and/or image(s), it returns L2-normalized embeddings, exactly
+matching the notebook's encode_text / encode_image + normalize steps.
+
+Request  (JSON):
+    { "inputs": { "text":  ["a red dress", ...],          # optional
+                  "image": ["<base64-encoded-image>", ...] } }  # optional
+
+Response (JSON):
+    { "text_embeddings":  [[... 768 floats ...], ...],    # if text sent
+      "image_embeddings": [[... 768 floats ...], ...] }   # if image sent
+
+Fusion (alpha*img + (1-alpha)*txt) is intentionally NOT done here — that
+happens in the n8n Code node, so this endpoint stays a clean encoder.
+"""
+
+import os
+import io
+import glob
+import base64
+
+import torch
+import open_clip
+from PIL import Image
+
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        # `path` is the local directory where the repo (including the .pt) is mounted.
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Architecture of the checkpoint. ViT-L-14 for the Marqo GCL model.
+        arch = os.environ.get("OPEN_CLIP_ARCH", "ViT-L-14")
+
+        # Build the architecture WITHOUT pretrained weights, then load the GCL checkpoint.
+        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
+            arch, pretrained=None
+        )
+        self.tokenizer = open_clip.get_tokenizer(arch)
+
+        ckpt_path = self._find_checkpoint(path)
+        self._load_weights(ckpt_path)
+
+        self.model = self.model.to(self.device).eval()
+
+    # ---- locate the checkpoint file inside the repo ----
+    def _find_checkpoint(self, path: str) -> str:
+        # Allow an explicit override via env var (filename or absolute path).
+        env = os.environ.get("CHECKPOINT_FILE")
+        if env:
+            return env if os.path.isabs(env) else os.path.join(path, env)
+
+        candidates = (
+            glob.glob(os.path.join(path, "*.pt"))
+            + glob.glob(os.path.join(path, "*.bin"))
+            + glob.glob(os.path.join(path, "*.safetensors"))
+        )
+        if not candidates:
+            raise FileNotFoundError(
+                f"No .pt/.bin/.safetensors checkpoint found in '{path}'. "
+                "Set CHECKPOINT_FILE env var to the filename."
+            )
+        # If several exist, prefer the largest (the model weights, not a config shard).
+        return max(candidates, key=os.path.getsize)
+
+    # ---- load weights, tolerating the 'full_states' training-checkpoint layout ----
+    def _load_weights(self, ckpt_path: str) -> None:
+        # The Marqo file is named "...full_states.pt": a full training state that
+        # wraps the model weights under a 'state_dict' key and may prefix keys with
+        # 'module.' (from DDP training). open_clip.load_checkpoint normalizes both.
+        try:
+            open_clip.load_checkpoint(self.model, ckpt_path, strict=False)
+            return
+        except Exception as e:
+            print(f"[handler] open_clip.load_checkpoint fallback triggered: {e}")
+
+        # Manual fallback. weights_only=False is required because a full training
+        # state contains non-tensor objects (optimizer state, epoch, etc.).
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+
+        if isinstance(ckpt, dict):
+            state_dict = ckpt.get("state_dict", ckpt.get("model", ckpt))
+        else:
+            state_dict = ckpt
+
+        # Strip a leading 'module.' from DDP-saved keys.
+        state_dict = {
+            (k[len("module."):] if k.startswith("module.") else k): v
+            for k, v in state_dict.items()
+        }
+
+        missing, unexpected = self.model.load_state_dict(state_dict, strict=False)
+        if missing:
+            print(f"[handler] missing keys (first 5): {list(missing)[:5]}")
+        if unexpected:
+            print(f"[handler] unexpected keys (first 5): {list(unexpected)[:5]}")
+
+    @torch.no_grad()
+    def __call__(self, data: dict) -> dict:
+        # HF passes {"inputs": ..., "parameters": ...}; be lenient about shape.
+        inputs = data.get("inputs", data) or {}
+        if not isinstance(inputs, dict):
+            return {"error": "Send an object under 'inputs' with 'text' and/or 'image'."}
+
+        out = {}
+
+        # ---- text embeddings ----
+        texts = inputs.get("text")
+        if texts:
+            if isinstance(texts, str):
+                texts = [texts]
+            tokens = self.tokenizer(texts).to(self.device)
+            feats = self.model.encode_text(tokens)
+            feats = feats / feats.norm(dim=-1, keepdim=True)
+            out["text_embeddings"] = feats.cpu().tolist()
+
+        # ---- image embeddings (base64-encoded bytes) ----
+        images = inputs.get("image")
+        if images:
+            if isinstance(images, str):
+                images = [images]
+            tensors = []
+            for b64 in images:
+                raw = base64.b64decode(b64)
+                img = Image.open(io.BytesIO(raw)).convert("RGB")
+                tensors.append(self.preprocess(img))
+            batch = torch.stack(tensors).to(self.device)
+            feats = self.model.encode_image(batch)
+            feats = feats / feats.norm(dim=-1, keepdim=True)
+            out["image_embeddings"] = feats.cpu().tolist()
+
+        if not out:
+            return {
+                "error": "Provide 'text' and/or 'image' (base64) under 'inputs'.",
+                "example": {"inputs": {"text": ["a red dress"]}},
+            }
+        return out

requirements.txt

@@ -0,0 +1,2 @@
+open_clip_torch>=2.24.0
+pillow