add handler

handler.py

@@ -1,38 +1,37 @@
-from typing import Dict, List
-import torch, base64, io
+import torch, open_clip
 from PIL import Image
-import open_clip
+from typing import Any, Dict
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model, _, preprocess = open_clip.create_model_and_transforms(
-    'ViT-B-32', pretrained='laion2b_s34b_b79K', device=device
-)
+class EndpointHandler:
+    def __init__(self, model_dir: str):
+        self.device = "cpu"
+        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
+            "ViT-B-32", pretrained="laion2b_s34b_b79K", device=self.device
+        )
+        self.tokenizer = open_clip.get_tokenizer("ViT-B-32")
 
-def _embed_image(img_b64: str) -> List[float]:
-    img = Image.open(io.BytesIO(base64.b64decode(img_b64))).convert("RGB")
-    tensor = preprocess(img).unsqueeze(0).to(device)
-    with torch.no_grad():
-        emb = model.encode_image(tensor)
-    return emb.squeeze().cpu().tolist()
+    def _encode_text(self, text: str):
+        tokens = self.tokenizer([text]).to(self.device)
+        with torch.no_grad():
+            return self.model.encode_text(tokens).cpu().numpy()[0].tolist()
 
-def _embed_text(text: str) -> List[float]:
-    tok = open_clip.tokenize([text]).to(device)
-    with torch.no_grad():
-        emb = model.encode_text(tok)
-    return emb.squeeze().cpu().tolist()
+    def _encode_image(self, image: Image.Image):
+        img = self.preprocess(image).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            return self.model.encode_image(img).cpu().numpy()[0].tolist()
 
-# === HF endpoint entrypoint ===
-def preprocess(payload: Dict):
-    return payload
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        if "image" in data:
+            if isinstance(data["image"], str):
+                import requests, io
+                resp = requests.get(data["image"])
+                img = Image.open(io.BytesIO(resp.content)).convert("RGB")
+            else:
+                img = Image.open(data["image"]).convert("RGB")
+            emb = self._encode_image(img)
+        elif "inputs" in data:
+            emb = self._encode_text(data["inputs"])
+        else:
+            raise ValueError("Provide 'image' or 'inputs'.")
+        return {"embedding": emb}
 
-def inference(payload: Dict):
-    if isinstance(payload, str) and payload.startswith("data:image"):
-        b64 = payload.split(",")[-1]
-        return {"vector": _embed_image(b64)}
-    elif isinstance(payload, str):
-        return {"vector": _embed_text(payload)}
-    else:
-        raise ValueError("Unsupported input")
-
-def postprocess(output):  # HF expects this even se passas direto
-    return output