ShadeNet-3.2-5M (model): weights, ONNX, app + card

.gitattributes

@@ -33,3 +33,28 @@ saved_model/**/* 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
+assets/103195344_5d2dc613a3_result.png filter=lfs diff=lfs merge=lfs -text
+assets/154871781_ae77696b77_result.png filter=lfs diff=lfs merge=lfs -text
+assets/1561658940_a947f2446a_result.png filter=lfs diff=lfs merge=lfs -text
+assets/157139628_5dc483e2e4_result.png filter=lfs diff=lfs merge=lfs -text
+assets/159712188_d530dd478c_result.png filter=lfs diff=lfs merge=lfs -text
+assets/160541986_d5be2ab4c1_result.png filter=lfs diff=lfs merge=lfs -text
+assets/160566014_59528ff897_result.png filter=lfs diff=lfs merge=lfs -text
+assets/160585932_fa6339f248_result.png filter=lfs diff=lfs merge=lfs -text
+assets/160792599_6a7ec52516_result.png filter=lfs diff=lfs merge=lfs -text
+assets/161669933_3e7d8c7e2c_result.png filter=lfs diff=lfs merge=lfs -text
+assets/2260560631_09093be4c6_result.png filter=lfs diff=lfs merge=lfs -text
+assets/2312984882_bec7849e09_result.png filter=lfs diff=lfs merge=lfs -text
+assets/241345721_3f3724a7fc_result.png filter=lfs diff=lfs merge=lfs -text
+assets/2453318633_550228acd4_result.png filter=lfs diff=lfs merge=lfs -text
+assets/252578659_9e404b6430_result.png filter=lfs diff=lfs merge=lfs -text
+assets/307994435_592f933a6d_result.png filter=lfs diff=lfs merge=lfs -text
+assets/3185645793_49de805194_result.png filter=lfs diff=lfs merge=lfs -text
+assets/326585030_e1dcca2562_result.png filter=lfs diff=lfs merge=lfs -text
+assets/3440104178_6871a24e13_result.png filter=lfs diff=lfs merge=lfs -text
+assets/487071033_27e460a1b9_result.png filter=lfs diff=lfs merge=lfs -text
+assets/atom_dictionary.png filter=lfs diff=lfs merge=lfs -text
+assets/compare_154871781_ae77696b77.png filter=lfs diff=lfs merge=lfs -text
+assets/compare_159712188_d530dd478c.png filter=lfs diff=lfs merge=lfs -text
+assets/compare_160585932_fa6339f248.png filter=lfs diff=lfs merge=lfs -text
+assets/hero.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,206 @@
 ---
 license: apache-2.0
+pipeline_tag: image-to-image
+library_name: pytorch
+tags:
+  - onnx
+  - inverse-rendering
+  - image-decomposition
+  - albedo
+  - normal-map
+  - depth-estimation
+  - pbr
+  - material-estimation
+  - flickr8k
+  - shadenet
+datasets:
+  - singam96/flickr8k_marigold_v2
+metrics:
+  - mse
+model-index:
+  - name: ShadeNet-3.2-5M
+    results:
+      - task:
+          type: image-to-image
+        dataset:
+          name: flickr8k_marigold_v2 (val split, 807 images)
+          type: singam96/flickr8k_marigold_v2
+        metrics:
+          - name: val/loss (weighted MSE + recon)
+            type: mse
+            value: 0.1671
 ---
