app.py

"""FaceAnything — Gradio demo (Hugging Face Space).

Upload up to 40 face images (a short clip, in order). The model reconstructs the
clip in a single feed-forward pass and the app returns:

  * canonical 2D video  — per-frame canonical facial-coordinate map (original | map)
  * depth 2D video      — per-frame JET depth map
  * normals 2D video    — per-frame surface-normal map (from depth)
  * a colorful 3D point-track point cloud (.ply) you can orbit in the 3D viewer,
    with a frame slider to scrub through the sequence, plus a downloadable .zip
    of every frame's track point cloud.

Two inference modes are exposed (the repo's `--process-mode`):
  * Joint (all-at-once) — all frames processed together: more 3D-consistent.
  * One-by-one          — each frame independently: more surface detail, less
                          memory (pairs well with a higher processing resolution).

The heavy lifting reuses the published `faceanything` package unchanged; this app
only orchestrates it and renders the requested outputs. The expensive Open3D
orbit-video renderer is intentionally NOT used — the canonical/depth/normals
videos and the track point clouds are produced from cheap NumPy ops.
"""
from __future__ import annotations

import os
import sys
import glob
import shutil
import tempfile
import traceback

import numpy as np


# --------------------------------------------------------------------------- #
# Locate the published FaceAnything source.
#
# The model code (`src/faceanything`, `src/depth_anything_3`) is vendored into
# this Space. For local testing against a source checkout, point FACEANYTHING_ROOT
# at it instead:
#   export FACEANYTHING_ROOT=/cluster/eriador/ukocasari/projects/FaceAnything
# --------------------------------------------------------------------------- #
APP_DIR = os.path.dirname(os.path.abspath(__file__))
FA_ROOT = os.environ.get("FACEANYTHING_ROOT", APP_DIR)


def _ensure_faceanything_importable():
    """Make `import faceanything` work, trying a few sensible source locations."""
    try:
        import faceanything  # noqa: F401  (already installed / vendored)
        return
    except Exception:
        pass
    for cand in (os.path.join(APP_DIR, "src"), os.path.join(FA_ROOT, "src")):
        if os.path.isdir(os.path.join(cand, "faceanything")) and cand not in sys.path:
            sys.path.insert(0, cand)


_ensure_faceanything_importable()

BASE_MODEL = os.environ.get("FACEANYTHING_BASE_MODEL", "depth-anything/DA3-GIANT-1.1")
GPU_DURATION = int(os.environ.get("FACEANYTHING_GPU_DURATION", "120"))
MAX_IMAGES = int(os.environ.get("FACEANYTHING_MAX_IMAGES", "40"))


# --------------------------------------------------------------------------- #
# Checkpoint (~15 GB). Recommended storage: a separate HF *model* repo, pulled
# once with `hf_hub_download` and cached (point HF_HOME at persistent storage,
# e.g. /data/.huggingface, so it survives restarts). Resolution order:
#   1. FACEANYTHING_CHECKPOINT — an explicit local file (used if it exists)
#   2. FACEANYTHING_CHECKPOINT_REPO — download <FILE> from this HF repo
#      (private repos: set the HF_TOKEN secret)
#   3. checkpoints/checkpoint.pt next to the app (e.g. committed via Git LFS)
# --------------------------------------------------------------------------- #
def _resolve_checkpoint():
    explicit = os.environ.get("FACEANYTHING_CHECKPOINT")
    if explicit and os.path.exists(explicit):
        return explicit

    repo = os.environ.get("FACEANYTHING_CHECKPOINT_REPO")
    if repo:
        from huggingface_hub import hf_hub_download

        return hf_hub_download(
            repo_id=repo,
            filename=os.environ.get("FACEANYTHING_CHECKPOINT_FILE", "checkpoint.pt"),
            repo_type=os.environ.get("FACEANYTHING_CHECKPOINT_REPO_TYPE", "model"),
            revision=os.environ.get("FACEANYTHING_CHECKPOINT_REVISION") or None,
            token=os.environ.get("HF_TOKEN") or None,
        )

    default = os.path.join(FA_ROOT, "checkpoints", "checkpoint.pt")
    if os.path.exists(default):
        return default

    raise FileNotFoundError(
        "No checkpoint found. Set FACEANYTHING_CHECKPOINT to a local file, or "
        "FACEANYTHING_CHECKPOINT_REPO to a Hugging Face repo id (add the HF_TOKEN "
        "secret if it is private), or place checkpoint.pt under checkpoints/."
    )


# Resolve (and download, if from a repo) at startup — on the CPU node, so the
# 15 GB transfer never counts against ZeroGPU compute time. A failure here is
# non-fatal: the UI still builds and the clear error surfaces on first run.
try:
    CHECKPOINT = _resolve_checkpoint()
    print(f"[faceanything] checkpoint ready: {CHECKPOINT}", flush=True)
except Exception as _ckpt_err:  # noqa: BLE001
    CHECKPOINT = None
    print(f"[faceanything] checkpoint not ready yet: {_ckpt_err}", flush=True)


# --------------------------------------------------------------------------- #
# ZeroGPU decorator — falls back to a no-op when `spaces` is unavailable
# (e.g. running on a plain GPU box / cluster), so the same file runs anywhere.
# --------------------------------------------------------------------------- #
try:
    import spaces

    GPU = spaces.GPU
except Exception:  # pragma: no cover - only on non-Spaces hosts

    def GPU(func=None, **_kwargs):
        if callable(func):
            return func

        def _deco(f):
            return f

        return _deco


import gradio as gr


# --------------------------------------------------------------------------- #
# Model (loaded once, lazily, inside the GPU context and cached across calls)
# --------------------------------------------------------------------------- #
_MODEL = None
_MODEL_DEVICE = None


def _get_model(device: str):
    global _MODEL, _MODEL_DEVICE, CHECKPOINT
    if _MODEL is not None and _MODEL_DEVICE == device:
        return _MODEL
    from faceanything.model import load_model

    # Re-resolve if the startup attempt failed (e.g. env was set afterwards).
    ckpt = CHECKPOINT or _resolve_checkpoint()
    CHECKPOINT = ckpt
    _MODEL = load_model(ckpt, base_model=BASE_MODEL, device=device)
    _MODEL_DEVICE = device
    return _MODEL


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tif", ".tiff", ".gif")


def _natural_key(name):
    """Sort key that orders frame_2 before frame_10 (numeric-aware)."""
    import re

    return [int(c) if c.isdigit() else c.lower() for c in re.split(r"(\d+)", name)]


def _to_entry(f):
    """Normalize a Gradio file value (str / dict / FileData-like) to
    ``(temp_path, original_name)``. The original name drives ordering + extension;
    the temp path is what we actually copy from."""
    if isinstance(f, str):
        return f, os.path.basename(f)
    if isinstance(f, dict):
        path = f.get("path") or f.get("name")
        orig = f.get("orig_name") or (os.path.basename(path) if path else None)
        return path, orig
    path = getattr(f, "path", None) or getattr(f, "name", None)
    orig = getattr(f, "orig_name", None) or (os.path.basename(path) if path else None)
    return path, orig


