app.py

"""OpenDelight ZeroGPU demo.

Wraps the OpenDelight facial delighting pipeline (matting -> skin mask ->
FFHQ alignment -> delighting network -> detail enhancer -> inverse alignment)
into a single-image Gradio app that runs on Hugging Face Spaces with ZeroGPU.

All model weights and the pure-Python ibug wheels are pulled from a companion
Hugging Face model repo at runtime; nothing in the upstream OpenDelight source
tree is modified.
"""

import os
import sys
import glob
import importlib
import tempfile
import threading

CODE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, CODE_DIR)
# OpenDelight's MAE encoder loads a CWD-relative weight path; anchor CWD to the
# app directory so it resolves regardless of where the process is launched.
os.chdir(CODE_DIR)

WEIGHTS_REPO = os.environ.get("OPENDELIGHT_WEIGHTS_REPO", "suvadityamuk/opendelight-weights")
HF_TOKEN = os.environ.get("HF_TOKEN")

# ---------------------------------------------------------------------------
# Import the DAViD matting runtime in isolation.
#
# The matting modules use absolute imports (`from utils import ...`,
# `from pixelwise_estimator import ...`) and ship their own `utils` module,
# which collides with OpenDelight's top-level `utils.py` (needed by the MAE
# encoder). We load the segmenter first with `matting/runtime` on the path,
# then purge those names from sys.modules so the root `utils` can load cleanly.
# ---------------------------------------------------------------------------
def _load_soft_foreground_segmenter():
    runtime_dir = os.path.join(CODE_DIR, "matting", "runtime")
    clash = ("utils", "pixelwise_estimator", "soft_foreground_segmenter")
    saved = {k: sys.modules.pop(k) for k in clash if k in sys.modules}
    sys.path.insert(0, runtime_dir)
    try:
        module = importlib.import_module("soft_foreground_segmenter")
        return module.SoftForegroundSegmenter
    finally:
        if runtime_dir in sys.path:
            sys.path.remove(runtime_dir)
        for k in clash:
            sys.modules.pop(k, None)
        sys.modules.update(saved)


SoftForegroundSegmenter = _load_soft_foreground_segmenter()

import cv2
import numpy as np
import yaml
import torch
import gradio as gr
import spaces
from PIL import Image
from huggingface_hub import hf_hub_download

from synthetic.align_utils import align_face, inverse_align_face

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

IMG_SIZE = 512
BORDER_SCALE = 0.85
DATA_IDX = 0

# ---------------------------------------------------------------------------
# Weight provisioning
# ---------------------------------------------------------------------------
def _dl(filename):
    # Local override for offline development; unused on the Space.
    local_dir = os.environ.get("OPENDELIGHT_WEIGHTS_DIR")
    if local_dir:
        path = os.path.join(local_dir, filename)
        if os.path.exists(path):
            return path
    return hf_hub_download(repo_id=WEIGHTS_REPO, filename=filename, token=HF_TOKEN)


def _ensure_mae_weights():
    """OpenDelight's MAE encoder loads a hardcoded relative path at construction
    time. Place the mirrored file exactly where the unmodified code expects it.
    The partial init is immediately overwritten by the full delight checkpoint,
    but the file must exist for the untouched code path to run."""
    target = os.path.join(CODE_DIR, "pretrained", "mae", "mae_visualize_vit_base.pth")
    if not os.path.exists(target):
        os.makedirs(os.path.dirname(target), exist_ok=True)
        src = _dl("mae_visualize_vit_base.pth")
        if os.path.abspath(src) != os.path.abspath(target):
            import shutil
            shutil.copy(src, target)


def _seed_facer_cache():
    """Pre-place the mirrored farl parser weights into the torch hub cache so
    facer loads them offline instead of downloading from GitHub at runtime."""
    fname = "face_parsing.farl.lapa.main_ema_136500_jit191.pt"
    ckpt_dir = os.path.join(torch.hub.get_dir(), "checkpoints")
    os.makedirs(ckpt_dir, exist_ok=True)
    target = os.path.join(ckpt_dir, fname)
    if not os.path.exists(target):
        src = _dl(fname)
        if os.path.abspath(src) != os.path.abspath(target):
            import shutil
            shutil.copy(src, target)