+
+# ShadeNet-3.2 5M
+
+A lightweight inverse-rendering model: one photo in, **albedo + relative depth + surface normals + shading** out (8 channels, 5.0M params). Successor of [ShadeNet-2](https://huggingface.co/singam96/ShadeNet-2-20M) — **4× smaller, better depth and normals**.
+
+[![hero](assets/hero.png)](assets/hero.png)
+
+## TL;DR
+
+| | |
+|---|---|
+| Params | **5.0M** (3.2M trainable + 1.8M frozen MobileNetV2 trunk) |
+| Input | RGB `[1, 3, H, W]` in `[-1, 1]` (384px trained; any multiple of 16) |
+| Output | 8ch `[1, 8, H, W]` in `[-1, 1]`: albedo, relative depth, normals, shading |
+| Val L1 (807 imgs) | albedo 0.695 · depth 0.217 · normal 0.581 |
+| Formats | fp32 ONNX (20MB), fp16 ONNX (10MB), torch checkpoint (44MB) |
+
+## Examples
+
+[![154871781](assets/154871781_ae77696b77_result.png)](assets/154871781_ae77696b77_result.png)
+[![157139628](assets/157139628_5dc483e2e4_result.png)](assets/157139628_5dc483e2e4_result.png)
+[![159712188](assets/159712188_d530dd478c_result.png)](assets/159712188_d530dd478c_result.png)
+[![160541986](assets/160541986_d5be2ab4c1_result.png)](assets/160541986_d5be2ab4c1_result.png)
+[![160566014](assets/160566014_59528ff897_result.png)](assets/160566014_59528ff897_result.png)
+[![160585932](assets/160585932_fa6339f248_result.png)](assets/160585932_fa6339f248_result.png)
+[![160792599](assets/160792599_6a7ec52516_result.png)](assets/160792599_6a7ec52516_result.png)
+[![161669933](assets/161669933_3e7d8c7e2c_result.png)](assets/161669933_3e7d8c7e2c_result.png)
+
+*Each grid: input | albedo | shading / depth | normal | recon (albedo×shading). Click any image for full size.*
+
+## Results
+
+Full 807-image val split, per-map L1 (the comparable metric across versions —
+the headline `val/loss` formula changed between v2 and v3):
+
+| Map (val L1) | ShadeNet-2 (20M) | **ShadeNet-3.2 (5M)** | change |
+|---|---|---|---|
+| Albedo | 0.708 | **0.695** | −1.7% |
+| Depth (SSI-aligned) | 0.247 | **0.217** | **−12%** |
+| Normal | 0.696 | **0.581** | **−16%** |
+
+[![curves](assets/training_curves.png)](assets/training_curves.png)
+
+Checkpoint variants (full val, pruned top-32 dictionary):
+
+| Weights | val/loss | albedo L1 | depth L1 | normal L1 |
+|---|---|---|---|---|
+| best.ckpt (raw) | **0.1669** | 0.7013 | 0.2241 | **0.5709** |
+| best.ckpt (EMA) | 0.1671 | **0.6952** | **0.2174** | 0.5807 |
+
+Shipped ONNX uses the **EMA** weights. All page outputs and the Space run a
+**3-pass multi-scale median** (scales 0.875/1.0/1.125) — a mild denoise; the
+model's raw single-pass albedo is sharper than its pseudo-labels, so this
+trades a little detail for lower variance.
+
+## ShadeNet-2 vs ShadeNet-3.2
+
+Same inputs (ShadeNet-2 top rows, ShadeNet-3.2 bottom rows), maps only. Both
+use their shipped weights; ShadeNet-3.2 runs the 3-pass multi-scale median.
+
+[![compare 1](assets/compare_154871781_ae77696b77.png)](assets/compare_154871781_ae77696b77.png)
+[![compare 2](assets/compare_159712188_d530dd478c.png)](assets/compare_159712188_d530dd478c.png)
+[![compare 3](assets/compare_160585932_fa6339f248.png)](assets/compare_160585932_fa6339f248.png)
+
+## Architecture
+
+**ParallelUNet generator (4.98M params)** + spectral-norm GroupNorm PatchGAN discriminator (2.77M, training only):
+
+- Dual parallel encoders — vanilla UNet path plus a **frozen MobileNetV2** feature trunk, fused at every decoder level
+- **Depthwise-separable factorized convs** (1×3 + 3×1) throughout; full H/32 bottleneck; reflect padding
+- **Patch-dictionary output tail**: 16×16 tiles softmax-addressed over **32 learned per-channel atoms** (pruned from 1024 — the top-32 hold 99.4% of addressing mass), blended back into the signal before tanh
+- Single-pass RGB → 8ch output; EMA weight shadow (shipped weights are EMA)
+
+## Patch dictionary
+
+The tail softmax-addresses 32 learned 16×16 atoms per tile (kept from 1024
+after measuring per-atom selection: only ~34 atoms are ever used, ~31 cover
+99% of the mass). Shown below per output channel — each panel is the 8×4 atom
+grid, shared grayscale scale.
+
+[![atom dictionary](assets/atom_dictionary.png)](assets/atom_dictionary.png)
+
+## Output maps
+
+| Map | Channels | Range | Description |
+|---|---|---|---|
+| Albedo | 3 `[0:3]` | [−1, 1] | Reflectance / diffuse color, lighting factored out |
+| Depth | 1 `[3:4]` | [−1, 1] | **Relative** depth (0=near), affine-ambiguous |
+| Normal | 3 `[4:7]` | [−1, 1] | Surface normals, unit-length regularised |
+| Shading | 1 `[7:8]` | [−1, 1] | Grayscale irradiance; `input ≈ albedo × shading` |
+| Recon | — | — | `albedo × shading` re-rendering (diagnostic, not a head) |
+
+## Files
+
+```
+├── app.py                    # Gradio Space app (fp16 ONNX)
+├── inference.py              # Standalone torch CLI
+├── inference_utils.py        # Grid visualisation (numpy/PIL)
+├── model.py                  # Standalone generator architecture
+├── requirements.txt
+├── checkpoints/shadenet32.ckpt     # Torch weights, EMA (44MB)
+└── onnx/
+    ├── model.onnx                  # fp32, EMA (20MB) — GPU via CUDA EP
+    └── model_fp16.onnx             # fp16, EMA (10MB) — CPU
+```
+
+## Usage
+
+### Gradio Space
+
+Try it in your browser — no installation: **[singam96/ShadeNet-3.2-5M Space](https://huggingface.co/spaces/singam96/ShadeNet-3.2-5M)**.
+
+This repo ships `app.py`, the Space entrypoint. To recreate it: New Space → Gradio SDK → point at this repo.
+
+### Torch CLI
+
+```bash
+pip install torch torchvision pillow numpy
+python inference.py photo.jpg --output_dir ./output
+# --checkpoint ./checkpoints/shadenet32.ckpt --image-size 512 --no-ema to disable EMA
+```
+
+### ONNX (CPU)
+
+```bash
+pip install onnxruntime pillow numpy
+python - <<'EOF'
+import onnxruntime as ort, numpy as np
+from PIL import Image
+from inference_utils import build_grid, pil_to_np, resize_pad
+sess = ort.InferenceSession("onnx/model_fp16.onnx", providers=["CPUExecutionProvider"])
+img = resize_pad(Image.open("photo.jpg"), 512)
+out = sess.run(None, {"input_rgb": pil_to_np(img).astype(np.float32)})[0]
+build_grid(img, out).save("result.png")
+EOF
+```
+
+Input: `[1, 3, H, W]` in `[-1, 1]` (any H, W; multiples of 16 recommended).
+Output: `[1, 8, H, W]` in `[-1, 1]`.
+
+## Training
+
+Trained from scratch on [`singam96/flickr8k_marigold_v2`](https://huggingface.co/datasets/singam96/flickr8k_marigold_v2) (8077 Flickr8k photos with Marigold-V2 pseudo-labels), 384px, fp32, single GTX 1650, early-stopped on `val/loss` (patience 5) at epoch 12. The 1024-atom dictionary was pruned to its top-32 by addressing mass (no retraining; output error vs full ~5e-4 mean) for the release.
+
+Losses: scale-invariant MSE on albedo (per-channel std alignment) + **scale-shift-invariant** MSE on depth after least-squares alignment (decoded Marigold depth is relative) + **Sobel gradient-matching** on depth (edge crispness) + MSE on normals + self-supervised **reconstruction coupling** (`albedo×shading ≈ input`, the shading head's only supervision) + LSGAN + normal unit-length penalty. Weight decay 1e-4 (patch dictionary exempt).
+
+## Limitations
+
+- Depth is **relative**, not metric — don't read meters off it.
+- Shading assumes **white light**; strongly colored illumination (sunsets, neon) leaks into albedo.
+- Normals are noisy in foliage/sky — those pseudo-labels were noisy too.
+- Occasional localized artifacts in albedo/shading (learned prior pockets).
+- No shadows/global illumination — relighting-style use is approximate.
+- Uncertainty is not provided: confidently-wrong pseudo-labels are fitted confidently.
+
+## Attribution
+
+Supervision labels come from **Marigold V2** (Ke et al.) applied to **Flickr8k** (Hodosh et al.):
+
+- Marigold: *Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation* — Ke, Obukhov, Metzger, Daudt, Schindler, Schindler (CVPR 2024)
+- Flickr8k: *Framing Image Description as a Ranking Task* — Hodosh, Young, Hockenmaier (2013)
+
+This model (weights + code) is Apache-2.0; upstream dataset/model terms still apply to their artifacts.
+
+## Citation
+
+```bibtex
+@software{shadenet32,
+  author = {Sachin},
+  title = {ShadeNet-3.2: single-image inverse rendering (5M)},
+  year = {2026},
+  url = {https://huggingface.co/singam96/ShadeNet-3.2-5M}
+}
+```

app.py

@@ -0,0 +1,59 @@
+"""ShadeNet-3.2 Gradio Space (CPU-friendly: fp16 ONNX backend)."""
+import os
+
+import gradio as gr
+import numpy as np
+import onnxruntime as ort
+import spaces
+from PIL import Image
+
+from inference_utils import (build_grid, extract_maps, pil_to_np, resize_pad,
+                             run_ensemble_onnx)
+
+MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                           "onnx", "model_fp16.onnx")
+IMAGE_SIZE = 512
+
+_sess = ort.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
+
+
+@spaces.GPU(duration=60)
+def decompose(image: Image.Image):
+    """Returns (grid, input, albedo, shading, depth, normal, recon)."""
+    img_rgb = resize_pad(image.convert("RGB"), IMAGE_SIZE)
+    out = run_ensemble_onnx(_sess, img_rgb, IMAGE_SIZE)
+    maps = extract_maps(out, img_rgb)
+    return (build_grid(img_rgb, out), maps["input"], maps["albedo"],
+            maps["shading"], maps["depth"], maps["normal"], maps["recon"])
+
+
+EXAMPLES = []
+_exdir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                      "assets", "examples")
+if os.path.isdir(_exdir):
+    EXAMPLES = sorted(os.path.join(_exdir, f) for f in os.listdir(_exdir)
+                      if f.lower().endswith((".png", ".jpg", ".jpeg")))
+
+demo = gr.Interface(
+    fn=decompose,
+    inputs=gr.Image(type="pil", label="Input photo"),
+    outputs=[
+        gr.Image(type="pil", label="Overview grid"),
+        gr.Image(type="pil", label="Input (resized)"),
+        gr.Image(type="pil", label="Albedo"),
+        gr.Image(type="pil", label="Shading"),
+        gr.Image(type="pil", label="Depth (relative: dark=near)"),
+        gr.Image(type="pil", label="Normal"),
+        gr.Image(type="pil", label="Recon (albedo x shading)"),
+    ],
+    title="ShadeNet-3.2: single-image inverse rendering (7M)",
+    description=("Decomposes a photo into albedo, relative depth, surface "
+                 "normals and shading (grayscale irradiance). Successor of "
+                 "ShadeNet-2 (20M), 3x smaller with better depth/normal "
+                 "accuracy. Runs the fp16 ONNX model."),
+    examples=EXAMPLES,
+    cache_examples=False,
+)
+
+if __name__ == "__main__":
+    demo.launch()

checkpoints/shadenet32.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b660a6d955f2df63304312d0fe491bb6a0bd9e71fcdff69a7398a9b19fe3fd0
+size 44034045

config.json

@@ -0,0 +1,18 @@
+{
+  "model_type": "shadenet",
+  "model_name": "ShadeNet-3.2-5M",
+  "architecture": "ParallelUNet-v3",
+  "params_total": 4977208,
+  "params_trainable": 3165496,
+  "image_size": 384,
+  "in_ch": 3,
+  "out_ch": 8,
+  "maps": {
+    "albedo": [0, 3],
+    "depth": [3, 4],
+    "normal": [4, 7],
+    "shading": [7, 8]
+  },
+  "framework": "pytorch",
+  "onnx": ["onnx/model.onnx", "onnx/model_fp16.onnx"]
+}

inference.py

@@ -0,0 +1,68 @@
+"""ShadeNet-3.2 torch inference (standalone CLI).
+
+Usage:
+    pip install -r requirements.txt
+    python inference.py input.jpg --output_dir ./output
+    python inference.py ./photos --checkpoint ./checkpoints/shadenet32.ckpt
+"""
+import argparse
+import os
+
+import torch
+from PIL import Image
+
+from model import load_shadenet32
+from inference_utils import build_grid, pil_to_np, resize_pad
+
+IMAGE_SIZE = 512
+ENS_SCALES = (0.875, 1.0, 1.125)
+
+
+@torch.no_grad()
+def run(model, img_rgb: Image.Image, device: str,
+        image_size: int = IMAGE_SIZE, scales=ENS_SCALES) -> "torch.Tensor":
+    """3-pass multi-scale median (matches the model card / Space)."""
+    outs = []
+    for s in scales:
+        side = max(16, int(round(image_size * s)) // 16 * 16)  # dict needs /16
+        scaled = resize_pad(img_rgb, side)
+        x = torch.from_numpy(pil_to_np(scaled)).to(device)
+        o = model(x)
+        if side != image_size:
+            o = torch.nn.functional.interpolate(
+                o, size=(image_size, image_size), mode="bilinear",
+                align_corners=False)
+        outs.append(o.cpu())
+    return torch.stack(outs, dim=0).median(dim=0).values
+
+
+def main():
+    ap = argparse.ArgumentParser(description="ShadeNet-3.2 inference")
+    ap.add_argument("input", help="Image file or folder")
+    ap.add_argument("--checkpoint", default="./checkpoints/shadenet32.ckpt")
+    ap.add_argument("--output_dir", default="./output")
+    ap.add_argument("--image-size", type=int, default=IMAGE_SIZE)
+    ap.add_argument("--no-ema", action="store_true",
+                    help="Use raw weights instead of EMA")
+    args = ap.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = load_shadenet32(args.checkpoint, device, use_ema=not args.no_ema)
+
+    if os.path.isdir(args.input):
+        files = sorted(os.path.join(args.input, f) for f in os.listdir(args.input)
+                       if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp")))
+    else:
+        files = [args.input]
+    os.makedirs(args.output_dir, exist_ok=True)
+    for fp in files:
+        img_rgb = resize_pad(Image.open(fp), args.image_size)
+        out = run(model, img_rgb, device).numpy()
+        base = os.path.splitext(os.path.basename(fp))[0]
+        build_grid(img_rgb, out).save(
+            os.path.join(args.output_dir, f"{base}_result.png"))
+        print(f"Saved: {base}_result.png")
+
+
+if __name__ == "__main__":
+    main()

inference_utils.py

@@ -0,0 +1,102 @@
+"""Standalone inference utilities (numpy/PIL only — no torch needed)."""
+import numpy as np
+from PIL import Image, ImageDraw
+
+ALBEDO = slice(0, 3)
+DEPTH = slice(3, 4)
+NORMAL = slice(4, 7)
+SHADING = slice(7, 8)
+
+
+def resize_pad(img_rgb: Image.Image, image_size: int) -> Image.Image:
+    """Scale + center-crop to a square (RGB in, RGB out)."""
+    img_rgb = img_rgb.convert("RGB")
+    w, h = img_rgb.size
+    size = int(image_size)
+    scale = max(size / float(w), size / float(h))
+    img_r = img_rgb.resize((max(1, int(round(w * scale))),
+                            max(1, int(round(h * scale)))), Image.BICUBIC)
+    left = (img_r.size[0] - size) // 2
+    top = (img_r.size[1] - size) // 2
+    return img_r.crop((left, top, left + size, top + size))
+
+
+def pil_to_np(img_rgb: Image.Image) -> np.ndarray:
+    """PIL RGB -> [1, 3, H, W] float32 in [-1, 1]."""
+    arr = np.array(img_rgb.convert("RGB"), dtype=np.float32).transpose(2, 0, 1)
+    return arr[np.newaxis] / 255.0 * 2.0 - 1.0
+
+
+def _to_01(output: np.ndarray) -> np.ndarray:
+    return np.clip((output + 1.0) / 2.0, 0, 1)
+
+
+def _tile(arr01: np.ndarray, ts: int) -> Image.Image:
+    return Image.fromarray(
+        (arr01 * 255.0 + 0.5).astype("uint8")).resize((ts, ts), Image.BICUBIC)
+
+
+def _resize_map(out: np.ndarray, size: int) -> np.ndarray:
+    """Bilinear-resize a [1, C, H, W] float map to [1, C, size, size]."""
+    c = out.shape[1]
+    res = np.empty((1, c, size, size), dtype=np.float32)
+    for i in range(c):
+        im = Image.fromarray(out[0, i].astype(np.float32), mode="F")
+        res[0, i] = np.asarray(im.resize((size, size), Image.BILINEAR))
+    return res
+
+
+ENS_SCALES = (0.875, 1.0, 1.125)
+
+
+def run_ensemble_onnx(sess, img_rgb: Image.Image, image_size: int = 512,
+                      scales=ENS_SCALES) -> np.ndarray:
+    """3-pass multi-scale median (torch-free). Returns [1, 8, S, S]."""
+    outs = []
+    for s in scales:
+        side = max(16, int(round(image_size * s)) // 16 * 16)  # dict needs /16
+        scaled = resize_pad(img_rgb, side)
+        x = pil_to_np(scaled).astype(np.float32)
+        o = sess.run(None, {"input_rgb": x})[0]
+        outs.append(o if side == image_size else _resize_map(o, image_size))
+    return np.median(np.stack(outs, axis=0), axis=0).astype(np.float32)
+
+
+def extract_maps(output: np.ndarray, img_rgb: Image.Image,
+                 ts: int = 256) -> dict:
+    """Split a [1, 8, H, W] output ([-1, 1]) into labelled PIL tiles."""
+    assert output.shape[1] == 8, f"want 8 channels, got {output.shape[1]}"
+    o = _to_01(output)
+    alb = o[:, ALBEDO][0].transpose(1, 2, 0)
+    sh = o[:, SHADING][0, 0]
+    d = o[:, DEPTH][0, 0]
+    span = d.max() - d.min()
+    dg = (d - d.min()) / (span if span > 1e-6 else 1.0)
+    return {
+        "input": img_rgb.resize((ts, ts), Image.BICUBIC),
+        "albedo": _tile(alb, ts),
+        "shading": _tile(np.repeat(sh[..., None], 3, axis=2), ts),
+        "depth": _tile(dg, ts).convert("RGB"),
+        "normal": _tile(o[:, NORMAL][0].transpose(1, 2, 0), ts),
+        "recon": _tile(np.clip(alb * sh[..., None], 0, 1), ts),
+    }
+
+
+def build_grid(img_rgb: Image.Image, output: np.ndarray,
+               ts: int = 256) -> Image.Image:
+    """3x2 labelled grid: input|albedo|shading / depth|normal|recon."""
+    maps = extract_maps(output, img_rgb, ts)
+    grid = Image.new("RGB", (ts * 3, ts * 2))
+    grid.paste(maps["input"], (0, 0))
+    grid.paste(maps["albedo"], (ts, 0))
+    grid.paste(maps["shading"], (ts * 2, 0))
+    grid.paste(maps["depth"], (0, ts))
+    grid.paste(maps["normal"], (ts, ts))
+    grid.paste(maps["recon"], (ts * 2, ts))
+    draw = ImageDraw.Draw(grid)
+    for (x, y, name) in [(4, 4, "input"), (ts + 4, 4, "albedo"),
+                         (ts * 2 + 4, 4, "shading"), (4, ts + 4, "depth"),
+                         (ts + 4, ts + 4, "normal"),
+                         (ts * 2 + 4, ts + 4, "recon")]:
+        draw.text((x, y), name, fill=(255, 255, 0))
+    return grid

model.py

@@ -0,0 +1,338 @@
+"""ShadeNet-3.2 generator (standalone, no Lightning dependency).
+
+ParallelUNet v3: dual parallel encoders (vanilla UNet + frozen MobileNetV2),
+depthwise-separable factorized convs, full H/32 bottleneck, and a patch
+dictionary output tail -> 8ch intrinsic maps in [-1, 1]:
+  [0:3] albedo | [3:4] relative depth (0=near) | [4:7] normal | [7:8] shading
+"""
+import math
+
+import torch
+import torch.nn as nn
+import torchvision
+from torchvision.models.feature_extraction import create_feature_extractor
+
+OUT_CH = 8
+WIDTH_MULT = 0.9
+
+
+def _num_groups(ch: int) -> int:
+    g = min(32, ch)
+    while ch % g != 0:
+        g -= 1
+    return g
+
+
+def _scale(width_mult):
+    """Channel scaler: width_mult rounded to multiples of 8 (GroupNorm-safe)."""
+    w = float(width_mult)
+
+    def C(n):
+        return max(8, int(round(n * w / 8.0) * 8))
+
+    return C
+
+
+class ChannelLinear(nn.Module):
+    """Per-channel learnable affine: y = x * weight + bias."""
+
+    def __init__(self, channels: int, init_scale: float = 0.01):
+        super().__init__()
+        self.weight = nn.Parameter(
+            torch.empty(1, channels, 1, 1).uniform_(-init_scale, init_scale) + 1.0
+        )
+        self.bias = nn.Parameter(
+            torch.empty(1, channels, 1, 1).uniform_(-init_scale, init_scale)
+        )
+
+    def forward(self, x):
+        return x * self.weight + self.bias
+
+
+class DoubleConv(nn.Module):
+    """Factorized (1,3)+(3,1) conv pair with GN+ELU, reflect padding."""
+
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        g = _num_groups(out_ch)
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, (1, 3), padding=(0, 1),
+                      padding_mode="reflect", bias=False),
+            nn.GroupNorm(g, out_ch),
+            nn.ELU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, (3, 1), padding=(1, 0),
+                      padding_mode="reflect", bias=False),
+            nn.GroupNorm(g, out_ch),
+            nn.ELU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class DSDoubleConv(nn.Module):
+    """Depthwise-separable DoubleConv: same geometry, ~3x fewer params."""
+
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        g = _num_groups(out_ch)
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_ch, in_ch, (1, 3), padding=(0, 1),
+                      padding_mode="reflect", groups=in_ch, bias=False),
+            nn.Conv2d(in_ch, out_ch, 1, bias=False),
+            nn.GroupNorm(g, out_ch),
+            nn.ELU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, (3, 1), padding=(1, 0),
+                      padding_mode="reflect", groups=out_ch, bias=False),
+            nn.Conv2d(out_ch, out_ch, 1, bias=False),
+            nn.GroupNorm(g, out_ch),
+            nn.ELU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Down(nn.Module):
+    def __init__(self, in_ch, out_ch, block=DSDoubleConv):
+        super().__init__()
+        self.pool = nn.AvgPool2d(2)
+        self.conv = block(in_ch, out_ch)
+
+    def forward(self, x):
+        return self.conv(self.pool(x))
+
+
+class Up(nn.Module):
+    def __init__(self, in_ch, out_ch, skip_ch=None, block=DSDoubleConv):
+        super().__init__()
+        skip_ch = skip_ch or in_ch
+        self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2)
+        self.conv = block(out_ch + skip_ch, out_ch)
+
+    def forward(self, x, skip):
+        x = self.up(x)
+        return self.conv(torch.cat([skip, x], dim=1))
+
+
+class PatchDictionaryBias(nn.Module):
+    """Soft patch-dictionary prior: tile -> softmax over atoms -> blend."""
+
+    def __init__(self, channels=OUT_CH, patch=16, n_atoms=1024, hidden=128,
+                 deblock=True, blur_radius=2):
+        super().__init__()
+        self.patch = int(patch)
+        self.n_atoms = int(n_atoms)
+        self.atoms = nn.Parameter(
+            torch.zeros(n_atoms, channels, patch, patch))
+        self.desc = nn.Conv2d(channels, channels, 4, stride=4,
+                              groups=channels, bias=False)
+        d = channels * (patch // 4) * (patch // 4)
+        self.addr = nn.Sequential(
+            nn.Linear(d, hidden),
+            nn.ELU(inplace=True),
+            nn.Linear(hidden, n_atoms),
+        )
+        self.blend = nn.Conv2d(2 * channels, channels, 1, bias=True)
+        self.blur_recon = bool(deblock)
+        self.blur_radius = int(blur_radius)
+        r = self.blur_radius
+        row = torch.tensor([math.comb(2 * r, i) for i in range(2 * r + 1)],
+                           dtype=torch.float32)
+        row = row / row.sum()
+        k2 = (row[:, None] * row[None, :])[None, None].expand(
+            channels, 1, -1, -1).contiguous()
+        self.register_buffer("blur_k", k2)
+
+    def _gauss_blur(self, t):
+        r = self.blur_radius
+        t = torch.nn.functional.pad(t, (r, r, r, r), mode="reflect")
+        return torch.nn.functional.conv2d(t, self.blur_k, groups=t.shape[1])
+
+    def _deblock(self, recon, b, c, nh, nw, p):
+        img = (recon.reshape(b, nh, nw, c, p, p)
+               .permute(0, 3, 1, 4, 2, 5)
+               .reshape(b, c, nh * p, nw * p))
+        r = self.blur_radius
+        blurred = self._gauss_blur(img)
+        H, W = img.shape[-2:]
+        ih = torch.arange(H, device=img.device)
+        iw = torch.arange(W, device=img.device)
+        mh = ((ih % p) < r) | ((ih % p) >= p - r)
+        mw = ((iw % p) < r) | ((iw % p) >= p - r)
+        mask = (mh[:, None] | mw[None, :]).to(img.dtype)[None, None]
+        img = img + mask * (blurred - img)
+        return (img.reshape(b, c, nh, p, nw, p)
+                .permute(0, 2, 4, 1, 3, 5)
+                .reshape(b * nh * nw, c, p, p))
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        p = self.patch
+        nh, nw = h // p, w // p
+        tiles = (x.reshape(b, c, nh, p, nw, p)
+                 .permute(0, 2, 4, 1, 3, 5)
+                 .reshape(b * nh * nw, c, p, p))
+        feat = self.desc(tiles).flatten(1)
+        wts = torch.softmax(self.addr(feat), dim=1)
+        recon = torch.einsum("na,achw->nchw", wts, self.atoms)
+        if self.blur_recon:
+            recon = self._deblock(recon, b, c, nh, nw, p)
+        both = torch.cat([tiles, recon], dim=1)
+        out = self.blend(both).reshape(b, nh, nw, c, p, p)
+        return out.permute(0, 3, 1, 4, 2, 5).reshape(b, c, nh * p, nw * p)
+
+
+class VanillaEncoder(nn.Module):
+    def __init__(self, in_ch=3, width_mult=1.0):
+        super().__init__()
+        C = _scale(width_mult)
+        self.inc = DoubleConv(in_ch, C(64))
+        self.down1 = Down(C(64), C(128))
+        self.down2 = Down(C(128), C(256))
+        self.down3 = Down(C(256), C(512))
+        self.down4 = Down(C(512), C(256))
+        self.down5 = Down(C(256), C(512))
+        self.mem1 = ChannelLinear(C(64))
+        self.mem2 = ChannelLinear(C(128))
+        self.mem3 = ChannelLinear(C(256))
+        self.mem4 = ChannelLinear(C(512))
+        self.mem5 = ChannelLinear(C(256))
+        self.mem6 = ChannelLinear(C(512))
+
+    def forward(self, x):
+        u1 = self.mem1(self.inc(x))
+        u2 = self.mem2(self.down1(u1))
+        u3 = self.mem3(self.down2(u2))
+        u4 = self.mem4(self.down3(u3))
+        u5 = self.mem5(self.down4(u4))
+        u6 = self.mem6(self.down5(u5))
+        return u1, u2, u3, u4, u5, u6
+
+
+class MobileEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+        self.register_buffer("shift", 1.0 - 2.0 * mean)
+        self.register_buffer("scale", 1.0 / (2.0 * std))
+        mbnet = torchvision.models.mobilenet_v2(weights="IMAGENET1K_V1")
+        self.trunk = create_feature_extractor(
+            mbnet,
+            return_nodes={
+                "features.0": "m0", "features.2": "m1", "features.6": "m2",
+                "features.13": "m3", "features.17": "m4",
+            },
+        )
+        for p in self.trunk.parameters():
+            p.requires_grad = False
+
+    def forward(self, x):
+        feats = self.trunk((x + self.shift) * self.scale)
+        return feats["m0"], feats["m1"], feats["m2"], feats["m3"], feats["m4"]
+
+
+class FusionBottleneck(nn.Module):
+    def __init__(self, width_mult=1.0):
+        super().__init__()
+        C = _scale(width_mult)
+        self.m_proj = nn.Conv2d(320, C(96), 1, bias=False)
+        self.fusion = DSDoubleConv(C(512) + C(96), C(512))
+        self.mem = ChannelLinear(C(512))
+
+    def forward(self, u6, m4):
+        return self.mem(self.fusion(torch.cat([u6, self.m_proj(m4)], dim=1)))
+
+
+class FusedDecoder(nn.Module):
+    def __init__(self, width_mult=1.0):
+        super().__init__()
+        C = _scale(width_mult)
+        self.up0 = Up(C(512), C(256), skip_ch=C(256) + 96)
+        self.up1 = Up(C(256), C(256), skip_ch=C(512) + 32)
+        self.up2 = Up(C(256), C(256), skip_ch=C(256) + 24)
+        self.up3 = Up(C(256), C(128), skip_ch=C(128) + 32)
+        self.up4 = Up(C(128), C(64), skip_ch=C(64))
+        self.mem0 = ChannelLinear(C(256))
+        self.mem1 = ChannelLinear(C(256))
+        self.mem2 = ChannelLinear(C(256))
+        self.mem3 = ChannelLinear(C(128))
+        self.mem4 = ChannelLinear(C(64))
+
+    def forward(self, b, uni, mob):
+        u1, u2, u3, u4, u5 = uni
+        m0, m1, m2, m3 = mob
+        d0 = self.mem0(self.up0(b, torch.cat([u5, m3], dim=1)))
+        d1 = self.mem1(self.up1(d0, torch.cat([u4, m2], dim=1)))
+        d2 = self.mem2(self.up2(d1, torch.cat([u3, m1], dim=1)))
+        d3 = self.mem3(self.up3(d2, torch.cat([u2, m0], dim=1)))
+        return self.mem4(self.up4(d3, u1))
+
+
+class IntrinsicHead(nn.Module):
+    def __init__(self, out_ch=OUT_CH, width_mult=1.0):
+        super().__init__()
+        C = _scale(width_mult)
+        self.head = DoubleConv(C(64) + 3, C(32))
+        self.head_out = nn.Conv2d(C(32), out_ch, 3, padding=1,
+                                  padding_mode="reflect")
+
+    def forward(self, d4, x):
+        return self.head_out(self.head(torch.cat([d4, x], dim=1)))
+
+
+class DictTail(nn.Module):
+    """Patch-dictionary prior (blends internally, deblocked) -> tanh."""
+
+    def __init__(self, out_ch=OUT_CH, deblock=True, dict_atoms=1024):
+        super().__init__()
+        self.dict = PatchDictionaryBias(channels=out_ch, patch=16,
+                                        n_atoms=dict_atoms, deblock=deblock)
+
+    def forward(self, logits):
+        return torch.tanh(self.dict(logits))
+
+
+class ParallelUNet(nn.Module):
+    def __init__(self, in_ch=3, out_ch=OUT_CH, width_mult=WIDTH_MULT,
+                 bias_res=384, deblock=True, dict_atoms=1024):
+        super().__init__()
+        self.encoder = VanillaEncoder(in_ch, width_mult)
+        self.mobile = MobileEncoder()
+        self.bottleneck = FusionBottleneck(width_mult)
+        self.decoder = FusedDecoder(width_mult)
+        self.head = IntrinsicHead(out_ch, width_mult)
+        self.tail = DictTail(out_ch, deblock=deblock, dict_atoms=dict_atoms)
+
+    def forward(self, x):
+        u1, u2, u3, u4, u5, u6 = self.encoder(x)
+        m0, m1, m2, m3, m4 = self.mobile(x)
+        b = self.bottleneck(u6, m4)
+        d4 = self.decoder(b, (u1, u2, u3, u4, u5), (m0, m1, m2, m3))
+        return self.tail(self.head(d4, x))
+
+
+def load_shadenet32(checkpoint_path, device="cpu", use_ema=True,
+                    width_mult=WIDTH_MULT) -> ParallelUNet:
+    """Build the generator and load shadenet32.ckpt (Lightning or raw format).
+
+    Strips the `generator.` prefix from Lightning checkpoints and applies the
+    EMA shadow unless use_ema=False.
+    """
+    model = ParallelUNet(out_ch=OUT_CH, width_mult=width_mult)
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    sd = ckpt.get("state_dict", ckpt)
+    if any(k.startswith("generator.") for k in sd):
+        sd = {k[len("generator."):]: v for k, v in sd.items()
+              if k.startswith("generator.")}
+    model.load_state_dict(sd, strict=False)
+    if use_ema:
+        ema = ckpt.get("ema_generator") or {}
+        if ema:
+            model.load_state_dict(ema, strict=False)
+            print(f"Using EMA weights ({len(ema)} tensors).")
+        else:
+            print("No EMA shadow in checkpoint, using raw weights.")
+    model.eval().to(device)
+    return model

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:364b1676cdc6e3bab1d3035ea575f513a81cffe8ec34180509999d2ff46d4f9d
+size 20053917

onnx/model_fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5375b0c34610fb73f0df22b29e03dc3049f69827882f62f4e1fa74f7b40b2b93
+size 10208498

requirements-space.txt

@@ -0,0 +1,5 @@
+gradio
+onnxruntime
+pillow
+numpy
+spaces

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+pillow
+numpy
+onnxruntime
+gradio