def _sniff_ext(path):
    """Detect an image extension from file content (Gradio temp files often have
    no usable extension). Returns a safe default of .png if undetectable."""
    try:
        from PIL import Image

        with Image.open(path) as im:
            fmt = (im.format or "").lower()
        return {"jpeg": ".jpg", "png": ".png", "webp": ".webp", "bmp": ".bmp",
                "gif": ".gif", "tiff": ".tif", "mpo": ".jpg"}.get(fmt, ".png")
    except Exception:
        return ".png"


def _extract_video(video_path, max_frames, out_dir):
    """Decode the first ``max_frames`` frames of a video. Uses cv2.VideoCapture,
    which (unlike imageio's extension-based plugin pick) robustly decodes webcam
    recordings — those were yielding only a single frame otherwise."""
    import cv2

    os.makedirs(out_dir, exist_ok=True)
    paths = []
    cap = cv2.VideoCapture(video_path)
    try:
        while len(paths) < int(max_frames):
            ok, frame = cap.read()
            if not ok:
                break
            p = os.path.join(out_dir, f"frame_{len(paths):04d}.png")
            cv2.imwrite(p, frame)  # BGR ndarray -> correct-RGB PNG on disk
            paths.append(p)
    finally:
        cap.release()
    return paths


def _prepare_inputs(files, video, max_frames, workdir):
    """Normalize the upload (an image set OR a video) into an ordered list of
    frame paths.

    A video takes precedence: it is decoded and its first ``max_frames`` frames
    are used. For images we don't glob by extension (Gradio temp files often lack
    one): files are natural-sorted by their original name (temporal order),
    capped, then copied as ``frame_XXXX.<ext>`` with a content-sniffed extension.
    """
    if video:
        vpath = video if isinstance(video, str) else _to_entry(video)[0]
        if vpath and os.path.exists(vpath):
            paths = _extract_video(vpath, max_frames,
                                   os.path.join(workdir, "video_frames"))
            if len(paths) <= 1:  # fall back to imageio if cv2 read too few frames
                from faceanything.io_utils import load_frame_paths
                try:
                    alt, _ = load_frame_paths(
                        vpath, max_frames=int(max_frames), stride=1,
                        work_dir=os.path.join(workdir, "video_frames_io"))
                    if len(alt) > len(paths):
                        paths = alt
                except Exception:
                    pass
            if paths:
                return paths

    if not files:
        raise gr.Error("Please upload images or a video.")

    entries = []
    for f in files:
        path, orig = _to_entry(f)
        if path and os.path.exists(path):
            entries.append((path, orig or os.path.basename(path)))
    if not entries:
        raise gr.Error("Could not read the uploaded files — please re-upload your images.")

    entries.sort(key=lambda e: _natural_key(e[1]))
    entries = entries[:int(max_frames)]

    img_dir = os.path.join(workdir, "images")
    os.makedirs(img_dir, exist_ok=True)
    out = []
    for i, (path, orig) in enumerate(entries):
        ext = os.path.splitext(orig)[1].lower()
        if ext not in _IMAGE_EXTS:
            ext = _sniff_ext(path)
        dst = os.path.join(img_dir, f"frame_{i:04d}{ext}")
        shutil.copy(path, dst)
        out.append(dst)
    if not out:
        raise gr.Error("No valid images found in the upload.")
    return out


def _srgb_to_linear(cols_u8):
    """sRGB uint8 (0-255) -> linear uint8. glTF COLOR_0 vertex colors are
    interpreted as *linear* and the viewer re-applies the display gamma, so our
    sRGB image colors must be linearized first or the points render washed-out."""
    c = np.asarray(cols_u8, np.float32) / 255.0
    lin = np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)
    return np.clip(lin * 255.0, 0, 255).astype(np.uint8)


def _points_to_glb(path, points, colors, max_points=1_000_000):
    """Write a colored point cloud as a ``.glb`` — the format gradio's Model3D
    renders as points (a vertex-only ``.ply`` is treated as an empty solid mesh).
    ``points`` must already be in glTF axes; sRGB colors are linearized for glTF's
    linear color space. The full cloud is kept (``max_points`` is only an
    extreme-size safety cap, matching DA3's default) so it renders dense, not
    sparse."""
    import trimesh

    pts = np.asarray(points, np.float32)
    cols = np.asarray(colors)
    finite = np.isfinite(pts).all(axis=1)
    pts, cols = pts[finite], cols[finite]
    if pts.shape[0] == 0:  # keep the viewer from erroring on an empty frame
        pts = np.zeros((1, 3), np.float32)
        cols = np.full((1, 3), 200, np.uint8)
    if pts.shape[0] > max_points:
        idx = np.random.default_rng(0).choice(pts.shape[0], max_points, replace=False)
        pts, cols = pts[idx], cols[idx]
    if cols.dtype != np.uint8:
        cols = np.clip(cols, 0, 255).astype(np.uint8)
    rgb = _srgb_to_linear(cols[:, :3])
    rgba = np.concatenate(
        [rgb, np.full((rgb.shape[0], 1), 255, np.uint8)], axis=1)
    scene = trimesh.Scene()
    scene.add_geometry(trimesh.points.PointCloud(vertices=pts, colors=rgba))
    scene.export(path)
    return path


# --------------------------------------------------------------------------- #
# Face + hair segmentation (FacePerceiver/facer).
#
# The colorful tracks should land only on the facial area and hair, not on the
# neck / shoulders / clothing. facer's face parser (CelebAMask-HQ classes) gives
# us exactly that: we keep every class except background / neck / necklace /
# cloth / hat and use it to restrict the track seeds and recoloring.
# --------------------------------------------------------------------------- #
_FACER = {}


def _get_face_detector(device):
    """Lazily build & cache facer's RetinaFace detector (used for the face crop)."""
    if not _FACER.get("detector"):
        import facer
        _FACER["detector"] = facer.face_detector("retinaface/mobilenet", device=device)
    return _FACER["detector"]


def _get_face_parser(device):
    """Lazily build & cache facer's face detector + parser."""
    if not _FACER.get("parser"):
        import facer
        _FACER["parser"] = facer.face_parser("farl/celebm/448", device=device)
    return _get_face_detector(device), _FACER["parser"]


