Robust Marigold output unwrap + deps compatible with HF toolkit

handler.py

@@ -1,6 +1,7 @@
 import io
 import base64
 import inspect
+import dataclasses
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
@@ -11,9 +12,9 @@ import requests
 from diffusers import MarigoldIntrinsicsPipeline, MarigoldNormalsPipeline
 
 
-# -------------------------
+# =======================
 # Image IO helpers
-# -------------------------
+# =======================
 
 def _is_data_url(s: str) -> bool:
     return isinstance(s, str) and s.startswith("data:image/") and ";base64," in s
@@ -45,13 +46,12 @@ def _decode_inputs_to_image(payload_inputs: Union[str, Dict[str, str]]) -> Image
         raise ValueError("inputs must be a string (url/data-url/base64) or a dict with image_url/image_b64")
 
     s = payload_inputs.strip()
-    if s.startswith("http://") or s.startswith("https://"):
+    if s.startswith(("http://", "https://")):
         return _url_to_pil(s)
     if _is_data_url(s):
         s = _strip_data_url_header(s)
         return _b64_to_pil(s)
 
-    # raw base64
     try:
         return _b64_to_pil(s)
     except Exception as e:
@@ -63,14 +63,13 @@ def _png_b64_from_pil(img: Image.Image) -> str:
     return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("utf-8")
 
 
-# -------------------------
+# =======================
 # Tensor helpers
-# -------------------------
+# =======================
 
 def _to_float_tensor(img: Any) -> torch.Tensor:
     """
-    Accept PIL / numpy / torch; return float32 tensor with last dim=3 when possible.
-    Scales from [0,255] to [0,1] if needed.
+    Accept PIL / numpy / torch; return float32 tensor HxWxC in [0,1] when possible.
     """
     if isinstance(img, torch.Tensor):
         t = img
@@ -84,9 +83,10 @@ def _to_float_tensor(img: Any) -> torch.Tensor:
         raise TypeError(f"Unexpected prediction type: {type(img)}")
 
     t = t.float()
-    if t.max() > 1.5:  # likely uint8 0..255
+    if t.max() > 1.5:  # likely 0..255
         t = t / 255.0
-    # If channels-first, move to HWC
+
+    # channels-first -> HWC
     if t.ndim == 3 and t.shape[0] in (1, 3) and t.shape[-1] not in (1, 3):
         t = t.permute(1, 2, 0).contiguous()
     return t
@@ -101,54 +101,76 @@ def _png_b64_from_float_img_neg1to1(img: torch.Tensor) -> str:
     return _png_b64_from_float_img01((img + 1.0) * 0.5)
 
 
-# -------------------------
+# =======================
 # Output unwrapping
-# -------------------------
+# =======================
+
+def _maybe_dataclass_to_dict(obj: Any) -> Optional[Dict[str, Any]]:
+    try:
+        if dataclasses.is_dataclass(obj):
+            return dataclasses.asdict(obj)
+    except Exception:
+        pass
+    # diffusers outputs often have .to_dict()
+    if hasattr(obj, "to_dict") and callable(getattr(obj, "to_dict")):
+        try:
+            return obj.to_dict()
+        except Exception:
+            pass
+    # last resort: use __dict__
+    if hasattr(obj, "__dict__") and isinstance(obj.__dict__, dict):
+        return dict(obj.__dict__)
+    return None
 
 def _unwrap_field(obj: Any, names: List[str]) -> Optional[Any]:
     """
-    Try to extract a named field from:
-      - dataclass-like (attributes)
+    Extract a named field from:
+      - dataclass-like (converted to dict)
       - dict-like
-      - tuple-like via known order
-      - diffusers outputs exposing .images
-    Returns first match or None.
+      - attribute access
+      - diffusers common: .images (first)
+      - tuple/list fallbacks by known order
     """
-    # Attribute access
+    # dataclass / to_dict / __dict__
+    asd = _maybe_dataclass_to_dict(obj)
+    if isinstance(asd, dict):
+        for n in names:
+            if n in asd:
+                return asd[n]
+
+    # attribute access
     for n in names:
         if hasattr(obj, n):
             return getattr(obj, n)
 
-    # Dict-like
+    # dict-like
     if isinstance(obj, dict):
         for n in names:
             if n in obj:
                 return obj[n]
 
-    # Diffusers common: images list
+    # diffusers: images list
     if hasattr(obj, "images"):
         imgs = getattr(obj, "images")
         if isinstance(imgs, (list, tuple)) and len(imgs) > 0:
-            # If caller asked for 'normals'/'albedo' etc but only .images is present,
-            # best effort is to return images[0].
             return imgs[0]
 
-    # Tuple/list fallback for known orders
+    # list/tuple positional guesses
     if isinstance(obj, (list, tuple)):
-        # Intrinsics often: [albedo, diffuse, residual] (or similar)
-        if len(obj) >= 3 and names == ["albedo"]:
+        # Intrinsics often 3 outputs
+        if len(obj) >= 3 and names == ["albedo", "pred_albedo", "albedo_img"]:
             return obj[0]