# ---------------------------------------------------------------------------
# Model construction (module scope -> weights download at container startup)
# ---------------------------------------------------------------------------
def _build_delight_model():
    from model.delight_base_net import DelightBaseModel
    with open(os.path.join(CODE_DIR, "config", "delight_base_network.yaml")) as f:
        cfg = yaml.safe_load(f)
    model = DelightBaseModel(enc_cfg=cfg["encoder"], dec_cfg=cfg["decoder"], device=DEVICE).to(DEVICE)
    weight = torch.load(_dl("base_delight_network.pth"), map_location=DEVICE)
    model.load_state_dict(weight, strict=True)
    model.eval()
    return model, cfg["encoder"]["type"]


def _build_enhancer():
    from model.detail_enhance_unet import UNet_Enhancer
    model = UNet_Enhancer(
        n_channels=6, n_classes=3, output_confidence=False, output_shadow_mask=False
    ).to(DEVICE)
    weight = torch.load(_dl("unet_enhancer.pth"), map_location=DEVICE)
    model.load_state_dict(weight, strict=True)
    model.eval()
    return model


class LandmarksDetectorIBug:
    def __init__(self, device):
        from ibug.face_detection import RetinaFacePredictor
        from ibug.face_alignment import FANPredictor

        det_model = RetinaFacePredictor.get_model("resnet50")
        det_model.weights = _dl("Resnet50_Final.pth")
        self.face_detector = RetinaFacePredictor(threshold=0.8, device=device, model=det_model)

        lmk_model = FANPredictor.get_model("2dfan2_alt")
        lmk_model.weights = _dl("2dfan2_alt.pth")
        # Disable TorchScript tracing: it runs a forward pass at construction,
        # which fails on ZeroGPU at module scope (no real GPU attached yet).
        lmk_cfg = FANPredictor.create_config(use_jit=False)
        self.landmark_detector = FANPredictor(device=device, model=lmk_model, config=lmk_cfg)

    def detect(self, images_bgr):
        detected_faces = self.face_detector(images_bgr, rgb=False)
        if len(detected_faces) == 0:
            return None
        landmarks, _ = self.landmark_detector(images_bgr, detected_faces, rgb=False)
        return landmarks


print("[OpenDelight] Provisioning weights and loading models ...", flush=True)
_ensure_mae_weights()
_seed_facer_cache()

import facer

DELIGHT_MODEL, ENC_TYPE = _build_delight_model()
ENHANCER = _build_enhancer()
LANDMARKS = LandmarksDetectorIBug(device=DEVICE)
MATTING = SoftForegroundSegmenter(
    onnx_model=_dl("foreground-segmentation-model-vitl16_384.onnx"),
    providers=["CPUExecutionProvider"],
)

# facer loads a TorchScript model via `torch.jit.load(map_location="cuda")`,
# which performs a real CUDA init that is not covered by the `spaces` torch
# monkeypatch. On ZeroGPU the GPU is only attached inside `@spaces.GPU`, so the
# parser must be created lazily on first use from within a GPU call.
_FACE_PARSER = None
_FACE_PARSER_LOCK = threading.Lock()


def _get_face_parser():
    global _FACE_PARSER
    if _FACE_PARSER is None:
        with _FACE_PARSER_LOCK:
            if _FACE_PARSER is None:
                _FACE_PARSER = facer.face_parser("farl/lapa/448", device=DEVICE)
    return _FACE_PARSER


print("[OpenDelight] Ready.", flush=True)


# ---------------------------------------------------------------------------
# Tensor <-> image helpers (mirrors test.py)
# ---------------------------------------------------------------------------
def _cv2_to_torch(img_cv2, device):
    img = torch.from_numpy(img_cv2).permute(2, 0, 1).float() / 255.0
    if img.shape[0] == 3:
        return img[None, [2, 1, 0], ...].to(device)
    if img.shape[0] == 1:
        return img[None, ...].to(device)
    raise NotImplementedError("cv2_to_torch supports 1 or 3 channel images only.")


