app.py

"""PointDiT — pixel-space diffusion for monocular geometry estimation (ICML 2026).

Single image in, dense 3D point map out.  This app runs the authors' reference
implementation (https://github.com/google-research/pointdit, Apache-2.0) with the
released PointDiT-H/16 512x512 checkpoint, following `scripts/demo_h_512.sh`
one-to-one: full-frame resize to the training token budget, 3-step Euler sampling
from zeros, resize back, sky removal, flying-point removal.
"""

import os
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from types import SimpleNamespace

import spaces  # noqa: F401  -- must be imported before torch / any CUDA touching import

import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
from PIL import Image

# ----------------------------------------------------------------------------------
# Model / checkpoint selection (PointDiT-H at 512x512, stage-2 "mixdata" weights).
# ----------------------------------------------------------------------------------
POINTDIT_REPO_ID = "haofeixu/pointdit"
POINTDIT_CKPT = "pointdith-512-mixdata-nodinov3-cb01dd3b.pth"
MODEL_NAME = "PointDiT-H/16"
FEATURE_EMBEDDING = "dinov3_vith16plus"
PROJ_DROPOUT = 0.2  # value used by scripts/demo_h_512.sh (inactive at eval time)

# The frozen DINOv3 encoder is stripped from the released checkpoints because the
# official weights are gated; fetch the identical file from a public mirror.
DINOV3_MIRROR_ID = "Fanqi-Lin-IR/dinov3_vith16plus"
DINOV3_WEIGHTS_NAME = "dinov3_vith16plus_pretrain_lvd1689m-7c1da9a5.pth"

PATCH_SIZE = 16
TRAIN_IMG_SIZE = 512
DEFAULT_TOKENS = (TRAIN_IMG_SIZE // PATCH_SIZE) ** 2  # 1024 tokens the model was trained on
MAX_VIEWER_POINTS = 500_000

# ----------------------------------------------------------------------------------
# Dependencies that are fetched rather than vendored: the PointDiT reference code and
# the DINOv3 encoder code (the latter is under the DINOv3 License and is not
# redistributed here).
# ----------------------------------------------------------------------------------
DEPS = Path(os.environ.get("POINTDIT_DEPS", "/tmp/pointdit_deps"))
POINTDIT_DIR = DEPS / "pointdit"
DINOV3_DIR = DEPS / "dinov3"
DINOV3_WEIGHTS_DIR = DEPS / "dinov3_weights"


def _clone(url: str, dest: Path) -> Path:
    if (dest / ".git").exists():
        return dest
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f"[setup] cloning {url} -> {dest}", flush=True)
    subprocess.run(["git", "clone", "--depth", "1", url, str(dest)], check=True)
    return dest


_clone("https://github.com/google-research/pointdit.git", POINTDIT_DIR)
_clone("https://github.com/facebookresearch/dinov3.git", DINOV3_DIR)

DINOV3_WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
dinov3_weights_path = hf_hub_download(
    repo_id=DINOV3_MIRROR_ID,
    filename=DINOV3_WEIGHTS_NAME,
    local_dir=str(DINOV3_WEIGHTS_DIR),
)
checkpoint_path = hf_hub_download(repo_id=POINTDIT_REPO_ID, filename=POINTDIT_CKPT)

# model.py reads both of these when it builds the frozen encoder through torch.hub.
os.environ["DINOV3_REPO"] = str(DINOV3_DIR)
os.environ["DINOV3_WEIGHTS_DIR"] = str(DINOV3_WEIGHTS_DIR)
sys.path.insert(0, str(POINTDIT_DIR))

import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
import trimesh  # noqa: E402

from denoiser import Denoiser  # noqa: E402  (PointDiT reference implementation)
from util.viz_depth import viz_depth_tensor  # noqa: E402
from util.viz_pointcloud import depth_flying_points, save_single_point_cloud  # noqa: E402

# ----------------------------------------------------------------------------------
# Build the model.  Every value below is the default of the reference argument parser,
# overridden only where scripts/demo_h_512.sh overrides it.
# ----------------------------------------------------------------------------------
ARGS = SimpleNamespace(
    model=MODEL_NAME,
    img_size=TRAIN_IMG_SIZE,
    in_channels=3,
    attn_dropout=0.0,
    proj_dropout=PROJ_DROPOUT,
    attention_type="torch",
    feature_embedding_type=FEATURE_EMBEDDING,
    dinov3_use_intermediate_layers=True,
    dinov3_num_intermediate_layers=4,
    feature_embedding_lr_scale=0.0,
    P_mean=-0.8,
    P_std=0.8,
    t_eps=5e-2,
    noise_scale=1.0,
    ema_decay1=0.9999,
    ema_decay2=0.9996,
    num_sampling_steps=3,
    generate_noise_scale=0.0,  # ODE starts from zeros -> deterministic output
    sample_t_eps=0.0,
)

print(f"[setup] building {MODEL_NAME} ...", flush=True)
model = Denoiser(ARGS)

