Update handler.py

handler.py

@@ -1,14 +1,12 @@
-# handler.py  (repo root)
 import io, base64, torch, open_clip
 from PIL import Image
-# optional: from open_clip import fuse_conv_bn_sequential   # if you want re‑param
 
 class EndpointHandler:
     """
-    MobileCLIP‑B ('datacompdr') zero‑shot classifier with per‑process
+    MobileCLIP‑B ('datacompdr') zero‑shot classifier with a per‑process
     text‑embedding cache.
 
-    Expected client JSON:
+    Client JSON must look like:
       {
         "inputs": {
           "image": "<base64 PNG/JPEG>",
@@ -17,13 +15,11 @@ class EndpointHandler:
       }
     """
 
+    # ---------- init (runs once per container) ----------
     def __init__(self, path=""):
-        # Load the exact weights your local run uses
         self.model, _, self.preprocess = open_clip.create_model_and_transforms(
             "mobileclip_b", pretrained="datacompdr"
         )
-        # Optional: fuse conv+bn for speed
-        # self.model = fuse_conv_bn_sequential(self.model).eval()
         self.model.eval()
 
         self.tokenizer = open_clip.get_tokenizer("mobileclip_b")
@@ -32,6 +28,7 @@ class EndpointHandler:
 
         self.cache: dict[str, torch.Tensor] = {}   # prompt → embedding
 
+    # ----------------- inference ------------------------
     def __call__(self, data):
         payload = data.get("inputs", data)
         img_b64 = payload["image"]
@@ -43,7 +40,7 @@ class EndpointHandler:
         img = Image.open(io.BytesIO(base64.b64decode(img_b64))).convert("RGB")
         img_t = self.preprocess(img).unsqueeze(0).to(self.device)
 
-        # Text embeddings with cache
+        # Text embeddings (cached)
         new = [l for l in labels if l not in self.cache]
         if new:
             tok = self.tokenizer(new).to(self.device)
@@ -65,6 +62,7 @@ class EndpointHandler:
             for l, p in sorted(zip(labels, probs), key=lambda x: x[1], reverse=True)
         ]
 
+
 # import io, base64, torch
 # from PIL import Image
 # import open_clip