Blame for app.py

"""ShadeNet-5 Gradio Space: photo -> albedo, relative depth, surface normals.

Runs the 4/8-bit "MIX" ONNX export of the model on the free CPU tier
(onnxruntime only; torch is used just for the rotation resampling of the
optional enhanced mode).
"""
import json
import os
import time

import gradio as gr
import numpy as np
import onnxruntime as ort
import torch
from huggingface_hub import hf_hub_download
from PIL import Image, ImageOps

from ensemble import mean_of, median_of, tta_passes, unit_normals
from framing import crop_to_box, display_size, letterbox, resize_maps

try:                                   # present on ZeroGPU Spaces only
    import spaces
except ImportError:                    # local tests, plain CPU hosts
    spaces = None

if spaces is not None:
    # ZeroGPU hosting (the free tier for Gradio Spaces) wants one function
    # decorated with @spaces.GPU. It is never called: all inference below runs
    # on the Space's CPU through ONNX Runtime, so no GPU quota is used.
    @spaces.GPU(duration=1)
    def _zerogpu_placeholder():
        return None

MODEL_REPO = "singam96/ShadeNet-5-35M"
ONNX_FILE = "onnx/shadenet5_mix_int4_int8.onnx"
REFINER_FILE = "onnx/shadenet5_refiner.onnx"
ROT_DEG = 4.0                      # the card's enhanced recipe: mirror + +-4 deg

# config.json also tells us the input size (and is what the Hub counts as a
# download of the model).
with open(hf_hub_download(MODEL_REPO, "config.json")) as f:
    CFG = json.load(f)
SIZE = int(CFG["image_size"])

def cpu_budget() -> int:
    """CPUs this container may use (its cgroup quota), not the host's core
    count: asking ONNX Runtime for more threads than the quota made it about
    9x slower in a measurement (64 host cores, 12 CPUs of quota)."""
    try:
        quota, period = open("/sys/fs/cgroup/cpu.max").read().split()
        if quota != "max":
            return max(1, min(8, int(int(quota) / int(period))))
    except Exception:
        pass
    try:
        return max(1, min(8, len(os.sched_getaffinity(0))))
    except Exception:
        return 2


THREADS = int(os.environ.get("ORT_THREADS", cpu_budget()))
torch.set_num_threads(THREADS)
_so = ort.SessionOptions()
_so.log_severity_level = 3
_so.intra_op_num_threads = THREADS
SESS = ort.InferenceSession(hf_hub_download(MODEL_REPO, ONNX_FILE), _so,
                            providers=["CPUExecutionProvider"])
RSESS = ort.InferenceSession(hf_hub_download(MODEL_REPO, REFINER_FILE), _so,
                             providers=["CPUExecutionProvider"])

MODE_SINGLE = "Single pass"
MODE_REFINE = "Single pass + refiner (recommended)"
MODE_SIX = "Six passes (mirror + rotations, slower)"


def to_tensor(img: Image.Image) -> torch.Tensor:
    arr = np.asarray(img, dtype=np.float32) / 255.0
    return torch.from_numpy(arr.transpose(2, 0, 1)[None] * 2.0 - 1.0)


def refine(c: torch.Tensor) -> torch.Tensor:
    out = SESS.run(None, {"input_rgb": np.ascontiguousarray(c.numpy())})[0]
    return torch.from_numpy(out)


def run_refiner(photo: torch.Tensor, maps: torch.Tensor) -> torch.Tensor:
    """The small refiner (trained to turn a single pass into the six-pass result)."""
    out = RSESS.run(None, {"photo": np.ascontiguousarray(photo.numpy()),
                           "maps": np.ascontiguousarray(maps.numpy())})[0]
    return torch.from_numpy(out)


def to_img(x01: torch.Tensor) -> Image.Image:
    if x01.shape[1] == 1:
        x01 = x01.expand(-1, 3, -1, -1)
    arr = (x01[0].permute(1, 2, 0).clamp(0, 1).numpy() * 255).round()
    return Image.fromarray(arr.astype("uint8"))


