Fix handler

.gitattributes

@@ -1,34 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

handler.py

@@ -1,33 +1,41 @@
 # handler.py
+
 import io
-import torch
 import numpy as np
 from PIL import Image
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+from huggingface_inference_toolkit.handler import BaseHandler
 
-# This will be called once on startup
-def initialize():
-    global mask_generator
-    # HF will mount your model files under /mnt/models
-    checkpoint = "/mnt/models/pytorch_model.bin"
-    sam = sam_model_registry["vit_b"](checkpoint=checkpoint)
-    mask_generator = SamAutomaticMaskGenerator(sam)
+class EndpointHandler(BaseHandler):
+    def __init__(self, model_dir):
+        """
+        Called once at startup. 
+        The model files are mounted under /mnt/models by the Inference Endpoint.
+        """
+        super().__init__(model_dir)
+        checkpoint = "/mnt/models/pytorch_model.bin"
+        sam = sam_model_registry["vit_b"](checkpoint=checkpoint)
+        self.mask_generator = SamAutomaticMaskGenerator(sam)
 
-# This handles each incoming request
-def inference(request):
-    # expect multipart/form-data with field name "image"
-    image_bytes = request.files["image"].read()
-    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    img_np = np.array(img)
-    masks = mask_generator.generate(img_np)
+    async def __call__(self, request):
+        """
+        Called on every HTTP request.
+        Expecting multipart/form-data with an 'image' field.
+        """
+        form = await request.form()
+        image_bytes = form["image"].file.read()
+        
+        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        img_np = np.array(img)
+        
+        masks = self.mask_generator.generate(img_np)
+        combined = np.zeros(img_np.shape[:2], dtype=np.uint8)
+        for m in masks:
+            combined[m["segmentation"]] = 255
 
-    # combine all masks into one binary mask
-    combined = np.zeros(img_np.shape[:2], dtype=np.uint8)
-    for m in masks:
-        combined[m["segmentation"]] = 255
+        out = io.BytesIO()
+        Image.fromarray(combined).save(out, format="PNG")
+        out.seek(0)
 
-    # serialize to PNG
-    out = io.BytesIO()
-    Image.fromarray(combined).save(out, format="PNG")
-    out.seek(0)
-    return {"mask_png": out.read()}
+        # Return a JSON-able dict; binary data will be base64-encoded by the toolkit
+        return {"mask_png": out.getvalue()}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:892c410e496344e527255ccdcb2cb7244a609acb5389c7c4fdba1288f861c579
+size 374979480

requirements.txt

@@ -1,3 +1,4 @@
+# requirements.txt
 torch
 numpy
 Pillow