Upload 2 files

83d47de verified 4 months ago

5.59 kB

"""
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

	"""
	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 (alphaimg + (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