# The encoder is also loaded explicitly (as main.py does), so a hub-side change can
# never leave us conditioning on a randomly initialised DINOv3.
encoder_sd = torch.load(dinov3_weights_path, map_location="cpu")
model.net.y_embedder.load_state_dict(encoder_sd, strict=True)
del encoder_sd
print(f"[setup] loaded DINOv3 {FEATURE_EMBEDDING} encoder", flush=True)

checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
incompatible = model.load_state_dict(checkpoint["model"], strict=False)
missing = [k for k in incompatible.missing_keys if not k.startswith("net.y_embedder.")]
if missing or incompatible.unexpected_keys:
    raise RuntimeError(
        f"checkpoint mismatch: missing={missing} unexpected={incompatible.unexpected_keys}"
    )
checkpoint.pop("model")

# Evaluation uses the first EMA copy (engine.evaluate_img2point switches to it).
ema_state_dict = checkpoint.get("model_ema1")
if ema_state_dict is not None:
    with torch.no_grad():
        for name, param in model.named_parameters():
            if name in ema_state_dict:
                param.copy_(ema_state_dict[name])
    print("[setup] switched to EMA weights", flush=True)
model.ema_params1 = None
model.ema_params2 = None
del checkpoint, ema_state_dict

model.eval()
model.requires_grad_(False)
model.to("cuda")
print("[setup] model ready", flush=True)


# ----------------------------------------------------------------------------------
# Inference
# ----------------------------------------------------------------------------------
@spaces.GPU(duration=20)  # measured on ZeroGPU: ~0.15 s default, ~1 s at the max settings
def _generate_point_map(image01: np.ndarray, num_steps: int, target_tokens: int):
    """Run the PointDiT sampler on one RGB image and return its dense 3D point map.

    Args:
        image01: (H, W, 3) float32 RGB image with values in [0, 1].
        num_steps: number of Euler sampling steps.
        target_tokens: token budget the full frame is resized to before sampling.

    Returns:
        (point_map, seconds) where point_map is (H, W, 3) float32 XYZ per pixel.
    """
    image = torch.from_numpy(image01).permute(2, 0, 1)[None].to("cuda", torch.float32)
    image = image * 2.0 - 1.0  # model input range
    _, _, height, width = image.shape

    # Full-frame resize, aspect ratio preserved, to the patch-aligned resolution closest
    # to the token budget the model was trained on (engine.py, wild-image branch).
    resize_factor = ((target_tokens * PATCH_SIZE**2) / (height * width)) ** 0.5
    new_h = max(PATCH_SIZE, int(round(height * resize_factor / PATCH_SIZE)) * PATCH_SIZE)
    new_w = max(PATCH_SIZE, int(round(width * resize_factor / PATCH_SIZE)) * PATCH_SIZE)
    original_size = None
    if (new_h, new_w) != (height, width):
        original_size = (height, width)
        image = F.interpolate(image, size=(new_h, new_w), mode="bilinear", align_corners=False)

    model.steps = int(num_steps)
    torch.cuda.synchronize()
    started = time.perf_counter()
    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
        point_map = model.generate(image)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - started

    point_map = point_map.float()
    if original_size is not None:
        point_map = F.interpolate(point_map, size=original_size, mode="nearest")
    return point_map[0].permute(1, 2, 0).cpu().numpy(), elapsed


def _resize_for_output(image: np.ndarray, max_size: int) -> np.ndarray:
    height, width = image.shape[:2]
    longest = max(height, width)
    if longest <= max_size:
        return image
    scale = max_size / longest
    new_size = (max(1, int(round(width * scale))), max(1, int(round(height * scale))))
    return np.asarray(Image.fromarray(image).resize(new_size, Image.LANCZOS))


