Add debug logging for payload tracing

handler.py

@@ -1,4 +1,5 @@
-# Fixed handler for BiRefNet endpoint — now supports base64 + URLs + file paths
+# handler.py — BiRefNet endpoint handler
+# Fully instrumented for debugging input structure and format.
 
 from typing import Dict, Any, Tuple
 import os
@@ -12,11 +13,11 @@ import torch
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
 
-torch.set_float32_matmul_precision(["high", "highest"][0])
+torch.set_float32_matmul_precision("high")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # ======================================================
-# Utility Functions
+# Utility functions
 # ======================================================
 def refine_foreground(image, mask, r=90):
     if mask.size != image.size:
@@ -24,16 +25,13 @@ def refine_foreground(image, mask, r=90):
     image = np.array(image) / 255.0
     mask = np.array(mask) / 255.0
     estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r)
-    image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
-    return image_masked
-
+    return Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
 
 def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
     alpha = alpha[:, :, None]
     F, blur_B = FB_blur_fusion_foreground_estimator(image, image, image, alpha, r)
     return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]
 
-
 def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
     if isinstance(image, Image.Image):
         image = np.array(image) / 255.0
@@ -43,15 +41,13 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
     blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))
     blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
     F = blurred_F + alpha * (image - alpha * blurred_F - (1 - alpha) * blurred_B)
-    F = np.clip(F, 0, 1)
-    return F, blurred_B
-
+    return np.clip(F, 0, 1), blurred_B
 
 # ======================================================
 # Preprocessing
 # ======================================================
-class ImagePreprocessor():
-    def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
+class ImagePreprocessor:
+    def __init__(self, resolution: Tuple[int, int] = (1024, 1024)):
         self.transform_image = transforms.Compose([
             transforms.Resize(resolution),
             transforms.ToTensor(),
@@ -61,7 +57,6 @@ class ImagePreprocessor():
     def proc(self, image: Image.Image) -> torch.Tensor:
         return self.transform_image(image)
 
-
 # ======================================================
 # Model and Endpoint
 # ======================================================
@@ -81,68 +76,79 @@ usage_to_weights_file = {
     'General-legacy': 'BiRefNet-legacy'
 }
 
-usage = 'General'
-if usage in ['General-Lite-2K']:
-    resolution = (2560, 1440)
-elif usage in ['General-reso_512']:
-    resolution = (512, 512)
-elif usage in ['General-HR', 'Matting-HR']:
-    resolution = (2048, 2048)
-else:
-    resolution = (1024, 1024)
-
+usage = "General"
+resolution = (1024, 1024)
 half_precision = True
 
-
-class EndpointHandler():
-    def __init__(self, path=''):
+# ======================================================
+# Endpoint Handler
+# ======================================================
+class EndpointHandler:
+    def __init__(self, path=""):
         self.birefnet = AutoModelForImageSegmentation.from_pretrained(
-            '/'.join(('zhengpeng7', usage_to_weights_file[usage])),
+            f"zhengpeng7/{usage_to_weights_file[usage]}",
             trust_remote_code=True
         )
-        self.birefnet.to(device)
-        self.birefnet.eval()
+        self.birefnet.to(device).eval()
         if half_precision:
             self.birefnet.half()
         print("✅ BiRefNet model loaded successfully.")
 
     def __call__(self, data: Dict[str, Any]):
-        """
-        Accepts either:
-          - URL (http:// or https://)
-          - Base64 (raw or data:image/...;base64,...)
-          - File path
-        """
         image_src = data.get("inputs")
+
+        # ================= DEBUG LOGS =================
+        print("\n==============================")
+        print("🧩 DEBUG: Incoming data structure")
+        print(f"Type of data: {type(data)}")
+        print(f"Keys: {list(data.keys()) if isinstance(data, dict) else 'N/A'}")
+        print(f"Type of inputs: {type(image_src)}")
+        if isinstance(image_src, str):
+            print(f"  Length: {len(image_src)}")
+            print(f"  Starts with: {repr(image_src[:120])}")
+        elif isinstance(image_src, bytes):
+            print(f"  Bytes length: {len(image_src)}")
+        else:
+            print(f"  Value preview: {repr(image_src)[:200]}")
+        print("==============================\n", flush=True)
+        # ===============================================
+
         if image_src is None:
             raise ValueError("Missing 'inputs' key in request payload")
 
-        # ✅ Handle base64 or data URI
-        if isinstance(image_src, str):
-            if image_src.startswith("data:image"):
-                header, b64data = image_src.split(",", 1)
-                image_ori = Image.open(io.BytesIO(base64.b64decode(b64data)))
-            elif image_src[:4] in ("/9j/", "iVBOR", "R0lG", "UklG"):
-                image_ori = Image.open(io.BytesIO(base64.b64decode(image_src)))
-            elif image_src.startswith("http"):
-                response = requests.get(image_src)
-                image_ori = Image.open(io.BytesIO(response.content))
-            elif os.path.isfile(image_src):
-                image_ori = Image.open(image_src)
+        # ✅ Decode base64 / data URI / URL / file path
+        try:
+            if isinstance(image_src, (bytes, bytearray)):
+                image_ori = Image.open(io.BytesIO(image_src))
+            elif isinstance(image_src, str):
+                image_src = image_src.strip()
+
+                if image_src.startswith("data:image"):
+                    header, b64data = image_src.split(",", 1)
+                    image_bytes = base64.b64decode(b64data)
+                    image_ori = Image.open(io.BytesIO(image_bytes))
+                elif any(image_src.startswith(pfx) for pfx in ("iVBOR", "/9j/", "R0lG", "UklG")):
+                    image_bytes = base64.b64decode(image_src)
+                    image_ori = Image.open(io.BytesIO(image_bytes))
+                elif image_src.startswith("http"):
+                    response = requests.get(image_src)
+                    image_ori = Image.open(io.BytesIO(response.content))
+                elif os.path.isfile(image_src):
+                    image_ori = Image.open(image_src)
+                else:
+                    raise ValueError(f"Unsupported input string format: {image_src[:40]}...")
             else:
-                raise ValueError("Unsupported input string format.")
-        else:
-            # Assume it's an array-like
-            image_ori = Image.fromarray(image_src)
+                image_ori = Image.fromarray(np.array(image_src))
+
+        except Exception as e:
+            print(f"❌ ERROR decoding input: {e}")
+            raise
 
-        image = image_ori.convert('RGB')
+        image = image_ori.convert("RGB")
 
-        # Preprocess
-        image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
-        image_proc = image_preprocessor.proc(image)
-        image_proc = image_proc.unsqueeze(0)
+        image_preprocessor = ImagePreprocessor(resolution=resolution)
+        image_proc = image_preprocessor.proc(image).unsqueeze(0)
 
-        # Predict
         with torch.no_grad():
             preds = self.birefnet(
                 image_proc.to(device).half() if half_precision else image_proc.to(device)
@@ -154,7 +160,6 @@ class EndpointHandler():
         image_masked = refine_foreground(image, pred_pil)
         image_masked.putalpha(pred_pil.resize(image.size))
 
-        # Return as base64 for easy JSON transport
         buffer = io.BytesIO()
         image_masked.save(buffer, format="PNG")
         encoded_result = base64.b64encode(buffer.getvalue()).decode("utf-8")