def predict(image, mode, median, progress=gr.Progress()):
    if image is None:
        raise gr.Error("Upload a photo first.")
    photo = ImageOps.exif_transpose(image).convert("RGB")
    # Fit the WHOLE photo into the square input (long side -> 448, the rest
    # mirrored from the photo's own edge) and cut the padding off the maps again.
    square, box = letterbox(photo, SIZE)
    x = to_tensor(square)
    ow, oh = display_size(*photo.size)
    enhanced = mode == MODE_SIX
    use_median = bool(enhanced and median)
    diff = None
    t0 = time.time()
    with torch.no_grad():
        if enhanced:
            passes = tta_passes(
                refine, x, flip=True, rot_deg=ROT_DEG,
                progress=lambda i, n: progress((i, n), desc="forward pass"))
            mean_maps = mean_of(passes)
            maps = median_of(passes) if use_median else mean_maps
        elif mode == MODE_REFINE:
            single = refine(x)
            maps = run_refiner(x, single).float().clamp(-1.0, 1.0)
        else:
            maps = refine(x)
        maps = resize_maps(crop_to_box(unit_normals(maps.float()), box), ow, oh)
        if use_median:                     # where the two reductions disagree
            ref = resize_maps(crop_to_box(unit_normals(mean_maps.float()), box), ow, oh)
            d = ((maps - ref).abs().mean(dim=1, keepdim=True) * 8).clamp(0, 1)
            diff = to_img(d)
    dt = time.time() - t0
    albedo = to_img((maps[:, 0:3] + 1) / 2)
    depth = to_img((maps[:, 3:4] + 1) / 2)
    normal = to_img((maps[:, 4:7] + 1) / 2)
    how = {MODE_SINGLE: "1 forward pass",
           MODE_REFINE: "1 forward pass + the small refiner",
           MODE_SIX: "6 forward passes (" + ("median" if use_median else "mean") + ")"}[mode]
    note = (f"{how} in {dt:.1f} s on CPU; the whole {photo.size[0]}x{photo.size[1]} "
            f"photo is kept (model input {SIZE}x{SIZE}, long side scaled to {SIZE}).")
    if median and not enhanced:
        note += " Median only applies to the six-pass mode."
    return photo.resize((ow, oh), Image.BICUBIC), albedo, depth, normal, diff, note


EXAMPLES = [[os.path.join("examples", f), MODE_REFINE, False]
            for f in sorted(os.listdir("examples"))
            if f.lower().endswith((".jpg", ".png"))]

with gr.Blocks(title="ShadeNet-5") as demo:
    gr.Markdown(
        "# ShadeNet-5: albedo, depth and normals from one photo\n"
        "A 35M-parameter feed-forward network on a frozen DINOv2 backbone. "
        "This Space runs the 4/8-bit **MIX** ONNX export on a free CPU. "
        "Depth is relative (per-image scale and shift); normals are "
        "x right, y up, z toward the viewer.\n\n"
        "**Single pass** is one forward pass. **Single pass + refiner** adds a "
        "small trained UNet (0.91M parameters) that cleans up the maps, mostly "
        "the normals, for a fraction of a second more: it recovers about "
        "three quarters of what the slower mode gains. **Six passes** averages the "
        "mirrored and rotated (+-4 degrees) photo and is the best but about six "
        "times slower; tick **Median** to take the median of the six instead of "
        "the mean (marginally worse in the card's measurements) and see the "
        "difference image. "
        f"[Model card and files](https://huggingface.co/{MODEL_REPO}). "
        "The example photos are Flickr8k photos shown for illustration only; "
        "see the model card for their licence status.")
    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="pil", label="Photo", sources=["upload", "clipboard"])
            mode = gr.Radio([MODE_SINGLE, MODE_REFINE, MODE_SIX], value=MODE_REFINE,
                            label="Mode")
            med = gr.Checkbox(False, label="Median instead of mean (six passes only)",
                              interactive=False)
            btn = gr.Button("Decompose", variant="primary")
            msg = gr.Markdown()
        with gr.Column():
            with gr.Row():
                o_in = gr.Image(label="Input")
                o_alb = gr.Image(label="Albedo")
            with gr.Row():
                o_dep = gr.Image(label="Depth (lighter = farther)")
                o_nor = gr.Image(label="Normal")
            o_diff = gr.Image(label="Median minus mean (x8, all maps; black = identical)")
    mode.change(lambda m: gr.update(interactive=(m == MODE_SIX), value=False), mode, med)
    outs = [o_in, o_alb, o_dep, o_nor, o_diff, msg]
    btn.click(predict, [inp, mode, med], outs)
    gr.Examples(EXAMPLES, [inp, mode, med], outs, fn=predict, cache_examples=False)

if __name__ == "__main__":
    demo.queue(max_size=8).launch()