def predict(
    image: np.ndarray,
    num_steps: int = 3,
    max_size: int = 800,
    target_tokens: int = DEFAULT_TOKENS,
    remove_sky: bool = True,
    sky_threshold: float = 2.9,
    remove_flying_points: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Estimate a dense 3D point map (and depth map) from a single RGB image.

    Args:
        image: input RGB image.
        num_steps: Euler sampling steps of the pixel-space diffusion model.
        max_size: longest side of the output point map / depth map, in pixels.
        target_tokens: token budget used for the prediction (1024 = training budget).
        remove_sky: drop points the model places on the sky dome.
        sky_threshold: point-norm above which a pixel counts as sky.
        remove_flying_points: drop points floating between depth discontinuities.

    Returns:
        3D point cloud (.glb), colorized depth map, downloadable files, run info.
    """
    if image is None:
        raise gr.Error("Please provide an input image.")
    if image.ndim == 2:
        image = np.stack([image] * 3, axis=-1)
    image = np.ascontiguousarray(image[..., :3])
    image = _resize_for_output(image, int(max_size))
    image01 = image.astype(np.float32) / 255.0

    point_map, elapsed = _generate_point_map(image01, int(num_steps), int(target_tokens))

    # Sky detection runs on the raw (normalized) model output, exactly as in engine.py.
    sky_mask = None
    if remove_sky:
        sky_mask = np.linalg.norm(point_map, axis=-1) > float(sky_threshold)

    depth = point_map[..., 2]
    depth_for_viz = np.where(sky_mask, 1e6, depth) if sky_mask is not None else depth
    # The wild-image prediction is zero-centered with unknown scale: shifted_depth=True.
    depth_vis = viz_depth_tensor(depth_for_viz, shifted_depth=True, colormap="plasma")

    keep = np.ones(depth.shape, dtype=bool)
    if remove_flying_points:
        keep &= ~depth_flying_points(depth, rtol=0.04)
    if sky_mask is not None:
        keep &= ~sky_mask
    if not keep.any():  # pathological input: keep everything rather than emit nothing
        keep = np.ones(depth.shape, dtype=bool)

    out_dir = Path(tempfile.mkdtemp(prefix="pointdit_"))
    ply_path = out_dir / "pointcloud.ply"
    glb_path = out_dir / "pointcloud.glb"
    depth_path = out_dir / "depth.png"
    npz_path = out_dir / "pointmap.npz"

    save_single_point_cloud(point_map, image01, str(ply_path), transform_to_gl=True, mask=keep)
    Image.fromarray(depth_vis).save(depth_path)
    np.savez_compressed(
        npz_path, points=point_map.astype(np.float32), mask=keep, image=image
    )

    # Viewer copy: same OpenCV -> OpenGL flip as the .ply, subsampled so the browser
    # does not have to swallow multi-million-point scenes.
    points = point_map.reshape(-1, 3)[keep.reshape(-1)]
    colors = (image01.reshape(-1, 3)[keep.reshape(-1)] * 255).astype(np.uint8)
    points = points * np.array([1.0, -1.0, -1.0], dtype=np.float32)
    if len(points) > MAX_VIEWER_POINTS:
        idx = np.linspace(0, len(points) - 1, MAX_VIEWER_POINTS).astype(np.int64)
        points, colors = points[idx], colors[idx]
    trimesh.PointCloud(points, colors=colors).export(glb_path)

    info = (
        f"**{MODEL_NAME}** · {num_steps} sampling step(s) · "
        f"{image.shape[1]}×{image.shape[0]} output · {len(points):,} points · "
        f"sampling took **{elapsed:.2f}s**"
    )
    files = [str(glb_path), str(ply_path), str(depth_path), str(npz_path)]
    return str(glb_path), depth_vis, files, info


# ----------------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------------
DESCRIPTION = """
<h1 align="center">PointDiT — single image → 3D point map</h1>
<p align="center">
A plain ViT that denoises raw 3D point-map patches directly in pixel space, conditioned on
frozen DINOv3 image tokens — no point tokenizer, no latent diffusion, no camera intrinsics.
<br>
<a href="https://huggingface.co/papers/2607.02515">Paper</a> ·
<a href="https://haofeixu.github.io/pointdit/">Project page</a> ·
<a href="https://github.com/google-research/pointdit">Code</a> ·
<a href="https://huggingface.co/haofeixu/pointdit">Weights</a> — ICML 2026
</p>
"""

with gr.Blocks(title="PointDiT") as demo:
    gr.HTML(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="numpy", image_mode="RGB", label="Input image")
            run_button = gr.Button("Estimate geometry", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                num_steps = gr.Slider(
                    1, 8, value=3, step=1, label="Sampling steps",
                    info="More steps sharpen fine detail (the paper's demo uses 3).",
                )
                max_size = gr.Slider(
                    384, 1280, value=800, step=64, label="Output resolution (longest side)",
                    info="The prediction is upsampled back to this size.",
                )
                target_tokens = gr.Slider(
                    256, 4096, value=DEFAULT_TOKENS, step=256, label="Token budget",
                    info="1024 matches the 512×512 training resolution.",
                )
                remove_sky = gr.Checkbox(value=True, label="Remove sky points")
                sky_threshold = gr.Slider(
                    1.0, 5.0, value=2.9, step=0.05, label="Sky threshold (point norm)"
                )
                remove_flying_points = gr.Checkbox(
                    value=True, label="Remove flying points at depth edges"
                )

        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("3D point map"):
                    model_3d = gr.Model3D(
                        label="Predicted point cloud",
                        display_mode="point_cloud",
                        clear_color=[0.0, 0.0, 0.0, 1.0],
                        height="55vh",
                    )
                    run_info = gr.Markdown("")
                with gr.Tab("Depth"):
                    depth_image = gr.Image(
                        type="numpy", label="Inverse depth (plasma)", format="png",
                        interactive=False,
                    )
                with gr.Tab("Downloads"):
                    output_files = gr.File(
                        type="filepath",
                        file_count="multiple",
                        label="Point cloud, depth map and raw point map",
                    )

    gr.Examples(
        examples=[
            ["examples/01_facade.png"],
            ["examples/02_office.jpg"],
            ["examples/03_traffic.jpg"],
            ["examples/04_mountain.jpg"],
        ],
        inputs=[input_image],
        outputs=[model_3d, depth_image, output_files, run_info],
        fn=predict,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples",
    )

    run_button.click(
        fn=predict,
        inputs=[
            input_image,
            num_steps,
            max_size,
            target_tokens,
            remove_sky,
            sky_threshold,
            remove_flying_points,
        ],
        outputs=[model_3d, depth_image, output_files, run_info],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), mcp_server=True)