-        if len(obj) >= 2 and names == ["diffuse_shading", "shading"]:
+        if len(obj) >= 2 and names == ["diffuse_shading", "shading", "pred_shading", "shading_diffuse"]:
             return obj[1]
-        if len(obj) >= 3 and names == ["residual"]:
+        if len(obj) >= 3 and names == ["residual", "pred_residual", "specular_residual"]:
             return obj[2]
 
     return None
 
 
-# -------------------------
+# =======================
 # Endpoint
-# -------------------------
+# =======================
 
 class EndpointHandler:
     def __init__(self, path: Optional[str] = None):
@@ -159,7 +181,7 @@ class EndpointHandler:
         self.normals_pipe: Optional[MarigoldNormalsPipeline] = None
         self.intrin_pipe: Optional[MarigoldIntrinsicsPipeline] = None
 
-        # Config per our discussion
+        # Config
         self.fixed_processing_res = 768
         self.default_intr_steps = 4
 
@@ -167,7 +189,7 @@ class EndpointHandler:
         self.model_id_intrin  = "prs-eth/marigold-iid-lighting-v1-1"
 
     def initialize(self, context: Optional[Dict[str, Any]] = None) -> None:
-        # Lazy load on first use
+        # Lazy load at first call
         pass
 
     # -------- loaders --------
@@ -188,7 +210,6 @@ class EndpointHandler:
     def _run_normals(self, pil_img: Image.Image, processing_res: int) -> Dict[str, Any]:
         self._ensure_normals_loaded()
 
-        # Only pass processing_res if supported
         sig = inspect.signature(self.normals_pipe.__call__)
         kwargs = {"image": pil_img}
         if "processing_res" in sig.parameters:
@@ -196,18 +217,25 @@ class EndpointHandler:
 
         out = self.normals_pipe(**kwargs)
 
-        # Robust extraction
-        normals = _unwrap_field(out, ["normals"])
+        # Try common field names used across versions
+        normals = _unwrap_field(out, ["normals", "pred_normals", "normal_map"])
         if normals is None:
-            # Some builds expose a single image via .images[0]
-            normals = _unwrap_field(out, ["images"])  # will give images[0] in our helper
-
+            # if still not found, see if the object itself is an image/tensor
+            asd = _maybe_dataclass_to_dict(out)
+            if isinstance(asd, dict):
+                # pick first image-like thing from the dict
+                for v in asd.values():
+                    try:
+                        _ = _to_float_tensor(v)
+                        normals = v
+                        break
+                    except Exception:
+                        continue
         if normals is None:
-            # Last resort: maybe it's a raw tensor/ndarray/PIL already
-            normals = out
+            # Last resort: the entire object is not directly convertible; fail clearly
+            raise TypeError(f"Unexpected prediction type: {type(out)}")
 
-        normals_t = _to_float_tensor(normals)  # expect [-1,1] or [0,1]
-        # Heuristic: if range looks 0..1, map to [-1,1] visualization anyway
+        normals_t = _to_float_tensor(normals)
         maxv, minv = float(normals_t.max().item()), float(normals_t.min().item())
         if maxv <= 1.01 and minv >= -0.01:
             vis_b64 = _png_b64_from_float_img01(normals_t)
@@ -229,12 +257,20 @@ class EndpointHandler:
 
         out = self.intrin_pipe(image=pil_img, num_inference_steps=int(steps))
 
-        # Try attributes/dict/tuple/tensor in a robust way
-        albedo   = _unwrap_field(out, ["albedo"])
-        diffuse  = _unwrap_field(out, ["diffuse_shading", "shading"])
-        residual = _unwrap_field(out, ["residual"])
+        # Keys that show up across versions
+        albedo   = _unwrap_field(out, ["albedo", "pred_albedo", "albedo_img"])
+        diffuse  = _unwrap_field(out, ["diffuse_shading", "shading", "pred_shading", "shading_diffuse"])
+        residual = _unwrap_field(out, ["residual", "pred_residual", "specular_residual"])
+
+        # If .images[0..2] only:
+        if any(v is None for v in (albedo, diffuse, residual)) and hasattr(out, "images"):
+            imgs = getattr(out, "images")
+            if isinstance(imgs, (list, tuple)) and len(imgs) >= 3:
+                albedo   = imgs[0] if albedo   is None else albedo
+                diffuse  = imgs[1] if diffuse  is None else diffuse
+                residual = imgs[2] if residual is None else residual
 
-        # Tensor stack fallback (e.g., [3,H,W,3] or similar)
+        # Tensor stack fallback
         if any(v is None for v in (albedo, diffuse, residual)):
             try:
                 t = _to_float_tensor(out)

requirements.txt

@@ -1,9 +1,8 @@
 torch>=2.1
 accelerate>=0.28
-diffusers==0.35.2
+diffusers==0.33.1
 transformers>=4.44
 peft>=0.17.0
-huggingface_hub[hf_transfer]==0.30.2
 Pillow>=10.3
 safetensors>=0.4.2
 numpy>=1.26