def _torch_to_cv2(img_tensor):
    img = (img_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype("uint8")
    return img[..., [2, 1, 0]]


def _skin_matte(image_rgb, landmarks, matte):
    """Refine the foreground matte to the facial skin region (minus hair) using
    the FaRL face parser, following test.py::skin_mask_image."""
    lm = landmarks[0]
    kps = np.concatenate([
        np.mean(lm[36:42], axis=0, keepdims=True),
        np.mean(lm[42:48], axis=0, keepdims=True),
        lm[30:31],
        lm[48:49],
        lm[54:55],
    ], axis=0)
    kps = torch.from_numpy(kps).unsqueeze(0).to(DEVICE)

    image = facer.hwc2bchw(torch.from_numpy(np.ascontiguousarray(image_rgb))).to(DEVICE)
    with torch.inference_mode():
        faces = {"points": kps, "image_ids": torch.tensor([0], device=DEVICE)}
        faces = _get_face_parser()(image, faces)

    seg_probs = faces["seg"]["logits"].softmax(dim=1)[0]
    labels = seg_probs.argmax(dim=0)
    hair_mask = (labels == 10).float()
    face_mask = (labels >= 1).float() - hair_mask

    matte_t = torch.from_numpy(matte).to(DEVICE)
    refined = (matte_t * face_mask).clamp(0, 1)
    return refined.cpu().numpy()


@spaces.GPU(duration=60)
def _run_gpu(image_rgb, matte):
    """All GPU stages for a single image: landmarks, face parsing, alignment,
    delighting, enhancement, and inverse alignment. Runs in one ZeroGPU call."""
    image_bgr = np.ascontiguousarray(image_rgb[..., ::-1])

    landmarks = LANDMARKS.detect(image_bgr)
    if landmarks is None:
        return None

    refined_matte = _skin_matte(image_rgb, landmarks, matte)

    lm = landmarks[0]
    aligned_face, H = align_face(
        image_bgr.astype("float32"), output_size=IMG_SIZE, lm=lm, border_scale=BORDER_SCALE
    )
    aligned_face = aligned_face.astype("uint8")

    mask_3 = np.repeat((refined_matte * 255.0)[..., None], 3, axis=2).astype("float32")
    aligned_mask, _ = align_face(mask_3, output_size=IMG_SIZE, lm=lm, border_scale=BORDER_SCALE)
    aligned_mask = aligned_mask.astype("uint8")

    orig_size = (image_bgr.shape[1], image_bgr.shape[0])

    with torch.no_grad():
        aligned_face_t = _cv2_to_torch(aligned_face, DEVICE)
        aligned_mask_t = _cv2_to_torch(aligned_mask, DEVICE)
        input_face_t = aligned_face_t * aligned_mask_t

        input_512 = torch.nn.functional.interpolate(input_face_t, size=(512, 512), mode="bicubic")
        if ENC_TYPE == "mae_mix":
            output_t = DELIGHT_MODEL(input_512, torch.tensor([DATA_IDX]))
        else:
            output_t = DELIGHT_MODEL(input_512, None)
        output_t = torch.nn.functional.interpolate(output_t, size=(IMG_SIZE, IMG_SIZE), mode="bicubic")

        enhanced_t = ENHANCER(torch.cat([output_t, input_face_t], dim=1))["diffuse"]

    # Aligned (512) previews.
    delit_aligned = _torch_to_cv2(output_t)[..., ::-1].copy()      # -> RGB
    enhanced_aligned = _torch_to_cv2(enhanced_t)[..., ::-1].copy()  # -> RGB
    input_aligned = (input_face_t.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype("uint8")

    # Inverse-align back to the original resolution and composite alpha.
    enhanced_cv2 = _torch_to_cv2(enhanced_t)
    enh_res = inverse_align_face(enhanced_cv2.astype("float32"), H, orig_size).astype("uint8")

    res_mask = inverse_align_face(aligned_mask.astype("float32"), H, orig_size).astype("uint8")
    ones = np.ones_like(image_bgr, dtype="float32")
    inv_mask = inverse_align_face(ones, H, orig_size)
    inv_mask = (inv_mask[:, :, 0] > 0).astype("float32")

    alpha = (res_mask[:, :, 0].astype("float32") / 255.0) * inv_mask
    alpha = np.clip(alpha, 0, 1)

    enh_res_rgb = enh_res[..., ::-1].astype("float32")  # BGR->RGB
    rgba = np.dstack([enh_res_rgb, alpha * 255.0]).astype("uint8")

    white_bg = enh_res_rgb * alpha[..., None] + 255.0 * (1 - alpha[..., None])
    composited = white_bg.astype("uint8")

    matte_vis = (np.clip(matte, 0, 1) * 255).astype("uint8")

    return {
        "matte": matte_vis,
        "input_aligned": input_aligned,
        "delit_aligned": delit_aligned,
        "enhanced_aligned": enhanced_aligned,
        "composited": composited,
        "rgba": rgba,
    }


def predict(image_rgb):
    if image_rgb is None:
        raise gr.Error("Please provide an input image.")
    if image_rgb.ndim == 2:
        image_rgb = np.stack([image_rgb] * 3, axis=-1)
    image_rgb = np.ascontiguousarray(image_rgb[..., :3].astype("uint8"))

    # Matting runs on CPU (onnxruntime) outside the GPU worker.
    matte = MATTING.estimate_foreground_segmentation(
        np.ascontiguousarray(image_rgb[..., ::-1])
    ).astype("float32")
    if matte.ndim == 3:
        matte = matte[..., 0]

    out = _run_gpu(image_rgb, matte)
    if out is None:
        raise gr.Error("No face detected. Try a clearer, front-facing portrait.")

    rgba_path = os.path.join(tempfile.mkdtemp(), "delit_rgba.png")
    Image.fromarray(out["rgba"], mode="RGBA").save(rgba_path)

    gallery = [
        (image_rgb, "Input"),
        (out["matte"], "Foreground matte"),
        (out["input_aligned"], "Aligned + masked (512)"),
        (out["delit_aligned"], "Delit (network output)"),
        (out["enhanced_aligned"], "Delit + detail enhancer"),
        (out["composited"], "Composited on white (full res)"),
    ]
    return gallery, out["composited"], rgba_path


DESCRIPTION = """
# OpenDelight — Facial Appearance Delighting

Removes baked-in illumination and shadows from an in-the-wild face photo,
recovering a clean, evenly lit diffuse appearance. This demo runs the full
pipeline from [OpenDelight](https://github.com/yxuhan/OpenDelight)
(SIGGRAPH 2026): background matting, skin-region masking, FFHQ alignment,
the MAE-based delighting network, and a detail-enhancement UNet.

Upload a front-facing portrait (or pick an example) and press **Delight**.
"""

CITATION = """
```bibtex
@inproceedings{han2026opendelight,
    author = {Han, Yuxuan and Ming, Xin and Li, Tianxiao and Shen, Zhuofan and Zhang, Qixuan and Xu, Lan and Xu, Feng},
    title = {Learning a Delighting Prior for Facial Appearance Capture in the Wild},
    booktitle = {SIGGRAPH},
    year = {2026}
}
```
"""

with gr.Blocks(title="OpenDelight") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="numpy", label="Input portrait", image_mode="RGB")
            run_btn = gr.Button("Delight", variant="primary")
            gr.Examples(
                examples=sorted(glob.glob(os.path.join(CODE_DIR, "misc", "test_ffhq", "*.png"))),
                inputs=inp,
                cache_examples=False,
            )
        with gr.Column():
            out_main = gr.Image(type="numpy", label="Delit result (on white)")
            out_rgba = gr.File(label="Delit result (RGBA PNG)")
    gallery = gr.Gallery(label="Pipeline stages", columns=3, height="auto")
    gr.Markdown(CITATION)

    run_btn.click(predict, inputs=inp, outputs=[gallery, out_main, out_rgba])


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