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

PointDiT — single image → 3D point map

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.
Paper · Project page · Code · Weights — ICML 2026

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