def _face_hair_masks(images, device, log):
    """Per-frame boolean (H,W) mask of the facial area + hair via facer.

    Returns a list aligned with ``images`` (an entry is ``None`` when no face was
    detected for that frame), or ``None`` entirely when facer is unavailable —
    the caller then falls back to unrestricted tracks."""
    try:
        import torch
        import facer
    except Exception as e:  # facer / its deps not installed
        log.append(f"WARNING: facer unavailable ({e}); colorful tracks are not "
                   f"restricted to face + hair.")
        return None
    try:
        detector, parser = _get_face_parser(device)
    except Exception as e:
        log.append(f"WARNING: could not load facer models ({e}); colorful tracks "
                   f"are not restricted to face + hair.")
        return None

    def _is_excluded(name):
        n = name.lower()
        return any(b in n for b in ("background", "neck", "cloth", "hat"))

    masks, n_ok = [], 0
    for img in images:
        try:
            t = facer.hwc2bchw(torch.from_numpy(np.ascontiguousarray(img))).to(device)
            with torch.inference_mode():
                faces = detector(t)
                rects = faces.get("rects") if faces else None
                if rects is None or len(rects) == 0:
                    masks.append(None)
                    continue
                faces = parser(t, faces)
            seg = faces["seg"]
            labels = seg["label_names"]
            argmax = seg["logits"].softmax(dim=1).argmax(dim=1)  # (nfaces, H, W)
            keep = [ci for ci, nm in enumerate(labels) if not _is_excluded(nm)]
            m = torch.zeros(argmax.shape[-2:], dtype=torch.bool, device=argmax.device)
            for f in range(argmax.shape[0]):
                for ci in keep:
                    m |= (argmax[f] == ci)
            masks.append(m.cpu().numpy())
            n_ok += 1
        except Exception:
            masks.append(None)
    if n_ok == 0:
        log.append("WARNING: facer detected no faces; colorful tracks are not "
                   "restricted to face + hair.")
        return None
    log.append(f"Face + hair segmentation (facer): {n_ok}/{len(images)} frame(s).")
    return masks


# --------------------------------------------------------------------------- #
# Face-centric cropping (pixel3dmm-style).
#
# Mirrors SimonGiebenhain/pixel3dmm `get_cstm_crop` (scripts/run_cropping.py +
# preprocessing/pipnet_utils.py): detect a face box, square it, expand it ~1.42x
# (or 1.1x of the clip's union box when the face moves a lot), clamp to the image,
# crop and resize. For a clip we compute ONE static box (mean + union over frames)
# so the crop is temporally stable. We reuse facer's RetinaFace detector for the
# box (no extra PIPNet/FaceBoxes weights). Cropping focuses the model's pixels on
# the face instead of the background / body.
# --------------------------------------------------------------------------- #
def _cstm_crop_box(mean_b, max_b, img_h, img_w, scale=1.42):
    """pixel3dmm get_cstm_crop → (ymin, ymax, xmin, xmax). Boxes are (x, y, w, h)."""
    det = list(mean_b); s = scale
    if det[2] * scale * det[3] * scale < max_b[2] * 1.1 * max_b[3] * 1.1:
        det = list(max_b); s = 1.1
    xmin, ymin, dw, dh = det
    if dw > dh:                      # square it: grow the shorter side symmetrically
        ymin -= (dw - dh) / 2.0; dh = dw
    elif dw < dh:
        xmin -= (dh - dw) / 2.0; dw = dh
    xmax = xmin + dw - 1; ymax = ymin + dh - 1
    xmin -= dw * (s - 1) / 2.0; ymin -= dh * (s - 1) / 2.0   # expand by the scale
    xmax += dw * (s - 1) / 2.0; ymax += dh * (s - 1) / 2.0
    if xmin < 0 or ymin < 0:         # shift inside the image, preserving the square
        o = min(xmin, ymin); xmin -= o; ymin -= o
    if xmax > img_w - 1 or ymax > img_h - 1:
        o = max(xmax - (img_w - 1), ymax - (img_h - 1)); xmax -= o; ymax -= o
    xmin = max(int(round(xmin)), 0); ymin = max(int(round(ymin)), 0)
    xmax = min(int(round(xmax)), img_w - 1); ymax = min(int(round(ymax)), img_h - 1)
    return ymin, ymax, xmin, xmax


def _combine_face_hair_box(face_box, hair_bbox, img_h, img_w,
                           pad_top=0.06, pad_side=0.03, pad_bot=0.03,
                           max_aspect=1.5):
    """Square crop around the face + hair segmentation bbox. The model performs
    better on square inputs, so we square the box — but base it on the *tight*
    face+hair mask (not the expanded detection box) to keep the inherent
    left/right background (a head is taller than wide) to a minimum.

    side = larger padded box dim, but capped at ``max_aspect`` * head width so a
    very tall head doesn't produce huge side margins; the box is centered
    horizontally and anchored at the bottom, so the chin/face is always kept and
    only a little hair-top is dropped when the cap bites ("hair is mostly in").
    Falls back to the face box when no hair bbox is available.

    face_box: (ymin, ymax, xmin, xmax). hair_bbox: (x0, y0, x1, y1) or None."""
    if hair_bbox is None:
        return face_box
    hx0, hy0, hx1, hy1 = hair_bbox
    bw = max(hx1 - hx0, 1.0); bh = max(hy1 - hy0, 1.0)
    x0 = hx0 - pad_side * bw; x1 = hx1 + pad_side * bw
    y0 = hy0 - pad_top * bh;  y1 = hy1 + pad_bot * bh
    bw_p, bh_p = x1 - x0, y1 - y0
    side = min(max(bw_p, bh_p), max_aspect * bw_p, float(img_w), float(img_h))
    nx0 = (x0 + x1) / 2.0 - side / 2.0     # centered horizontally on the head
    ny0 = y1 - side                         # anchored at the bottom (keep the chin)
    nx0 = min(max(nx0, 0.0), img_w - side)
    ny0 = min(max(ny0, 0.0), img_h - side)
    xmin = int(round(nx0)); ymin = int(round(ny0)); side = int(round(side))
    return ymin, min(ymin + side, img_h - 1), xmin, min(xmin + side, img_w - 1)


