Initial commit: Marigold Lighting handler + requirements

.gitattributes

@@ -1,35 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.pt 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
-*.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

README.md

@@ -0,0 +1,39 @@
+---
+title: Marigold Intrinsic Image Decomposition
+emoji: 🏵️
+colorFrom: purple
+colorTo: blue
+sdk: gradio
+sdk_version: 4.44.1
+app_file: app.py
+pinned: true
+license: apache-2.0
+models:
+- prs-eth/marigold-iid-appearance-v1-1
+- prs-eth/marigold-iid-lighting-v1-1
+---
+
+This is a demo of Marigold-IID, the state-of-the-art intrinsic image decomposition model for images in the wild.
+We provide two models: 
+- Marigold-IID-Appearance which predicts albedo, metallic and roughness 
+- Marigold-IID-Lighting which predicts albedo, diffuse shading and non-diffuse residual
+
+Find out more in our paper titled ["Marigold: Affordable Adaptation of Diffusion-Based Image Generators for Image Analysis"](https://arxiv.org/abs/2505.09358)
+
+```
+@InProceedings{ke2023repurposing,
+  title={Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation},
+  author={Bingxin Ke and Anton Obukhov and Shengyu Huang and Nando Metzger and Rodrigo Caye Daudt and Konrad Schindler},
+  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+  year={2024}
+}
+
+@misc{ke2025marigold,
+  title={Marigold: Affordable Adaptation of Diffusion-Based Image Generators for Image Analysis},
+  author={Bingxin Ke and Kevin Qu and Tianfu Wang and Nando Metzger and Shengyu Huang and Bo Li and Anton Obukhov and Konrad Schindler},
+  year={2025},
+  eprint={2505.09358},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV}
+}
+```

handler.py

@@ -0,0 +1,169 @@
+import base64, io, os
+from typing import Any, Dict, List, Union
+
+import torch
+import numpy as np
+from PIL import Image
+
+from diffusers import MarigoldIntrinsicsPipeline
+from diffusers.utils import load_image
+
+
+def _decode_image(inp: Union[str, bytes]) -> Image.Image:
+    """
+    Accepts:
+      - https/http URL
+      - data: URI or bare base64 string of an image
+    """
+    if isinstance(inp, bytes):
+        return Image.open(io.BytesIO(inp)).convert("RGB")
+
+    if isinstance(inp, str) and (inp.startswith("http://") or inp.startswith("https://")):
+        return load_image(inp).convert("RGB")
+
+    # base64 (with or without data: prefix)
+    if isinstance(inp, str) and "," in inp and inp.strip().lower().startswith("data:"):
+        inp = inp.split(",", 1)[1]
+    data = base64.b64decode(inp)
+    return Image.open(io.BytesIO(data)).convert("RGB")
+
+
+def _b64_png(img: Image.Image) -> str:
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+
+class EndpointHandler:
+    """
+    Hugging Face Inference Endpoints entrypoint.
+    Loads Marigold *Intrinsic Image Decomposition* (Lighting) and exposes a JSON API.
+
+    Request JSON:
+      {
+        "inputs": "<image-url OR base64>",
+        "parameters": {
+          "num_inference_steps": 4,           # optional; DDIM or LCM schedules handle 1–50 (docs)
+          "processing_resolution": 768,       # optional; sensible default comes from model
+          "match_input_resolution": True,     # optional
+          "ensemble_size": 1,                 # >1 enables 'uncertainty'
+          "return_visualizations": true       # if true, returns PNGs for albedo/shading/residual
+        }
+      }
+
+    Response JSON:
+      {
+        "outputs": {
+          "targets": ["albedo", "diffuse_shading", "residual"],
+          "visualizations": { "albedo": "<b64png>", ... },     # if requested
+          "tensors_info": { "shape": [N, 3, H, W], "dtype": "float16/float32" },
+          "uncertainty_available": true/false
+        }
+      }
+    """
+
+    def __init__(self, model_path: str = ""):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Load the Lighting (IID) checkpoint with fp16 for speed, as in the docs.
+        # Model card: prs-eth/marigold-iid-lighting-v1-1
+        # https://huggingface.co/prs-eth/marigold-iid-lighting-v1-1
+        repo_id = os.environ.get("MARIGOLD_REPO_ID", "prs-eth/marigold-iid-lighting-v1-1")
+
+        self.pipe = MarigoldIntrinsicsPipeline.from_pretrained(
+            repo_id, variant="fp16", torch_dtype=torch.float16
+        ).to(self.device)
+
+        # Optional: speed tweak recommended by docs — tiny VAE
+        # (works across Marigold pipelines; comment out if you prefer default VAE)
+        try:
+            from diffusers import AutoencoderTiny
+            self.pipe.vae = AutoencoderTiny.from_pretrained(
+                "madebyollin/taesd", torch_dtype=torch.float16
+            ).to(self.device)
+        except Exception as _:
+            pass
+
+        # Disable progress bars in server context
+        self.pipe.set_progress_bar_config(disable=True)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        if not isinstance(data, dict) or "inputs" not in data:
+            return {"error": "Expected JSON with an 'inputs' field."}
+
+        params = data.get("parameters", {}) or {}
+        return_visuals = bool(params.get("return_visualizations", True))
+
+        # Map supported pipeline kwargs exactly as in diffusers:
+        # - num_inference_steps
+        # - processing_resolution
+        # - match_input_resolution
+        # - ensemble_size
+        # - generator (omitted here; could be added if you pass a seed)
+        call_kwargs = {}
+        for k in (
+            "num_inference_steps",
+            "processing_resolution",
+            "match_input_resolution",
+            "ensemble_size",
+            "output_uncertainty",
+            "output_latent",
+        ):
+            if k in params:
+                call_kwargs[k] = params[k]
+
+        try:
+            image = _decode_image(data["inputs"])
+        except Exception as e:
+            return {"error": f"Failed to decode 'inputs': {e}"}
+
+        # Run the pipeline. For Intrinsics, output contains:
+        #  - prediction: (num_targets * N) x 3 x H x W  in [0,1]
+        #  - target_properties: names + visualization hints
+        #  - uncertainty: present if ensemble_size > 2
+        try:
+            out = self.pipe(image, **call_kwargs)
+        except TypeError as e:
+            # Common pitfall: wrong arg names (e.g. 'processing_res' instead of 'processing_resolution')
+            return {"error": f"Pipeline call failed: {e}"}
+        except Exception as e:
+            return {"error": f"Inference failed: {e}"}
+
+        # Targets (order from target_properties)
+        target_props = getattr(out, "target_properties", None)
+        target_names: List[str] = []
+        if target_props and isinstance(target_props, dict):
+            target_names = list(target_props.get("target_names", []))
+
+        resp: Dict[str, Any] = {
+            "targets": target_names or ["albedo", "diffuse_shading", "residual"],
+            "tensors_info": None,
+            "uncertainty_available": getattr(out, "uncertainty", None) is not None,
+        }
+
+        # Shapes / dtype
+        pred = out.prediction  # torch.Tensor or np.ndarray
+        if isinstance(pred, torch.Tensor):
+            n, c, h, w = pred.shape
+            resp["tensors_info"] = {"shape": [int(n), int(c), int(h), int(w)], "dtype": str(pred.dtype).replace("torch.", "")}
+        elif isinstance(pred, np.ndarray):
+            h, w, c = pred.shape[-3:]
+            resp["tensors_info"] = {"shape": list(pred.shape), "dtype": str(pred.dtype)}
+
+        # Optional visualizations using the official helper
+        if return_visuals:
+            try:
+                vis_list = self.pipe.image_processor.visualize_intrinsics(
+                    out.prediction, target_props
+                )  # -> List[Dict[str, PIL.Image]]
+                # Usually returns a single dict with keys matching target_names
+                visuals = {}
+                if isinstance(vis_list, list) and len(vis_list) > 0 and isinstance(vis_list[0], dict):
+                    for k, v in vis_list[0].items():
+                        if isinstance(v, Image.Image):
+                            visuals[k] = _b64_png(v)
+                resp["visualizations"] = visuals
+            except Exception as e:
+                resp["visualizations_error"] = f"Failed to visualize: {e}"
+
+        return {"outputs": resp}

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.1
+accelerate>=0.28
+diffusers>=0.35.0
+transformers>=4.44
+Pillow>=10.3
+safetensors>=0.4.2
+numpy>=1.26