Spaces:

hugging-apps
/

pointdit-geometry

Runtime error

tighten ZeroGPU duration to measured runtime; pass theme to launch (Gradio 6)

47ac6d9 verified about 1 month ago

15.8 kB

"""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)

	"""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)