def _face_crop_frames(frame_paths, device, out_dir, log, process_res):
    """Crop every frame to a face-centric square (pixel3dmm-style), grown to also
    cover the hair (top of head + long hair on the sides) via the face+hair
    segmentation. Returns new frame paths; on any failure (facer missing / no face)
    returns the originals so the run never breaks."""
    import cv2
    try:
        import torch
        import facer
        detector = _get_face_detector(device)
    except Exception as e:
        log.append(f"WARNING: face crop unavailable ({e}); using full frames.")
        return frame_paths

    imgs = [cv2.imread(fp) for fp in frame_paths]
    rgb_imgs = [cv2.cvtColor(im, cv2.COLOR_BGR2RGB) if im is not None else None
                for im in imgs]
    sizes = [im.shape[:2] if im is not None else None for im in imgs]

    # face detection boxes (x, y, w, h)
    boxes = []
    for rgb in rgb_imgs:
        if rgb is None:
            boxes.append(None); continue
        try:
            t = facer.hwc2bchw(torch.from_numpy(np.ascontiguousarray(rgb))).to(device)
            with torch.inference_mode():
                faces = detector(t)
            rects = faces.get("rects") if faces else None
            if rects is None or len(rects) == 0:
                boxes.append(None); continue
            scores = faces.get("scores")
            bi = int(scores.argmax()) if scores is not None and len(scores) else 0
            x1, y1, x2, y2 = [float(v) for v in rects[bi].tolist()]
            boxes.append([x1, y1, x2 - x1, y2 - y1])        # x, y, w, h
        except Exception:
            boxes.append(None)

    valid = [b for b in boxes if b is not None]
    if not valid:
        log.append("WARNING: face crop found no faces; using full frames.")
        return frame_paths

    # face + hair mask → bbox per frame, so the crop encloses the hair, not just
    # the face detection box (which starts around the hairline).
    hair_masks = _face_hair_masks(rgb_imgs, device, log)

    def _mask_bbox(m):
        if m is None:
            return None
        ys, xs = np.nonzero(m)
        if not len(xs):
            return None
        return [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())]

    hair_bboxes = ([_mask_bbox(m) for m in hair_masks]
                   if hair_masks is not None else [None] * len(frame_paths))

    # One static box for the whole clip when every frame shares a resolution.
    uniq = set(s for s in sizes if s is not None)
    static_box = None
    if len(uniq) == 1:
        H, W = next(iter(uniq))
        xs = np.array([b[0] for b in valid]); ys = np.array([b[1] for b in valid])
        ws = np.array([b[2] for b in valid]); hs = np.array([b[3] for b in valid])
        x0, y0 = xs.min(), ys.min()
        mean_b = [xs.mean(), ys.mean(), ws.mean(), hs.mean()]
        max_b = [x0, y0, (xs + ws - x0).max(), (ys + hs - y0).max()]   # union box
        face_static = _cstm_crop_box(mean_b, max_b, H, W)
        hb = [b for b in hair_bboxes if b is not None]
        hair_static = ([min(b[0] for b in hb), min(b[1] for b in hb),
                        max(b[2] for b in hb), max(b[3] for b in hb)] if hb else None)
        static_box = _combine_face_hair_box(face_static, hair_static, H, W)

    out_size = int(min(1024, max(512, int(process_res))))
    crop_dir = os.path.join(out_dir, "cropped")
    os.makedirs(crop_dir, exist_ok=True)
    fallback = [np.mean([b[0] for b in valid]), np.mean([b[1] for b in valid]),
                np.mean([b[2] for b in valid]), np.mean([b[3] for b in valid])]
    new_paths, n_cropped = [], 0
    for i, fp in enumerate(frame_paths):
        im = imgs[i]
        if im is None:
            new_paths.append(fp); continue
        h, w = im.shape[:2]
        if static_box is not None:
            ymin, ymax, xmin, xmax = static_box
        else:
            b = boxes[i] if boxes[i] is not None else fallback
            face_box = _cstm_crop_box(b, b, h, w)
            ymin, ymax, xmin, xmax = _combine_face_hair_box(
                face_box, hair_bboxes[i], h, w)
        crop = im[ymin:ymax, xmin:xmax]
        out_fp = os.path.join(crop_dir, f"frame_{i:04d}.png")
        if crop.size == 0:
            cv2.imwrite(out_fp, im)
        else:
            # keep aspect ratio (the model resizes the longest side itself); only
            # downscale if the crop is larger than we need.
            ch, cw = crop.shape[:2]
            longest = max(ch, cw)
            if longest > out_size:
                sc = out_size / float(longest)
                crop = cv2.resize(crop, (max(1, round(cw * sc)), max(1, round(ch * sc))))
            cv2.imwrite(out_fp, crop)
            n_cropped += 1
        new_paths.append(out_fp)
    log.append(
        f"Face crop (pixel3dmm-style, hair-aware): {n_cropped}/{len(frame_paths)} "
        f"frame(s) → {out_size}x{out_size}"
        + (" static box." if static_box is not None else " per-frame."))
    return new_paths


