Fix: remove unsupported target_properties kwarg; add previews; robust inputs

handler.py

@@ -1,153 +1,124 @@
 import base64
 import io
+import json
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Tuple, Optional
 
-import numpy as np
-from PIL import Image
 import torch
-from diffusers import MarigoldIntrinsicsPipeline
-
-
-def _b64_to_pil(b64: str) -> Image.Image:
-    """Decode base64 (optionally data URL) → PIL RGB."""
-    s = b64.strip()
-    if s.startswith("data:image"):
-        s = s.split(",", 1)[1]
-    raw = base64.b64decode(s, validate=False)
-    return Image.open(io.BytesIO(raw)).convert("RGB")
-
-
-def _pil_to_b64_png(img: Image.Image) -> str:
-    """PIL → base64 PNG (no data URL prefix)."""
-    buf = io.BytesIO()
-    img.save(buf, format="PNG")
-    return base64.b64encode(buf.getvalue()).decode("utf-8")
+from PIL import Image
+from diffusers import MarigoldIntrinsicsPipeline, utils as diffusers_utils
 
 
 class EndpointHandler:
     """
-    Hugging Face Inference Toolkit-compatible handler.
-
-    Expects request JSON:
-      {
-        "inputs": "<base64 image or data URL>",
-        "parameters": {
-          "processing_res": 768,
-          "targets": ["albedo","diffuse_shading","residual"],  # optional subset
-          "num_inference_steps": <int>,                        # optional
-          "match_input_resolution": true                       # default true
-        }
-      }
-
-    Returns:
-      {
-        "outputs": {
-          "targets": [...],
-          "tensors_info": {"shape":[H,W,C], "dtype":"float32"},
-          "uncertainty_available": <bool>,
-          "visualizations": ["<b64_png_albedo>", "<b64_png_diffuse>", "<b64_png_residual>"],
-          "visualizations_error": null | "<msg>"
-        }
-      }
+    Hugging Face Inference Endpoints custom handler for Marigold Intrinsics (Lighting).
     """
 
