Return embedding bin file

app.py

@@ -5,23 +5,28 @@ from handler import EndpointHandler
 # 初始化你的 ONNX 模型
 handler = EndpointHandler(path=".")
 
-def api_call(image, parameters):
+def api_call(image, return_embeddings_only):
     data = {
         "inputs": image,      # PIL.Image
-        "parameters": parameters or {}
+        "parameters": {"return_embeddings_only": return_embeddings_only}
     }
     result = handler(data)[0]
+
+    # If embeddings file is returned
+    if "file" in result:
+        return result["file"]
+
     return result  # JSON only（含 base64 mask）
 
 demo = gr.Interface(
     fn=api_call,
     inputs=[
         gr.Image(type="pil", label="image"),
-        gr.JSON(label="parameters (point_coords, point_labels, return_mask_image)")
+        gr.Checkbox(label="Return embeddings only (as .bin file)", value=True)
     ],
-    outputs=gr.JSON(label="result"),
-    title="Edge SAM API",
-    description="Pure API Space — send image + parameters, get mask JSON."
+    outputs=[gr.File(label="embeddings.bin or result JSON")],
+    title="Edge SAM Encoder API",
+    description="Send image, get embeddings in float32 binary format (.bin file)."
 )
 
 if __name__ == "__main__":

handler.py

@@ -49,6 +49,14 @@ class EndpointHandler:
             # Encode
             embeddings = self.encoder.run(None, {'image': img_array})[0]
 
+            # Check if only embeddings are requested
+            if params.get("return_embeddings_only", False):
+                # Convert embeddings to float32 binary and save to temp file
+                embeddings_float32 = embeddings.astype(np.float32)
+                temp_file = "/tmp/embeddings.bin"
+                embeddings_float32.tofile(temp_file)
+                return [{"file": temp_file, "shape": list(embeddings.shape), "dtype": "float32"}]
+
             # Prepare prompts
             coords = np.array(params.get("point_coords", [[512, 512]]), dtype=np.float32)
             labels = np.array(params.get("point_labels", [1]), dtype=np.float32)