@GPU(duration=GPU_DURATION)
def run(
    files,
    video,
    mode,
    process_res,
    remove_bg,
    face_crop,
    conf_percentile,
    n_tracks,
    track_k,
    track_threshold,
    fps,
    max_frames,
    progress=gr.Progress(),
):
    """End-to-end inference + visualization. Returns the 5 outputs + viewer state."""
    import torch

    # Imported here so the UI still builds even if heavy deps are missing.
    from faceanything.predict import run_inference
    from faceanything.geometry import (
        point_cloud_from_depth,
        unproject_depth,
        pointmap_to_normals,
    )
    from faceanything.colorize import (
        depth_to_jet,
        normals_to_rgb,
        canonical_to_rgb,
    )
    from faceanything.tracking import compute_track_colors
    from faceanything.export import save_ply
    from faceanything.render import side_by_side, write_video

    device = "cuda" if torch.cuda.is_available() else "cpu"
    workdir = tempfile.mkdtemp(prefix="faceanything_demo_")
    log = []

    def _say(frac, msg):
        log.append(msg)
        progress(frac, desc=msg)

    try:
        _say(0.02, "Preparing inputs…")
        frame_paths = _prepare_inputs(
            files, video, min(int(max_frames), MAX_IMAGES), workdir)
        n_in = len(frame_paths)
        if n_in == 0:
            raise gr.Error("No valid images found in the upload.")
        log.append(f"{n_in} frame(s) | mode: {mode} | process_res: {int(process_res)}")

        # ---- face-centric crop (optional, pixel3dmm-style) ----
        # Done before background removal so the masks align with the cropped frames.
        if face_crop:
            _say(0.05, "Cropping to the face (pixel3dmm-style)…")
            frame_paths = _face_crop_frames(frame_paths, device, workdir, log,
                                            process_res)

        # ---- background removal (optional) ----
        mask_paths = None
        if remove_bg:
            _say(0.08, "Removing background (Robust Video Matting)…")
            from faceanything.background import generate_masks

            try:
                mask_paths = generate_masks(
                    frame_paths, os.path.join(workdir, "masks"), device=device
                )
            except Exception as bg_err:  # don't let RVM take down the whole run
                mask_paths = None
                log.append(f"WARNING: background removal failed ({bg_err}); "
                           f"reconstructing the full frame instead.")

        # ---- model + inference ----
        _say(0.15, "Loading model (first run downloads/loads the checkpoint)…")
        model = _get_model(device)
        _say(0.30, f"Running inference on {n_in} frame(s)…")
        pred = run_inference(
            model,
            frame_paths,
            mask_paths=mask_paths,
            process_res=int(process_res),
            monocular=False,  # always use predicted camera poses (world frame)
            conf_percentile=float(conf_percentile),
            per_frame=(mode == "One-by-one"),
        )
        N = int(pred.depth.shape[0])
        has_canon = pred.canonical is not None
        log.append(f"Inference done: {N} frame(s), depth {tuple(pred.depth.shape[1:])}, "
                   f"canonical: {has_canon}")

        # ---- per-frame clouds + global color ranges (mirrors run_inference.py) ----
        _say(0.55, "Building point clouds and color maps…")
        clouds = []
        for i in range(N):
            pts, rgb, canon, pix = point_cloud_from_depth(
                pred.depth[i], pred.images[i], pred.intrinsics[i],
                extrinsics=pred.extrinsics[i], valid_mask=pred.valid[i],
                deformation=pred.canonical[i] if has_canon else None,
            )
            depth_vals = pred.depth[i][pix[:, 0], pix[:, 1]]
            clouds.append(dict(points=pts, rgb=rgb, canonical=canon,
                               depth_vals=depth_vals, pix=pix))

        all_depth = (np.concatenate([c["depth_vals"] for c in clouds])
                     if N else np.zeros(1))
        dmin, dmax = (np.percentile(all_depth, [2, 98]) if all_depth.size else (0.0, 1.0))
        if dmax <= dmin:
            dmax = dmin + 1e-6
        canon_ranges = None
        if has_canon:
            allc = np.concatenate([c["canonical"] for c in clouds
                                   if c["canonical"] is not None])
            _, canon_ranges = canonical_to_rgb(allc.reshape(-1, 1, 3), None)

        # ---- 2D maps (image space) ----
        def frame2d(modality, i):
            v = pred.valid[i]
            if modality == "depth":
                return depth_to_jet(pred.depth[i], v, dmin, dmax)
            if modality == "normals":
                nmap = pointmap_to_normals(
                    unproject_depth(pred.depth[i], pred.intrinsics[i], None)[0])
                img = normals_to_rgb(nmap)
                img[~v] = 255
                return img
            if modality == "canonical":
                img, _ = canonical_to_rgb(pred.canonical[i], v, ranges=canon_ranges)
                return img
            raise ValueError(modality)

        vids_dir = os.path.join(workdir, "videos")
        os.makedirs(vids_dir, exist_ok=True)

        def make_2d_video(modality):
            seq = [side_by_side(pred.images[i], frame2d(modality, i)) for i in range(N)]
            seq = seq * 30 if len(seq) == 1 else seq  # avoid 1-frame videos
            out = os.path.join(vids_dir, f"{modality}_2d.mp4")
            write_video(seq, out, fps=int(fps))
            return out

        _say(0.65, "Rendering depth 2D video…")
        depth_vid = make_2d_video("depth")
        _say(0.72, "Rendering normals 2D video…")
        normals_vid = make_2d_video("normals")
        canonical_vid = None
        if has_canon:
            _say(0.79, "Rendering canonical 2D video…")
            canonical_vid = make_2d_video("canonical")

        # ---- colorful point tracks (canonical NN matching) ----
        tracks_zip = None
        tracks2d_vid = None
        view_glbs = []
        if has_canon:
            # face + hair mask (facer) → restrict the colorful tracks to the face
            # and hair, never the neck / shoulders / clothing.
            _say(0.83, "Segmenting face + hair (facer)…")
            face_masks = _face_hair_masks(pred.images, device, log)
            regions, seed_frame = None, 0
            if face_masks is not None:
                regions = []
                for i, c in enumerate(clouds):
                    pix = c["pix"]
                    fm = face_masks[i]
                    if fm is None:
                        regions.append(np.zeros(pix.shape[0], bool))
                    else:
                        regions.append(fm[pix[:, 0], pix[:, 1]])
                sizes = [int(r.sum()) for r in regions]
                if sizes and max(sizes) > 0:
                    seed_frame = int(np.argmax(sizes))  # seed where the face is biggest
                else:
                    regions = None  # nothing usable — don't restrict

            _say(0.86, f"Computing {int(n_tracks)} colorful point tracks…")
            track_colors, track_overlay = compute_track_colors(
                [dict(canonical=c["canonical"], rgb=c["rgb"], pix=c["pix"])
                 for c in clouds],
                n_tracks=int(n_tracks), k=int(track_k),
                threshold=float(track_threshold),
                regions=regions, seed_frame=seed_frame,
            )

            import trimesh
            # glTF alignment, shared across frames so the slider view stays stable:
            # orient to the first camera, flip Y/Z (OpenCV -> glTF), center by median.
            w2c0 = pred.extrinsics[0].astype(np.float64)
            A = np.diag([1.0, -1.0, -1.0, 1.0]) @ w2c0
            all_pts = (np.concatenate([c["points"] for c in clouds])
                       if N else np.zeros((1, 3)))
            center = np.median(trimesh.transform_points(all_pts, A), axis=0)
            T = np.eye(4); T[:3, 3] = -center
            A = T @ A

            # Downloadable .ply (repo coords): two colorings in the same zip.
            # Viewer .glb (glTF-aligned): one set track-colored, one set plain RGB
            # — the viewer toggles between them client-side (≤ 2·N ≈ 80 files).
            tracks_dir = os.path.join(workdir, "pointclouds", "tracks")
            points_dir = os.path.join(workdir, "pointclouds", "points")
            view_dir = os.path.join(workdir, "anim_glb")
            for d in (tracks_dir, points_dir, view_dir):
                os.makedirs(d, exist_ok=True)
            track_glbs, rgb_glbs = [], []
            for i in range(N):
                pts = clouds[i]["points"]
                save_ply(os.path.join(tracks_dir, f"frame_{i:04d}.ply"),
                         pts, track_colors[i])                  # colorful tracks
                save_ply(os.path.join(points_dir, f"frame_{i:04d}.ply"),
                         pts, clouds[i]["rgb"])                 # plain colored points
                aligned = trimesh.transform_points(pts, A)
                tg = os.path.join(view_dir, f"track_{i:04d}.glb")
                rg = os.path.join(view_dir, f"rgb_{i:04d}.glb")
                _points_to_glb(tg, aligned, track_colors[i])    # colorful tracks
                _points_to_glb(rg, aligned, clouds[i]["rgb"])   # image RGB colors
                track_glbs.append(tg)
                rgb_glbs.append(rg)
            view_glbs = track_glbs + rgb_glbs
            tracks_zip = shutil.make_archive(
                os.path.join(workdir, "pointclouds"), "zip",
                os.path.join(workdir, "pointclouds"))

            # bonus: 2D track overlay video (colorful seeds on the original frames)
            _say(0.93, "Rendering 2D track overlay video…")

            def _paint(img, pix, col, radius):
                H, W = img.shape[:2]
                for dr in range(-radius, radius + 1):
                    for dc in range(-radius, radius + 1):
                        rr = np.clip(pix[:, 0] + dr, 0, H - 1)
                        cc = np.clip(pix[:, 1] + dc, 0, W - 1)
                        img[rr, cc] = col

            t_seq = []
            for i in range(N):
                img = pred.images[i].copy()
                img[~pred.valid[i]] = 255
                pix, col = track_overlay[i]
                if pix.shape[0]:
                    _paint(img, pix, col, radius=max(2, round(img.shape[0] / 160)))
                t_seq.append(side_by_side(pred.images[i], img))
            t_seq = t_seq * 30 if len(t_seq) == 1 else t_seq
            tracks2d_vid = os.path.join(vids_dir, "tracks_2d.mp4")
            write_video(t_seq, tracks2d_vid, fps=int(fps))

        _say(1.0, "Done.")

        if not has_canon:
            log.append("WARNING: model produced no canonical output — canonical "
                       "video and point tracks were skipped (check the checkpoint).")

        status = "\n".join(f"• {m}" for m in log)
        # view_glbs is the per-frame track point cloud (.glb), glTF-aligned and
        # track-colored. They go to the hidden file list, whose URLs the
        # client-side three.js player preloads and animates (see VIEWER_JS).
        return (
            view_glbs or None,
            canonical_vid,
            depth_vid,
            normals_vid,
            tracks2d_vid,
            tracks_zip,
            status,
        )
    except gr.Error:
        raise
    except Exception as e:  # surface the traceback in the UI instead of a blank fail
        tb = traceback.format_exc()
        raise gr.Error(f"Inference failed: {e}\n\n{tb[-1500:]}")


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
DESCRIPTION = """
# Face Anything: 4D Face Reconstruction from Any Image Sequence

Upload **up to 40 face images** (a short clip, named so they sort in order).
The model jointly predicts depth and **canonical facial coordinates** in a single
feed-forward pass, from which we derive canonical / depth / normal maps and dense,
temporally-consistent **3D point tracks**.

[Project page](https://kocasariumut.github.io/FaceAnything/) ·
[arXiv](https://arxiv.org/abs/2604.19702) ·
[Code](https://github.com/kocasariumut/FaceAnything)
"""