-    def __init__(self, model_dir: Optional[str] = None):
-        # The toolkit passes model_dir but we load by model id to ensure the correct subtask.
-        self.model_id = os.environ.get("MARIGOLD_MODEL_ID", "prs-eth/marigold-iid-lighting-v1-1")
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
-        self.pipe: Optional[MarigoldIntrinsicsPipeline] = None
-        self.supported_targets: List[str] = ["albedo", "diffuse_shading", "residual"]
-
-    def setup(self):
-        # Load once
-        self.pipe = MarigoldIntrinsicsPipeline.from_pretrained(self.model_id, torch_dtype=self.dtype)
-        self.pipe = self.pipe.to(self.device)
-        # figure out supported target names from the checkpoint if provided
-        try:
-            if getattr(self.pipe, "target_properties", None):
-                names = self.pipe.target_properties.get("target_names")
-                if isinstance(names, list) and names:
-                    self.supported_targets = names
-        except Exception:
-            pass
-        try:
-            self.pipe.set_progress_bar_config(disable=True)
-        except Exception:
-            pass
-
+    def __init__(self, path: str = "."):
+        # Load the Lighting checkpoint
+        self.pipe = MarigoldIntrinsicsPipeline.from_pretrained(
+            "prs-eth/marigold-iid-lighting-v1-1",
+            variant="fp16",
+            torch_dtype=torch.float16,
+        ).to("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Optional: reduced steps (user can override via inputs)
+        self.default_steps = 4
+
+    # ------- helpers --------
+    def _load_image_from_any(self, obj: Any) -> Image.Image:
+        """
+        Accepts:
+          - str URL (http/https)
+          - base64 string (png/jpg)
+          - dict {"image_url": "..."}
+        Returns PIL.Image
+        """
+        if isinstance(obj, dict):
+            if "image_url" in obj:
+                return diffusers_utils.load_image(obj["image_url"])
+            # if someone sent {"image_b64": "..."}:
+            if "image_b64" in obj:
+                return Image.open(io.BytesIO(base64.b64decode(obj["image_b64"]))).convert("RGB")
+            raise ValueError("Unsupported dict format for 'inputs'. Provide 'image_url' or 'image_b64'.")
+
+        if isinstance(obj, str):
+            s = obj.strip()
+            if s.startswith("http://") or s.startswith("https://"):
+                return diffusers_utils.load_image(s)
+            # else assume base64 image
+            return Image.open(io.BytesIO(base64.b64decode(s))).convert("RGB")
+
+        raise ValueError("Unsupported 'inputs' type. Provide URL string, base64 string, or {'image_url': ...} dict.")
+
+    def _png_b64(self, pil_img: Image.Image) -> str:
+        buf = io.BytesIO()
+        pil_img.save(buf, format="PNG")
+        return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+    # ------- main entry --------
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        if self.pipe is None:
-            self.setup()
+        # Inference Endpoints expects the request JSON to have an "inputs" key.
+        if "inputs" not in data:
+            raise ValueError("Body must include an 'inputs' key")
 
-        # ---- Parse inputs ----
-        raw_inputs = data.get("inputs")
-        params = data.get("parameters") or {}
+        # Parse options (all optional)
+        params = data if isinstance(data, dict) else {}
+        num_steps = int(params.get("num_inference_steps", self.default_steps))
 
-        if not isinstance(raw_inputs, (str, bytes, bytearray)):
-            return {"error": "Send the image in 'inputs' as base64 string (or data URL)."}
+        # Note: do NOT pass unsupported kwargs like target_properties or processing_res
+        image = self._load_image_from_any(data["inputs"])
 
-        # decode base64 image
-        try:
-            if isinstance(raw_inputs, (bytes, bytearray)):
-                raw_inputs = raw_inputs.decode("utf-8")
-            image = _b64_to_pil(raw_inputs)
-        except Exception as e:
-            return {"error": f"Failed to decode 'inputs': {e}"}
-
-        processing_res = int(params.get("processing_res", getattr(self.pipe, "default_processing_resolution", 768) or 768))
-        num_steps = params.get("num_inference_steps", None)
-        match_input_resolution = bool(params.get("match_input_resolution", True))
-
-        # requested targets, keep order, intersect with supported
-        req_targets = params.get("targets", self.supported_targets)
-        targets = [t for t in req_targets if t in self.supported_targets]
-        if not targets:
-            targets = self.supported_targets
-
-        # target_properties override so the pipeline outputs only what we want, in order
-        target_props = dict(self.pipe.target_properties) if getattr(self.pipe, "target_properties", None) else {}
-        target_props["target_names"] = targets
-
-        # ---- Inference ----
         with torch.inference_mode():
-            out = self.pipe(
-                image=image,
-                processing_resolution=processing_res,
-                match_input_resolution=match_input_resolution,
+            intrinsics = self.pipe(
+                image,
                 num_inference_steps=num_steps,
-                target_properties=target_props,
-                output_type="np",
-                return_dict=True,
+                # If you later want to expose these (and they’re supported),
+                # you can add conditionals:
+                # processing_resolution=params.get("processing_resolution", None),
+                # match_input_resolution=params.get("match_input_resolution", None),
             )
 
-        pred = out.prediction  # numpy HxWxC float32
-        h, w, c = pred.shape
-        dtype_str = str(pred.dtype)
-        unc_avail = bool(getattr(out, "uncertainty", None) is not None)
+        # intrinsics.prediction is a tensor; lighting model yields 3 targets
+        pred = intrinsics.prediction  # shape: (T, H, W, 3) typically; T=3 for lighting
+        if isinstance(pred, torch.Tensor):
+            t = pred
+        else:
+            # Just in case (diffusers usually returns a dataclass with .prediction tensor)
+            raise RuntimeError("Unexpected intrinsics.prediction type")
 
-        # ---- Visualize ----
-        vis_b64 = None
-        vis_err = None
+        # Gather target names from pipeline properties
+        # For lighting these should be: ["albedo", "shading", "residual"]
         try:
-            # Uses the built-in visualizer for intrinsics
-            from diffusers.pipelines.marigold.image_processor import MarigoldImageProcessor
-            vis_imgs = MarigoldImageProcessor.visualize_intrinsics(
-                prediction=pred, target_properties=target_props, color_map="binary"
-            )
-            # to base64 PNGs (no data URL prefix to keep payload compact)
-            vis_b64 = [_pil_to_b64_png(im) for im in vis_imgs]
-        except Exception as e:
-            vis_b64 = None
-            vis_err = f"Failed to visualize: {e}"
+            target_names = list(self.pipe.target_properties.keys())
+        except Exception:
+            # Fallback (expected order for lighting model)
+            target_names = ["albedo", "shading", "residual"]
 
-        return {
+        out: Dict[str, Any] = {
             "outputs": {
-                "targets": targets,
-                "tensors_info": {"shape": [h, w, c], "dtype": dtype_str},
-                "uncertainty_available": unc_avail,
-                "visualizations": vis_b64,
-                "visualizations_error": vis_err,
+                "targets": target_names,
+                "tensors_info": {
+                    "shape": list(t.shape),
+                    "dtype": str(t.dtype).replace("torch.", ""),
+                },
+                "uncertainty_available": False,
             }
         }
+
+        # Try to attach base64 previews per target
+        try:
+            vis_list = self.pipe.image_processor.visualize_intrinsics(t, self.pipe.target_properties)
+            # vis_list is a list over batch (we run 1 image), where each item is a dict of PIL images per target
+            previews: Dict[str, str] = {}
+            if isinstance(vis_list, (list, tuple)) and len(vis_list) >= 1 and isinstance(vis_list[0], dict):
+                for name, pil_img in vis_list[0].items():
+                    # only include known target previews to avoid surprises
+                    if name in target_names and isinstance(pil_img, Image.Image):
+                        previews[name] = self._png_b64(pil_img)
+            if previews:
+                out["outputs"]["visualizations_png_b64"] = previews
+        except Exception as e:
+            out["outputs"]["visualizations_error"] = f"Failed to visualize: {e.__class__.__name__}: {e}"
+
+        return out