# --------------------------------------------------------------------------- #
# Custom 3D viewer (client-side three.js).
#
# gradio's Model3D re-fetches and re-parses a .glb from the server on every
# frame, so animating it flashes white (the next cloud isn't on the client yet).
# Instead we load *every* frame's .glb once into a three.js scene and animate by
# toggling which frame is visible — no per-frame network/parse, the points stay
# on screen the whole time, and the full (un-subsampled) cloud is kept.
# --------------------------------------------------------------------------- #
THREE_HEAD = """
<script src="https://cdn.jsdelivr.net/npm/three@0.137.0/build/three.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/three@0.137.0/examples/js/loaders/GLTFLoader.js"></script>
<script src="https://cdn.jsdelivr.net/npm/three@0.137.0/examples/js/controls/OrbitControls.js"></script>
"""

VIEWER_MARKUP = """
<div class="fa-viewer-root" style="width:100%;">
  <div class="fa-canvas-wrap" style="position:relative;width:100%;height:420px;background:#ffffff;border:1px solid #e5e7eb;border-radius:8px;overflow:hidden;">
    <div class="fa-overlay" style="position:absolute;inset:0;display:flex;align-items:center;justify-content:center;color:#6b7280;font-family:sans-serif;font-size:14px;text-align:center;padding:0 16px;">Run a reconstruction to view the 3D point tracks here.</div>
  </div>
  <div class="fa-controls" style="display:none;gap:10px;align-items:center;padding:8px 4px 2px;font-family:sans-serif;font-size:13px;flex-wrap:wrap;">
    <button class="fa-play" type="button" style="cursor:pointer;padding:4px 12px;border:1px solid #d1d5db;border-radius:6px;background:#f9fafb;">&#9654; Play</button>
    <label style="display:inline-flex;align-items:center;gap:5px;cursor:pointer;" title="On: colorful tracks on face + hair. Off: image RGB colors."><input class="fa-tracks" type="checkbox"> Colorful tracks</label>
    <span style="display:inline-flex;align-items:center;gap:6px;">Speed <input class="fa-speed" type="range" min="1" max="30" step="1" value="12" style="width:90px;vertical-align:middle;"><span class="fa-fps">12 fps</span></span>
    <input class="fa-scrub" type="range" min="0" max="0" step="1" value="0" style="flex:1;min-width:120px;vertical-align:middle;">
    <span class="fa-frame" style="min-width:64px;text-align:right;color:#374151;">&ndash; / &ndash;</span>
  </div>
</div>
"""

# Runs once when the HTML component mounts; sets up the three.js scene and
# exposes window.faViewer.load(items) for the file bridge below to call. It holds
# two parallel per-frame sets — colorful tracks and image-RGB — switched instantly
# client-side by the "Colorful tracks" checkbox.
VIEWER_JS = """
(function(){
  if (element.__faInit) return;
  element.__faInit = true;
  var wrap = element.querySelector('.fa-canvas-wrap');
  var overlay = element.querySelector('.fa-overlay');
  var controls = element.querySelector('.fa-controls');
  var playBtn = element.querySelector('.fa-play');
  var tracksEl = element.querySelector('.fa-tracks');
  var speedEl = element.querySelector('.fa-speed');
  var fpsEl = element.querySelector('.fa-fps');
  var scrubEl = element.querySelector('.fa-scrub');
  var frameEl = element.querySelector('.fa-frame');
  if (!wrap) return;
  var renderer, scene, camera, orbit, group;
  var setTracks = [], setRgb = [], mode = 'rgb';  // default: plain image RGB
  var cur = 0, playing = false, fps = 12, acc = 0, last = 0, loadToken = 0;

  function activeSet(){
    var a = (mode === 'rgb') ? setRgb : setTracks;
    if (!a.length) a = (mode === 'rgb') ? setTracks : setRgb;  // fall back if empty
    return a;
  }
  function waitThree(cb, tries){
    tries = tries || 0;
    if (window.THREE && THREE.GLTFLoader && THREE.OrbitControls) { cb(); }
    else if (tries > 200) { setOverlay('Could not load the 3D viewer (three.js) \\u2014 check your network / ad-blocker.'); }
    else { setTimeout(function(){ waitThree(cb, tries + 1); }, 60); }
  }
  function setOverlay(msg){
    if (!overlay) return;
    if (msg) { overlay.textContent = msg; overlay.style.display = 'flex'; }
    else { overlay.style.display = 'none'; }
  }
  function resize(){
    if (!renderer) return;
    var w = wrap.clientWidth || 1, h = wrap.clientHeight || 1;
    renderer.setSize(w, h, false);
    camera.aspect = w / h; camera.updateProjectionMatrix();
  }
  function applyStyle(root){
    root.traverse(function(o){
      if (o.isPoints && o.material){
        o.material.size = 2.5;
        o.material.sizeAttenuation = false;
        o.material.vertexColors = true;
        o.material.needsUpdate = true;
      }
    });
  }
  function firstPoints(root){
    var found = null;
    root.traverse(function(o){ if (!found && o.isPoints) found = o; });
    return found;
  }
  function hideAll(){
    var k;
    for (k = 0; k < setTracks.length; k++){ if (setTracks[k]) setTracks[k].visible = false; }
    for (k = 0; k < setRgb.length; k++){ if (setRgb[k]) setRgb[k].visible = false; }
  }
  function showFrame(i){
    var arr = activeSet();
    if (!arr.length) return;
    if (i < 0) i = 0;
    if (i > arr.length - 1) i = arr.length - 1;
    hideAll();
    if (arr[i]) arr[i].visible = true;
    cur = i;
    if (scrubEl) scrubEl.value = String(i);
    if (frameEl) frameEl.textContent = (i + 1) + ' / ' + arr.length;
  }
  function fitCamera(){
    var ref = setTracks[0] || setRgb[0];
    if (!ref) return;
    var box = new THREE.Box3().setFromObject(ref);
    if (box.isEmpty()) return;
    var c = box.getCenter(new THREE.Vector3());
    var s = box.getSize(new THREE.Vector3());
    var r = Math.max(s.x, s.y, s.z) * 0.5 || 0.5;
    var d = (r / Math.tan(camera.fov * Math.PI / 360)) * 1.15;
    orbit.target.copy(c);
    camera.near = Math.max(d / 200, 0.0005);
    camera.far = d * 50 + r * 20;
    camera.position.set(c.x, c.y, c.z + d);
    camera.updateProjectionMatrix();
    orbit.update();
  }
  function clearFrames(){
    var arrs = [setTracks, setRgb], a, k, o;
    for (a = 0; a < arrs.length; a++){
      for (k = 0; k < arrs[a].length; k++){
        o = arrs[a][k];
        if (!o) continue;
        group.remove(o);
        if (o.geometry) o.geometry.dispose();
        if (o.material) o.material.dispose();
      }
    }
    setTracks = []; setRgb = []; cur = 0;
  }
  function play(){ if (activeSet().length < 2) return; playing = true; last = 0; acc = 0; if (playBtn) playBtn.innerHTML = '\\u23F8 Pause'; }
  function pause(){ playing = false; if (playBtn) playBtn.innerHTML = '\\u25B6 Play'; }
  function toggle(){ if (playing) pause(); else play(); }
  function animate(ts){
    requestAnimationFrame(animate);
    if (orbit) orbit.update();
    var arr = activeSet();
    if (playing && arr.length > 1){
      if (!last) last = ts;
      acc += (ts - last); last = ts;
      var interval = 1000 / Math.max(1, fps);
      if (acc >= interval){
        var steps = Math.floor(acc / interval);
        acc -= steps * interval;
        showFrame((cur + steps) % arr.length);
      }
    } else { last = ts; }
    if (renderer && scene && camera) renderer.render(scene, camera);
  }
  function initThree(){
    if (renderer) return;
    renderer = new THREE.WebGLRenderer({ antialias: true, alpha: false });
    renderer.setPixelRatio(1);
    renderer.setClearColor(0xffffff, 1);
    if (THREE.sRGBEncoding) renderer.outputEncoding = THREE.sRGBEncoding;
    renderer.domElement.style.display = 'block';
    renderer.domElement.style.width = '100%';
    renderer.domElement.style.height = '100%';
    wrap.appendChild(renderer.domElement);
    scene = new THREE.Scene();
    scene.background = new THREE.Color(0xffffff);
    camera = new THREE.PerspectiveCamera(50, 1, 0.001, 1000);
    camera.position.set(0, 0, 2);
    orbit = new THREE.OrbitControls(camera, renderer.domElement);
    orbit.enableDamping = true; orbit.dampingFactor = 0.1;
    group = new THREE.Group();
    scene.add(group);
    resize();
    if (window.ResizeObserver) { new ResizeObserver(resize).observe(wrap); }
    else { window.addEventListener('resize', resize); }
    requestAnimationFrame(animate);
  }
  function classify(name){
    var lower = (name || '').toLowerCase();
    var kind = (lower.indexOf('rgb') !== -1) ? 'rgb' : 'tracks';
    var d = (lower.match(/[0-9]+/g) || []).join('');
    return { kind: kind, idx: d ? parseInt(d, 10) : 0 };
  }
  function compact(map){
    var keys = Object.keys(map).map(Number).sort(function(a, b){ return a - b; });
    var out = [];
    for (var i = 0; i < keys.length; i++){ if (map[keys[i]]) out.push(map[keys[i]]); }
    return out;
  }
  function finishLoad(token, byTracks, byRgb){
    if (token !== loadToken) return;
    setTracks = compact(byTracks);
    setRgb = compact(byRgb);
    if (!setTracks.length && !setRgb.length){ setOverlay('Failed to load the 3D point clouds.'); return; }
    var all = setTracks.concat(setRgb), k;
    for (k = 0; k < all.length; k++){ all[k].visible = false; group.add(all[k]); }
    if (tracksEl){
      if (!setRgb.length && setTracks.length) tracksEl.checked = true;
      else if (!setTracks.length) tracksEl.checked = false;
      mode = tracksEl.checked ? 'tracks' : 'rgb';
    }
    setOverlay('');
    if (controls) controls.style.display = 'flex';
    if (scrubEl){ scrubEl.min = '0'; scrubEl.max = String(Math.max(0, activeSet().length - 1)); scrubEl.value = '0'; }
    fitCamera();
    showFrame(0);
    if (activeSet().length > 1) play();
  }
  function load(items){
    waitThree(function(){
      initThree();
      var token = ++loadToken;
      pause();
      clearFrames();
      if (controls) controls.style.display = 'none';
      if (!items || !items.length){ setOverlay('No 3D point tracks for this run.'); return; }
      setOverlay('Loading 3D sequence\\u2026 0 / ' + items.length);
      var loader = new THREE.GLTFLoader();
      var byTracks = {}, byRgb = {};
      var done = 0, total = items.length;
      function tick(){
        done++;
        setOverlay('Loading 3D sequence\\u2026 ' + done + ' / ' + total);
        if (done === total) finishLoad(token, byTracks, byRgb);
      }
      items.forEach(function(it){
        var url = (it && it.url) ? it.url : it;
        var meta = classify((it && it.name) ? it.name : String(url));
        loader.load(url, function(gltf){
          if (token !== loadToken) return;
          var pts = firstPoints(gltf.scene);
          if (pts){ applyStyle(pts); (meta.kind === 'rgb' ? byRgb : byTracks)[meta.idx] = pts; }
          tick();
        }, undefined, function(){ if (token === loadToken){ tick(); } });
      });
    });
  }
  if (playBtn) playBtn.addEventListener('click', toggle);
  if (tracksEl) tracksEl.addEventListener('change', function(){ mode = tracksEl.checked ? 'tracks' : 'rgb'; showFrame(cur); });
  if (speedEl) speedEl.addEventListener('input', function(){ fps = parseInt(speedEl.value, 10) || 12; if (fpsEl) fpsEl.textContent = fps + ' fps'; });
  if (scrubEl) scrubEl.addEventListener('input', function(){ pause(); showFrame(parseInt(scrubEl.value, 10) || 0); });
  window.faViewer = { load: load, play: play, pause: pause, setFrame: showFrame };
  waitThree(function(){ initThree(); });
})();
"""

# Bridge: when the hidden file list (served .glb URLs) changes, hand the URL +
# filename of each to the three.js viewer (filename selects tracks vs RGB and the
# frame index). Runs purely client-side (no server round-trip).
ANIM_BRIDGE_JS = """
(files) => {
  try {
    var list = (files || []).map(function(f){
      if (!f) return null;
      var url = f.url || f.path;
      if (!url) return null;
      var name = f.orig_name || String(url).split('/').pop().split('?')[0];
      return { url: url, name: name };
    }).filter(Boolean);
    if (window.faViewer) { window.faViewer.load(list); }
  } catch (e) { console.error('faViewer load error', e); }
}
"""


def build_demo():
    with gr.Blocks(title="Face Anything") as demo:
        gr.Markdown(DESCRIPTION)

        with gr.Row():
            # ---------------- inputs ----------------
            with gr.Column(scale=1):
                files = gr.File(
                    label=f"Input images (up to {MAX_IMAGES}, in temporal order)",
                    file_count="multiple",
                    file_types=["image"],
                    type="filepath",
                )
                gallery = gr.Gallery(
                    label="Preview", columns=6, height=180, show_label=True,
                    object_fit="contain",
                )
                video = gr.Video(
                    label=f"…or upload a video (its first {MAX_IMAGES} frames are used)",
                )
                mode = gr.Radio(
                    choices=["Joint", "One-by-one"],
                    value="One-by-one",
                    label="Inference mode",
                    info="One-by-one: more surface detail, lower memory. "
                         "Joint (all-at-once): more 3D-consistent across frames.",
                )
                face_crop = gr.Checkbox(
                    value=True, label="Face crop",
                    info="Crop each frame to a face-centred square (pixel3dmm-style) "
                         "so the model focuses on the face. Uncheck for full frames.",
                )
                remove_bg = gr.Checkbox(
                    value=True, label="Remove background",
                    info="Robust Video Matting (recommended).",
                )
                process_res = gr.Slider(
                    252, 1036, value=504, step=14,
                    label="Processing resolution",
                    info="Higher = more detail (and more memory). Multiples of 14.",
                )

                with gr.Accordion("Point-track settings", open=False):
                    n_tracks = gr.Slider(10, 500, value=100, step=10,
                                         label="Number of tracks (seeds)")
                    track_k = gr.Slider(1, 100, value=25, step=1,
                                        label="Neighbours recolored per track (k)")
                    track_threshold = gr.Slider(
                        0.001, 0.1, value=0.01, step=0.001,
                        label="Canonical match threshold")

                with gr.Accordion("Advanced", open=False):
                    conf_percentile = gr.Slider(
                        0, 95, value=0, step=5,
                        label="Confidence percentile cut",
                        info="Drop the least-confident depth pixels (0 = keep all).")
                    fps = gr.Slider(1, 30, value=10, step=1, label="Output video FPS")
                    max_frames = gr.Slider(
                        1, MAX_IMAGES, value=MAX_IMAGES, step=1,
                        label="Max frames to use")

                run_btn = gr.Button("Reconstruct", variant="primary")

            # ---------------- outputs ----------------
            with gr.Column(scale=1):
                gr.Markdown("**3D point cloud with colorful tracks**  ·  "
                            "loads the whole sequence, then plays smoothly")
                # client-side three.js player: all frames preloaded once, then
                # animated by visibility toggle (no per-frame reload → no white
                # flashes; full, un-subsampled cloud). See VIEWER_JS above.
                viewer = gr.HTML(
                    value=VIEWER_MARKUP, head=THREE_HEAD, js_on_load=VIEWER_JS,
                    show_label=False,
                )
                # hidden: run() puts the per-frame .glb files here so gradio
                # serves them; ANIM_BRIDGE_JS hands their URLs to the viewer.
                anim_files = gr.File(file_count="multiple", visible=False)
                tracks_zip = gr.File(
                    label="Download point clouds (.zip: tracks/ + points/)")
                with gr.Tab("Normals (2D)"):
                    normals_vid = gr.Video(label="Surface-normal map")
                with gr.Tab("Depth (2D)"):
                    depth_vid = gr.Video(label="Depth map")
                with gr.Tab("Canonical (2D)"):
                    canonical_vid = gr.Video(label="Canonical facial-coordinate map")
                with gr.Tab("Tracks (2D)"):
                    tracks2d_vid = gr.Video(label="2D point-track overlay")
                status = gr.Textbox(label="Log", lines=6, interactive=False)

        # preview uploaded files in the gallery
        files.change(lambda fs: fs or [], inputs=files, outputs=gallery)

        run_btn.click(
            run,
            inputs=[files, video, mode, process_res, remove_bg, face_crop,
                    conf_percentile, n_tracks, track_k, track_threshold,
                    fps, max_frames],
            outputs=[anim_files, canonical_vid, depth_vid, normals_vid,
                     tracks2d_vid, tracks_zip, status],
            concurrency_limit=1,
        )
        # when the served .glb list changes, push the URLs to the three.js player
        anim_files.change(None, inputs=anim_files, outputs=None, js=ANIM_BRIDGE_JS)

        # ---------------- examples (thumbnail shown, click to load + run) ----------------
        ex40 = sorted(glob.glob(os.path.join(APP_DIR, "examples", "seq40", "*.png")))[:MAX_IMAGES]
        if ex40:
            run_inputs = [files, video, mode, process_res, remove_bg, face_crop,
                          conf_percentile, n_tracks, track_k, track_threshold,
                          fps, max_frames]
            run_outputs = [anim_files, canonical_vid, depth_vid, normals_vid,
                           tracks2d_vid, tracks_zip, status]

            def _thumb():
                return gr.Image(value=ex40[0], height=150, show_label=False,
                                interactive=False, container=False)

            gr.Markdown("### Examples")
            with gr.Row():
                with gr.Column(scale=1, min_width=150):
                    _thumb()
                    gr.Markdown("**NeRSemble** 40 images")
                    btn40 = gr.Button("Load & run", size="sm")
                with gr.Column(scale=1, min_width=150):
                    _thumb()
                    gr.Markdown("**NeRSemble** 1 image")
                    btn1 = gr.Button("Load & run", size="sm")
                with gr.Column(scale=3):  # spacer so the thumbnails stay small
                    pass

            # set the inputs, then run the pipeline (which reads the just-set values)
            btn40.click(lambda: (ex40, None, MAX_IMAGES),
                        outputs=[files, video, max_frames]).then(
                run, inputs=run_inputs, outputs=run_outputs, concurrency_limit=1)
            btn1.click(lambda: ([ex40[0]], None, 1),
                       outputs=[files, video, max_frames]).then(
                run, inputs=run_inputs, outputs=run_outputs, concurrency_limit=1)
    return demo


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