# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. import argparse import os import random import tempfile from pathlib import Path from typing import Any, Optional import gradio as gr import numpy as np import plotly.graph_objects as go import spaces import torch import trimesh from huggingface_hub import hf_hub_download from meshflow.pipelines import MeshFlowPipeline from meshflow.utils.dtype import AUTOCAST_DTYPE_CHOICES from meshflow.utils.mesh import ( _read_point_cloud_file, DEFAULT_NUM_VERTS, GEOMETRY_EXTS, Mesh, resolve_num_verts_for_mesh, ) from omegaconf import OmegaConf from PIL import Image REPO_ROOT = Path(__file__).resolve().parent CHECKPOINT_REPO_ID = "facebook/meshflow" CHECKPOINT_BUNDLE_DEFAULT = "meshflow" CHECKPOINT_BUNDLE_NUM_VERTS = "meshflow_w_num_verts_control" CHECKPOINT_BUNDLES = (CHECKPOINT_BUNDLE_DEFAULT, CHECKPOINT_BUNDLE_NUM_VERTS) DEFAULT_CHECKPOINT_BUNDLE = CHECKPOINT_BUNDLE_DEFAULT CHECKPOINT_CONFIG_FILENAME = "config.yaml" CHECKPOINT_WEIGHTS_FILENAME = "model.pth" GALLERY_DIR = REPO_ROOT / "assets" / "gallery" GALLERY_SURFACE_PC_DIR = GALLERY_DIR / "surface_pc" GALLERY_THUMBNAIL_DIR = GALLERY_DIR / "thumbnails" GALLERY_EXAMPLES_PER_PAGE = 8 GALLERY_THUMBNAIL_SIZE = 168 GALLERY_THUMBNAIL_PADDING = 6 GALLERY_THUMBNAIL_MAX_POINTS = 4096 GALLERY_THUMBNAIL_POINT_COLOR = (93, 164, 189) GALLERY_THUMBNAIL_CACHE_VERSION = 2 GALLERY_THUMBNAIL_VERSION_FILE = GALLERY_THUMBNAIL_DIR / ".cache_version" NUM_VERTS_MIN = 1024 NUM_VERTS_MAX = 4096 NUM_VERTS_STEP = 256 NUM_VERTS_CONTROL_NOTE = ( "`num_verts` is injected via `proj_cond_on_temb` as `num_verts / num_latents` " "(normalization uses `mesh_model.num_latents` from config, e.g. 4096). " "For `.glb` with fewer verts than `num_latents`, the file vertex count is used. " f"({NUM_VERTS_MIN}–{NUM_VERTS_MAX}, default {DEFAULT_NUM_VERTS})." ) PREVIEW_ROT_X_DEG = 90.0 PREVIEW_ROT = trimesh.transformations.rotation_matrix( np.deg2rad(PREVIEW_ROT_X_DEG), [1.0, 0.0, 0.0] ) PLOT_SCENE_BG = "#eef4f8" PLOT_MESH_COLOR = "#5da4bd" PLOT_MESH_OPACITY = 0.62 PLOT_WIRE_COLOR = "#1a3344" PLOT_WIRE_HALO_COLOR = "rgba(238, 244, 248, 0.9)" PLOT_WIRE_WIDTH = 1.8 PLOT_WIRE_HALO_WIDTH = 3.2 PLOT_AXIS_RANGE = 0.92 PLOT_MESH_AXIS_PADDING = 1.08 PLOT_WIRE_AXIS_PADDING = 1.12 PREVIEW_POINT_CLOUD_MAX_POINTS = 8192 INPUT_PREVIEW_PLOT_HEIGHT = 240 OUTPUT_PLOT_HEIGHT = 340 INPUT_PREVIEW_AXIS_RANGE = 0.78 INPUT_PREVIEW_AXIS_PADDING = 0.96 INPUT_PREVIEW_POINT_SIZE = 1.5 INPUT_PREVIEW_POINT_COLOR = "#4a5568" INPUT_PREVIEW_CAMERA = dict( eye=dict(x=0.0, y=-1.38, z=0.0), center=dict(x=0.0, y=0.0, z=0.0), up=dict(x=0.0, y=0.0, z=1.0), ) PLOT_CAMERA = dict( eye=dict(x=0.0, y=-1.75, z=0.0), center=dict(x=0.0, y=0.0, z=0.0), up=dict(x=0.0, y=0.0, z=1.0), ) APP_TITLE = "MeshFlow" APP_VENUE = "CVPR 2026 Highlight" APP_TAGLINE = ( "Generate artist-like meshes from surface point clouds in about one second." ) APP_TAB_TITLE = "MeshFlow Demo" PROJECT_PAGE_URL = "https://mesh-flow.github.io/" ARXIV_URL = "https://arxiv.org/pdf/2606.04621" GITHUB_URL = "https://github.com/facebookresearch/meshflow" HF_MODEL_URL = "https://huggingface.co/facebook/meshflow" PAPER_AUTHORS = ( ("https://weiyuli.xyz/", "Weiyu Li"), ("https://www.antoinetlc.com/", "Antoine Toisoul"), ("https://tmonnier.com/", "Tom Monnier"), ("https://shapovalov.ro/", "Roman Shapovalov"), ("https://www.linkedin.com/in/rakesh-r-3848538", "Rakesh Ranjan"), ("https://ece.hkust.edu.hk/pingtan", "Ping Tan"), ("https://www.robots.ox.ac.uk/~vedaldi/", "Andrea Vedaldi"), ) HOW_TO_STEPS = ( ("Upload", "Upload a point cloud or mesh, or choose an Example."), ("Generate", "Click Generate."), ("Download", "Preview the mesh and download the GLB file."), ) _BRAND_LETTERS = "".join( f'{ch}' for i, ch in enumerate(APP_TITLE) ) _HOWTO_LI = "".join(f"

{a} — {b}

" for a, b in HOW_TO_STEPS) _AUTHORS = "".join( f'{n}' for u, n in PAPER_AUTHORS ) _HEADER_LINKS = " ".join( f'{t}' for u, t in ( (PROJECT_PAGE_URL, "Project"), (GITHUB_URL, "GitHub"), (HF_MODEL_URL, "Model"), (ARXIV_URL, "Paper"), ) ) APP_HEADER_HTML = f""" """ _THEME_COLORS = dict( body_background_fill="#eef4f8", block_background_fill="#ffffff", block_border_color="#d4e0ea", body_text_color="#162432", input_background_fill="#f7fafc", input_border_color="#d4e0ea", background_fill_primary="#f4f8fb", background_fill_secondary="#ffffff", border_color_primary="#d4e0ea", block_label_background_fill="#edf6fa", block_label_text_color="#0f6d8f", block_title_background_fill="#edf6fa", block_title_text_color="#0f6d8f", button_secondary_background_fill="#f4f8fb", button_secondary_text_color="#243447", button_secondary_border_color="#d4e0ea", table_even_background_fill="#f7fafc", table_odd_background_fill="#ffffff", panel_background_fill="#ffffff", checkbox_background_color="#ffffff", checkbox_border_color="#d4e0ea", stat_background_fill="#f4f8fb", ) _THEME_KW: dict[str, str] = {} for _k, _v in _THEME_COLORS.items(): _THEME_KW[_k] = _v _THEME_KW[f"{_k}_dark"] = _v MESHFLOW_THEME = gr.themes.Soft( primary_hue=gr.themes.colors.cyan, secondary_hue=gr.themes.colors.blue, neutral_hue=gr.themes.colors.slate, font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"], font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"], ).set( **_THEME_KW, block_border_width="1px", block_label_text_size="*text_sm", block_label_text_weight="600", block_title_text_weight="600", block_shadow="none", block_shadow_dark="none", button_large_padding="12px 20px", button_primary_background_fill="#1484a8", button_primary_background_fill_hover="#0f6d8f", button_primary_text_color="#ffffff", slider_color="#1484a8", ) FORCE_LIGHT_MODE_HEAD = ( "" ) GALLERY_PAGINATION_SCROLL_JS = ( "" ) GALLERY_REVEAL_JS = ( "" ) APP_HEAD = FORCE_LIGHT_MODE_HEAD + GALLERY_PAGINATION_SCROLL_JS + GALLERY_REVEAL_JS CUSTOM_CSS = ( "@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700" "&family=Ubuntu:wght@500;600;700&display=swap');" "html,body,.gradio-container,.main,.contain,.app{scrollbar-gutter:stable;}" ":root{--mf-ink:#243447;--mf-muted:#5c7284;--mf-line:#d4e0ea;--mf-accent:#1484a8;" "--mf-accent-dark:#0f6d8f;--mf-accent-soft:#e8f4f9;--mf-panel-muted:#f4f8fb;" "--mf-shadow:0 1px 2px rgba(22,36,50,.04),0 6px 18px rgba(22,36,50,.035);color-scheme:light;}" ":root,:root.dark,.dark{--body-background-fill:#eef4f8!important;--block-background-fill:#fff!important;" "--body-text-color:#243447!important;--border-color-primary:#d4e0ea!important;color-scheme:light!important;}" ".gradio-container{max-width:1080px!important;margin:0 auto!important;padding:1.25rem 1rem 2rem!important;" "background:linear-gradient(180deg,#eef4f8,#f4f8fb)!important;color:var(--mf-ink)!important;" "font-family:Inter,ui-sans-serif,system-ui,sans-serif!important;}" ".mf-workspace>.gap,.mf-workspace>.form,.mf-workspace>.wrap{display:flex!important;" "flex-direction:column!important;gap:16px!important;width:100%;}" ".mf-workspace,.mf-gallery,.mf-gallery>.block,.mf-gallery>.form,.mf-gallery .wrap" "{display:block!important;width:100%!important;max-width:100%!important;overflow-anchor:none!important;}" ".mf-gallery{container-type:inline-size!important;opacity:0!important;transition:opacity .18s ease!important;}" ".mf-gallery.mf-gallery-ready{opacity:1!important;}" ".mf-gallery .gallery{display:flex!important;flex-wrap:wrap!important;gap:8px!important;" "width:100%!important;justify-content:center!important;align-content:flex-start!important;" "min-height:calc((100cqw - " + str((GALLERY_EXAMPLES_PER_PAGE - 1) * 8) + "px) / " + str(GALLERY_EXAMPLES_PER_PAGE) + ")!important;overflow-anchor:none!important;}" ".mf-gallery .paginate{margin-top:8px!important;overflow-anchor:none!important;}" ".mf-gallery .gallery-item{flex:0 0 calc((100% - " + str(GALLERY_EXAMPLES_PER_PAGE - 1) + " * 8px) / " + str(GALLERY_EXAMPLES_PER_PAGE) + ")!important;width:calc((100% - " + str(GALLERY_EXAMPLES_PER_PAGE - 1) + " * 8px) / " + str(GALLERY_EXAMPLES_PER_PAGE) + ")!important;max-width:calc((100% - " + str(GALLERY_EXAMPLES_PER_PAGE - 1) + " * 8px) / " + str(GALLERY_EXAMPLES_PER_PAGE) + ")!important;min-width:0!important;aspect-ratio:1!important;" "padding:0!important;overflow:hidden!important;background:#fff!important;" "border:1px solid transparent!important;box-sizing:border-box!important;" "display:flex!important;align-items:center!important;justify-content:center!important;}" ".mf-gallery .gallery-item:hover,.mf-gallery .gallery-item.selected" "{border-color:var(--mf-accent)!important;background:#fff!important;}" ".mf-gallery .gallery-item>*,.mf-gallery .gallery-item button,.mf-gallery .gallery-item .contain" "{width:100%!important;height:100%!important;min-height:0!important;padding:0!important;margin:0!important;" "display:flex!important;align-items:center!important;justify-content:center!important;}" ".mf-gallery .gallery-item img{box-sizing:border-box!important;width:100%!important;height:100%!important;" "max-width:100%!important;max-height:100%!important;object-fit:contain!important;" "object-position:center center!important;display:block!important;margin:0 auto!important;}" ".mf-viewer-grid{display:flex!important;flex-wrap:nowrap!important;gap:10px!important;width:100%!important;}" ".mf-viewer-grid>.block,.mf-viewer-grid>.form,.mf-viewer-grid>.column{flex:1 1 0!important;min-width:0!important;}" ".mf-output-col>.block,.mf-output-col>.form,.mf-output-col>.column{width:100%!important;max-width:100%!important;}" ".mf-output-col .mf-full-width-action,.mf-output-col .mf-full-width-action>.wrap," ".mf-output-col .mf-full-width-action>.form,.mf-output-col .mf-full-width-action button" "{width:100%!important;max-width:100%!important;box-sizing:border-box!important;}" ".mf-output-col .mf-generate-action{margin:0 0 .65rem!important;}" ".mf-output-col .mf-download-action{margin:.65rem 0 0!important;}" ".mf-viewer-grid .plot-container,.mf-viewer-grid .gr-panel{height:" + str(OUTPUT_PLOT_HEIGHT) + "px!important;min-height:" + str(OUTPUT_PLOT_HEIGHT) + "px!important;max-height:" + str(OUTPUT_PLOT_HEIGHT) + "px!important;}" ".mf-input-preview-plot,.mf-input-preview-plot>.block,.mf-input-preview-plot .wrap{min-height:" + str(INPUT_PREVIEW_PLOT_HEIGHT) + "px!important;}" ".mf-input-preview-plot .plot-container,.mf-input-preview-plot .gr-panel{height:" + str(INPUT_PREVIEW_PLOT_HEIGHT) + "px!important;min-height:" + str(INPUT_PREVIEW_PLOT_HEIGHT) + "px!important;max-height:" + str(INPUT_PREVIEW_PLOT_HEIGHT) + "px!important;}" ".mf-input-file,.mf-input-file>.block,.mf-input-file .wrap,.mf-input-file .file-preview" "{min-height:168px!important;max-height:168px!important;overflow:hidden!important;}" ".gradio-container .js-plotly-plot .modebar-container,.gradio-container .js-plotly-plot .modebar" "{top:auto!important;bottom:6px!important;right:6px!important;left:auto!important;}" ".gradio-container .js-plotly-plot .plotly-logomark" "{top:auto!important;bottom:6px!important;right:6px!important;left:auto!important;}" ".mf-app-header{display:flex;flex-direction:column;gap:.85rem;margin-bottom:1rem;padding:0;" "background:transparent;border:none;box-shadow:none;}" ".mf-header-wrap.block,.mf-header-wrap>.block,.mf-header-wrap .html-container{padding:0!important;" "margin:0!important;background:transparent!important;box-shadow:none!important;border:none!important;}" ".mf-header-wrap .mf-app-header{margin-bottom:0!important;}" ".mf-header-wrap p,.mf-app-header p{padding:0!important;margin-left:0!important;margin-right:0!important;" "text-indent:0!important;}" ".mf-app-header-top{display:flex;flex-wrap:wrap;justify-content:space-between;gap:1rem 1.5rem;}" ".mf-app-title-wrap{flex:1 1 280px;min-width:0;display:flex!important;flex-direction:column!important;" "align-items:flex-start!important;}" ".mf-app-title{margin:0;font-family:Ubuntu,Helvetica,sans-serif;font-size:clamp(1.65rem,3vw,1.9rem);" "font-weight:650;line-height:1.15;display:flex;flex-wrap:wrap;align-items:baseline;gap:.35rem .55rem;}" ".mf-venue-badge{color:var(--mf-muted);font-size:.58em;font-weight:550;white-space:nowrap;}" ".mf-brand-word{display:inline-flex;font-weight:700;white-space:nowrap;}" ".mf-brand-letter{display:inline-block;background:linear-gradient(100deg,#042f4b,#075985 14%,#0284c7 28%," "#06b6d4 42%,#67e8f9 52%,#f0fdff 58%,#38bdf8 66%,#0e7490 82%,#083344);background-size:260% 100%;" "background-position:0% 50%;background-attachment:fixed;-webkit-background-clip:text;background-clip:text;" "color:transparent;-webkit-text-fill-color:transparent;" "animation:mf-brand-flow 5.5s ease-in-out infinite alternate,mf-letter-wave 2.35s ease-in-out infinite;" "animation-delay:0s,calc(var(--m-i,0)*.065s);}" "@keyframes mf-brand-flow{0%{background-position:0% 50%}100%{background-position:100% 50%}}" "@keyframes mf-letter-wave{0%,100%{transform:translateY(0)}50%{transform:translateY(-.12em)}}" "@media (prefers-reduced-motion:reduce){.mf-brand-letter{animation:none;color:#0369a1;-webkit-text-fill-color:unset;background:none;}}" ".mf-app-tagline,.mf-authors{max-width:36rem;line-height:1.5;width:100%;}" ".mf-app-tagline{margin:.35rem 0 0;color:var(--mf-muted);font-size:.92rem;}" ".mf-authors{margin:.4rem 0 0;color:#4a6274;font-size:.84rem;}" ".mf-authors a{display:inline!important;margin:0!important;padding:0!important;" "color:var(--mf-accent-dark);font-weight:600;text-decoration:none;white-space:nowrap;}" ".mf-authors a:not(:last-child)::after{content:', ';color:#4a6274;font-weight:400;}" ".mf-header-wrap .mf-app-links{display:flex!important;flex-wrap:wrap!important;gap:.45rem!important;" "align-items:center!important;align-self:flex-start!important;flex-shrink:0!important;}" ".mf-header-wrap .mf-app-links a,.mf-header-wrap a.mf-link{box-sizing:border-box!important;" "display:inline-flex!important;align-items:center!important;justify-content:center!important;" "width:auto!important;min-width:0!important;max-width:none!important;height:auto!important;" "padding:.38rem .72rem!important;margin:0!important;border:1px solid var(--mf-line)!important;" "border-radius:8px!important;background:var(--mf-panel-muted)!important;" "color:var(--mf-ink)!important;font-size:.82rem!important;font-weight:600!important;" "font-family:inherit!important;line-height:1.2!important;text-decoration:none!important;" "white-space:nowrap!important;box-shadow:none!important;cursor:pointer!important;" "appearance:none!important;-webkit-appearance:none!important;}" ".mf-header-wrap .mf-app-links a:hover,.mf-header-wrap a.mf-link:hover{" "border-color:#a8d4e8!important;background:var(--mf-accent-soft)!important;" "color:var(--mf-accent-dark)!important;text-decoration:none!important;}" ".mf-howto-details{margin-top:.55rem;width:100%;}" ".mf-howto-details>summary{cursor:pointer;list-style:none;color:var(--mf-accent-dark);font-size:.84rem;font-weight:600;}" ".mf-howto-details>summary::-webkit-details-marker{display:none;}" ".mf-howto-details>summary::before{content:'▸';display:inline-block;width:.95em;margin-right:.2rem;}" ".mf-howto-details[open]>summary::before{transform:rotate(90deg);}" ".mf-howto-steps{margin:.55rem 0 0;padding-left:1.15rem;color:#4a6274;font-size:.84rem;line-height:1.55;}" "@media (max-width:768px){.mf-app-header-top{flex-direction:column}" ".mf-header-wrap .mf-app-links{width:100%!important;}}" ) def plotly_scene_layout( fig: go.Figure, axis_range: float | None = None, camera: dict | None = None, uirevision: str | None = None, ) -> go.Figure: scene_range = PLOT_AXIS_RANGE if axis_range is None else axis_range axis = dict(visible=False, showbackground=False, range=[-scene_range, scene_range]) fig.update_layout( margin=dict(l=0, r=0, b=0, t=0), showlegend=False, uirevision=uirevision, scene=dict( xaxis=axis, yaxis=axis, zaxis=axis, bgcolor=PLOT_SCENE_BG, aspectmode="cube", camera=camera or PLOT_CAMERA, ), ) return fig def mesh_to_plotly_solid( mesh: trimesh.Trimesh, axis_range: float | None = None ) -> go.Figure: verts = np.asarray(mesh.vertices, dtype=np.float32) faces = np.asarray(mesh.faces, dtype=np.int32) mesh_trace = go.Mesh3d( x=verts[:, 0], y=verts[:, 1], z=verts[:, 2], i=faces[:, 0], j=faces[:, 1], k=faces[:, 2], color=PLOT_MESH_COLOR, opacity=PLOT_MESH_OPACITY, flatshading=True, lighting=dict( ambient=0.72, diffuse=0.45, specular=0.08, roughness=0.85, fresnel=0.05 ), lightposition=dict(x=0.35, y=-0.6, z=1.8), showscale=False, ) return plotly_scene_layout( go.Figure(data=[mesh_trace]), axis_range=axis_range, uirevision="mesh-output" ) def mesh_to_plotly_wireframe( mesh: trimesh.Trimesh, axis_range: float | None = None ) -> go.Figure: verts = np.asarray(mesh.vertices, dtype=np.float32) faces = np.asarray(mesh.faces, dtype=np.int32) edges = np.unique( np.sort( np.vstack([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]]), axis=1 ), axis=0, ) wx, wy, wz = [], [], [] for i0, i1 in edges: p0, p1 = verts[i0], verts[i1] wx.extend((float(p0[0]), float(p1[0]), None)) wy.extend((float(p0[1]), float(p1[1]), None)) wz.extend((float(p0[2]), float(p1[2]), None)) common = dict(x=wx, y=wy, z=wz, mode="lines", hoverinfo="skip") traces = [ go.Scatter3d( **common, line=dict(color=PLOT_WIRE_HALO_COLOR, width=PLOT_WIRE_HALO_WIDTH) ), go.Scatter3d(**common, line=dict(color=PLOT_WIRE_COLOR, width=PLOT_WIRE_WIDTH)), ] return plotly_scene_layout( go.Figure(data=traces), axis_range=axis_range, uirevision="mesh-output", ) def resolve_geometry_path(upload: Any) -> Optional[str]: if upload is None: return None if isinstance(upload, str): return upload or None if isinstance(upload, (list, tuple)): return resolve_geometry_path(upload[0]) if upload else None if isinstance(upload, dict): return upload.get("path") or upload.get("name") path = getattr(upload, "path", None) if path: return str(path) name = getattr(upload, "name", None) return str(name) if name else None def preview_input(input_file: Any) -> tuple[dict, Optional[str]]: path = resolve_geometry_path(input_file) if path is None: return gr.update(value=None), None try: if Mesh.is_point_cloud_file(path): points = np.asarray(_read_point_cloud_file(path).numpy(), dtype=np.float32) if points.shape[0] == 0: raise ValueError("Point cloud is empty.") if points.shape[0] > PREVIEW_POINT_CLOUD_MAX_POINTS: idx = np.random.default_rng(0).choice( points.shape[0], PREVIEW_POINT_CLOUD_MAX_POINTS, replace=False ) points = points[idx] points = trimesh.transformations.transform_points( np.asarray(points, dtype=np.float64), PREVIEW_ROT ).astype(np.float32) centered = points - points.mean(axis=0) points = centered / max(float(np.linalg.norm(centered, axis=1).max()), 1e-6) preview_axis_range = max( INPUT_PREVIEW_AXIS_RANGE, float(np.max(np.abs(points))) * INPUT_PREVIEW_AXIS_PADDING, ) figure = plotly_scene_layout( go.Figure( data=[ go.Scatter3d( x=points[:, 0], y=points[:, 1], z=points[:, 2], mode="markers", marker=dict( size=INPUT_PREVIEW_POINT_SIZE, color=INPUT_PREVIEW_POINT_COLOR, opacity=0.88, ), hoverinfo="skip", ) ] ), axis_range=preview_axis_range, camera=INPUT_PREVIEW_CAMERA, uirevision="input-preview", ) else: mesh = Mesh.load_mesh(path, normalize=False, preprocess=True).to_trimesh() mesh = mesh.copy() mesh.apply_transform(PREVIEW_ROT) verts = np.asarray(mesh.vertices, dtype=np.float32) if verts.shape[0] == 0: raise ValueError("Mesh has no vertices.") centered = verts - verts.mean(axis=0) mesh.vertices = centered / max( float(np.linalg.norm(centered, axis=1).max()), 1e-6 ) preview_axis_range = max( INPUT_PREVIEW_AXIS_RANGE, float(np.max(np.abs(mesh.vertices))) * INPUT_PREVIEW_AXIS_PADDING, ) figure = mesh_to_plotly_solid(mesh, axis_range=preview_axis_range) figure.update_layout( scene_camera=INPUT_PREVIEW_CAMERA, uirevision="input-preview" ) except Exception as exc: raise gr.Error(f"Failed to load input preview: {exc}") from exc return gr.update(value=figure), path def pick_gallery_example( index: int | None, gallery_paths: list[str] ) -> tuple[dict, Optional[str], dict]: if index is None or index < 0 or index >= len(gallery_paths): return gr.update(value=None), None, gr.update(value=None) path = gallery_paths[index] preview_update, _ = preview_input(path) return gr.update(value=path), path, preview_update def clear_input_preview() -> tuple[dict, None]: return gr.update(value=None), None def discover_gallery_examples() -> tuple[list[str], list[str]]: if not GALLERY_DIR.is_dir() or not GALLERY_SURFACE_PC_DIR.is_dir(): return [], [] rebuild_all = True if GALLERY_THUMBNAIL_VERSION_FILE.is_file(): try: rebuild_all = ( int(GALLERY_THUMBNAIL_VERSION_FILE.read_text(encoding="utf-8").strip()) != GALLERY_THUMBNAIL_CACHE_VERSION ) except ValueError: rebuild_all = True paths, thumbs = [], [] for ply_path in sorted(GALLERY_SURFACE_PC_DIR.glob("*.ply")): thumb = GALLERY_THUMBNAIL_DIR / f"{ply_path.stem}.png" needs_refresh = ( rebuild_all or not thumb.exists() or thumb.stat().st_mtime < ply_path.stat().st_mtime ) if not needs_refresh: try: with Image.open(thumb) as im: needs_refresh = im.mode != "RGBA" except OSError: needs_refresh = True if needs_refresh: try: points = np.asarray( _read_point_cloud_file(str(ply_path)).numpy(), dtype=np.float32 ) except ValueError: if not thumb.exists(): continue else: if points.shape[0] > GALLERY_THUMBNAIL_MAX_POINTS: idx = np.random.default_rng(0).choice( points.shape[0], GALLERY_THUMBNAIL_MAX_POINTS, replace=False ) points = points[idx] size, padding = GALLERY_THUMBNAIL_SIZE, GALLERY_THUMBNAIL_PADDING span = size - 2 * padding points = trimesh.transformations.transform_points( np.asarray(points, dtype=np.float64), PREVIEW_ROT ).astype(np.float64) points = points - points.mean(axis=0) xs, zs, depth = points[:, 0], points[:, 2], points[:, 1] cx = 0.5 * (float(xs.min()) + float(xs.max())) cz = 0.5 * (float(zs.min()) + float(zs.max())) half = max(float(xs.max() - xs.min()), float(zs.max() - zs.min())) * 0.5 half = max(half, 1e-6) x_norm = (xs - cx) / half z_norm = (zs - cz) / half px = ((x_norm + 1.0) * 0.5 * span + padding).astype(np.int32) py = ((1.0 - (z_norm + 1.0) * 0.5) * span + padding).astype(np.int32) shade = 0.55 + 0.45 * (depth - depth.min()) / (np.ptp(depth) + 1e-6) canvas = np.zeros((size, size, 4), dtype=np.uint8) valid = (px >= 0) & (px < size) & (py >= 0) & (py < size) px, py, shade = px[valid], py[valid], shade[valid] order = np.argsort(shade) rgb = np.clip( np.array(GALLERY_THUMBNAIL_POINT_COLOR) * shade[order, None], 0, 255 ).astype(np.uint8) canvas[py[order], px[order]] = np.column_stack( [rgb, np.full(len(order), 255, dtype=np.uint8)] ) thumb.parent.mkdir(parents=True, exist_ok=True) Image.fromarray(canvas, mode="RGBA").save(thumb) if not thumb.exists(): continue paths.append(str(ply_path.resolve())) thumbs.append(str(thumb.resolve())) if rebuild_all and thumbs: GALLERY_THUMBNAIL_DIR.mkdir(parents=True, exist_ok=True) GALLERY_THUMBNAIL_VERSION_FILE.write_text( str(GALLERY_THUMBNAIL_CACHE_VERSION), encoding="utf-8" ) return paths, thumbs def resolve_model_path( model_path: Optional[str], checkpoint_bundle: str = DEFAULT_CHECKPOINT_BUNDLE, ) -> str: if checkpoint_bundle not in CHECKPOINT_BUNDLES: raise ValueError( f"checkpoint_bundle must be one of {CHECKPOINT_BUNDLES}, got {checkpoint_bundle!r}" ) def bundle_ok(root: Path) -> bool: return (root / CHECKPOINT_CONFIG_FILENAME).is_file() and ( root / CHECKPOINT_WEIGHTS_FILENAME ).is_file() if model_path: root = Path(model_path) if bundle_ok(root): return str(root.resolve()) raise FileNotFoundError( f"model_path must contain {CHECKPOINT_CONFIG_FILENAME} and " f"{CHECKPOINT_WEIGHTS_FILENAME}: {root}" ) local_default = REPO_ROOT / "ckpt" / checkpoint_bundle if bundle_ok(local_default): print(f"[MeshFlow] Using local checkpoint bundle at {local_default.resolve()}") return str(local_default.resolve()) cache_root = Path( os.environ.get("MESHFLOW_CACHE_DIR", Path.home() / ".cache" / "meshflow") ) bundle_dir = cache_root / checkpoint_bundle bundle_dir.mkdir(parents=True, exist_ok=True) for filename in (CHECKPOINT_CONFIG_FILENAME, CHECKPOINT_WEIGHTS_FILENAME): if not (bundle_dir / filename).is_file(): hf_path = f"{checkpoint_bundle}/{filename}" downloaded = hf_hub_download( repo_id=CHECKPOINT_REPO_ID, filename=hf_path, local_dir=str(cache_root), ) print( f"[MeshFlow] Downloaded {hf_path} from {CHECKPOINT_REPO_ID} to {downloaded}" ) if not bundle_ok(bundle_dir): raise FileNotFoundError( f"Failed to prepare checkpoint bundle at {bundle_dir} from " f"{CHECKPOINT_REPO_ID}/{checkpoint_bundle}/" ) print( f"[MeshFlow] Using checkpoint bundle at {bundle_dir.resolve()} " f"(from Hugging Face {CHECKPOINT_REPO_ID}/{checkpoint_bundle}/)" ) return str(bundle_dir.resolve()) def clamp_num_verts(value: int) -> int: value = int(value) if value < NUM_VERTS_MIN or value > NUM_VERTS_MAX: raise ValueError( f"num_verts must be between {NUM_VERTS_MIN} and {NUM_VERTS_MAX}, got {value}" ) remainder = (value - NUM_VERTS_MIN) % NUM_VERTS_STEP if remainder: value -= remainder return value # Module-level handle to the loaded pipeline. _PIPELINE: Optional[MeshFlowPipeline] = None @spaces.GPU() @torch.no_grad() def run_meshflow( loaded_num_verts: Optional[int], runtime_args: argparse.Namespace, geometry_path_state: Optional[str], input_file: Any, input_image: Optional[Any], steps: int, guidance_scale: float, seed: int, num_verts: Optional[int] = None, ) -> tuple[go.Figure, go.Figure, str, Optional[int]]: geometry_path = geometry_path_state or resolve_geometry_path(input_file) if not geometry_path: raise gr.Error( "Please upload a mesh (.glb/.obj/.stl/.ply) or point cloud (.ply/.pcd/.xyz/.npz)." ) ext = Path(geometry_path).suffix.lower() if ext not in GEOMETRY_EXTS: raise gr.Error( f"Unsupported format: {ext}. Supported: {', '.join(sorted(GEOMETRY_EXTS))}" ) supports_num_verts_scaling = getattr( runtime_args, "supports_num_verts_scaling", False ) global _PIPELINE pipeline = _PIPELINE # Move the pipeline on the allocated GPU. device = f"cuda:{runtime_args.gpu}" if torch.cuda.is_available() else "cpu" pipeline.to(device) proj_num_verts = None if supports_num_verts_scaling: if num_verts is None: raise ValueError( "num_verts is required when use_proj_cond_on_temb is enabled" ) num_verts = clamp_num_verts(num_verts) if loaded_num_verts != num_verts: pipeline = MeshFlowPipeline.from_pretrained( runtime_args.model_path, device=device, dtype=runtime_args.dtype, compile_models=runtime_args.compile, num_verts=num_verts, ) _PIPELINE = pipeline loaded_num_verts = num_verts proj_num_verts = resolve_num_verts_for_mesh( Path(geometry_path), num_verts, pipeline.num_latents ) out_mesh = pipeline.run( geometry_path, image=input_image, steps=int(steps), guidance_scale=float(guidance_scale), seed=int(seed), preprocess_image=False, disable_prog=False, num_verts=proj_num_verts, ) mesh = out_mesh.to_trimesh().copy() mesh.apply_transform(PREVIEW_ROT) out_verts = np.asarray(mesh.vertices, dtype=np.float64) if out_verts.size == 0: mesh_axis_range = PLOT_AXIS_RANGE * PLOT_MESH_AXIS_PADDING wire_axis_range = PLOT_AXIS_RANGE * PLOT_WIRE_AXIS_PADDING else: radius = float(np.max(np.abs(out_verts))) mesh_axis_range = max(PLOT_AXIS_RANGE, radius * PLOT_MESH_AXIS_PADDING) wire_axis_range = max(PLOT_AXIS_RANGE, radius * PLOT_WIRE_AXIS_PADDING) fd, download_path = tempfile.mkstemp(suffix=".glb", prefix="meshflow_") os.close(fd) mesh.export(download_path) return ( mesh_to_plotly_solid(mesh, axis_range=mesh_axis_range), mesh_to_plotly_wireframe(mesh, axis_range=wire_axis_range), download_path, loaded_num_verts, ) def build_ui( pipeline: MeshFlowPipeline, args: argparse.Namespace, default_num_verts: int, config_num_latents: int, supports_num_verts_scaling: bool, ) -> gr.Blocks: catalog_paths, catalog_thumbs = discover_gallery_examples() global _PIPELINE _PIPELINE = pipeline with gr.Blocks(title=APP_TAB_TITLE) as demo: num_verts_state = gr.State(default_num_verts) geometry_path_state = gr.State(None) runtime_args_state = gr.State(args) gallery_order_state = gr.State([]) gr.HTML(APP_HEADER_HTML, elem_classes="mf-header-wrap") gallery_dataset = None if catalog_paths: gallery_dataset = gr.Dataset( components=[ gr.Image( type="filepath", show_label=False, interactive=False, render=False, ) ], samples=[], type="index", layout="gallery", samples_per_page=GALLERY_EXAMPLES_PER_PAGE, label="Examples", container=True, elem_classes="mf-gallery", ) with gr.Column(elem_classes="mf-workspace"): with gr.Row(equal_height=False, elem_classes="mf-main-row"): with gr.Column(scale=4): input_file = gr.File( label="Input Geometry", file_types=list(GEOMETRY_EXTS), height=168, elem_classes="mf-input-file", ) input_preview_plot = gr.Plot( label="Input preview", visible=True, elem_classes="mf-input-preview-plot", ) with gr.Accordion("Advanced options", open=False): gr.Markdown( NUM_VERTS_CONTROL_NOTE, visible=supports_num_verts_scaling ) gr.Markdown( f"Normalization divisor: **num_latents = {config_num_latents}** " f"(from `mesh_model.num_latents` in config).", visible=supports_num_verts_scaling, ) num_verts = gr.Slider( NUM_VERTS_MIN, NUM_VERTS_MAX, value=default_num_verts, step=NUM_VERTS_STEP, label="num_verts (proj_cond numerator, num_verts / num_latents)", visible=supports_num_verts_scaling, ) seed = gr.Number( value=args.seed, precision=0, label="Random seed" ) guidance = gr.Slider( 1.0, 15.0, value=args.guidance_scale or pipeline.guidance_scale, step=0.1, label="Classifier-free guidance", ) steps = gr.Slider( 1, 100, value=args.steps or pipeline.num_inference_steps, step=1, label="Sampling steps", ) input_image = gr.Image( type="pil", label="Reference Image (optional)", height=240 ) with gr.Column(scale=6, elem_classes="mf-output-col"): run_btn = gr.Button( "Generate", variant="primary", elem_classes="mf-full-width-action mf-generate-action", ) with gr.Row(equal_height=True, elem_classes="mf-viewer-grid"): mesh_solid_plot = gr.Plot(label="Output Mesh", scale=1) mesh_wire_plot = gr.Plot(label="Output Wireframe", scale=1) mesh_download = gr.DownloadButton( "Download GLB", variant="secondary", elem_classes="mf-full-width-action mf-download-action", ) run_inputs = [ num_verts_state, runtime_args_state, geometry_path_state, input_file, input_image, steps, guidance, seed, ] if supports_num_verts_scaling: run_inputs.append(num_verts) run_btn.click( fn=run_meshflow, inputs=run_inputs, outputs=[mesh_solid_plot, mesh_wire_plot, mesh_download, num_verts_state], show_progress="full", ) input_file.upload( fn=preview_input, inputs=[input_file], outputs=[input_preview_plot, geometry_path_state], ) input_file.clear( fn=clear_input_preview, outputs=[input_preview_plot, geometry_path_state], queue=False, ) if gallery_dataset is not None: def shuffle_gallery(_paths=catalog_paths, _thumbs=catalog_thumbs): order = list(range(len(_paths))) random.shuffle(order) return gr.Dataset(samples=[[_thumbs[i]] for i in order]), [ _paths[i] for i in order ] demo.load( fn=shuffle_gallery, outputs=[gallery_dataset, gallery_order_state], queue=False, ) gallery_dataset.click( fn=pick_gallery_example, inputs=[gallery_dataset, gallery_order_state], outputs=[input_file, geometry_path_state, input_preview_plot], show_progress="hidden", queue=False, ) return demo def main() -> None: parser = argparse.ArgumentParser(description="MeshFlow Gradio demo") parser.add_argument( "--checkpoint_bundle", type=str, default=os.environ.get("MESHFLOW_CHECKPOINT_BUNDLE", DEFAULT_CHECKPOINT_BUNDLE), choices=CHECKPOINT_BUNDLES, help=( f"Checkpoint subfolder in {CHECKPOINT_REPO_ID} and ckpt/. " f"Use {CHECKPOINT_BUNDLE_NUM_VERTS!r} for num_verts control (default does not use vertex number condition)." ), ) parser.add_argument( "--model_path", type=str, default=None, help=( "Explicit model bundle directory (config.yaml + model.pth). " f"If omitted, use local ckpt// or download from " f"Hugging Face ({CHECKPOINT_REPO_ID}//)." ), ) parser.add_argument("--gpu", type=int, default=0) parser.add_argument( "--dtype", type=str, default="fp16", choices=AUTOCAST_DTYPE_CHOICES, help="autocast dtype: bf16, fp16, or fp32 (default: fp16)", ) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--steps", type=int, default=None) parser.add_argument("--guidance_scale", type=float, default=None) parser.add_argument( "--num_verts", type=int, default=None, help=( "initial proj_cond numerator (num_verts / mesh_model.num_latents). " "Only effective when denoiser_model.use_proj_cond_on_temb is enabled in config " f"({NUM_VERTS_MIN}-{NUM_VERTS_MAX}; default: {DEFAULT_NUM_VERTS})" ), ) parser.add_argument( "--compile", action=argparse.BooleanOptionalAction, default=False, help=( "torch.compile models for faster inference (CUDA only, default off). " "Leave off on ZeroGPU Spaces: TorchDynamo recompiles on every cold GPU " ), ) parser.add_argument("--server_name", type=str, default="0.0.0.0") parser.add_argument("--server_port", type=int, default=7860) args = parser.parse_args() args.model_path = resolve_model_path(args.model_path, args.checkpoint_bundle) cfg = OmegaConf.load(Path(args.model_path) / CHECKPOINT_CONFIG_FILENAME) use_proj_cond = bool(cfg.system.denoiser_model.get("use_proj_cond_on_temb", False)) config_num_latents = int(cfg.system.mesh_model.num_latents) if not use_proj_cond and args.num_verts is not None: print( "[MeshFlow] Ignoring --num_verts: denoiser_model.use_proj_cond_on_temb is disabled in config" ) device = f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu" if use_proj_cond: default_num_verts = clamp_num_verts(args.num_verts or DEFAULT_NUM_VERTS) pipeline = MeshFlowPipeline.from_pretrained( args.model_path, device=device, dtype=args.dtype, compile_models=args.compile, num_verts=default_num_verts, ) else: default_num_verts = int(cfg.data.n_verts) pipeline = MeshFlowPipeline.from_pretrained( args.model_path, device=device, dtype=args.dtype, compile_models=args.compile, ) args.supports_num_verts_scaling = use_proj_cond demo = build_ui( pipeline, args, default_num_verts, config_num_latents, use_proj_cond ) allowed_paths = [] if GALLERY_DIR.is_dir(): allowed_paths.extend( str(p) for p in (GALLERY_SURFACE_PC_DIR, GALLERY_THUMBNAIL_DIR) if p.is_dir() ) demo.launch( server_name=args.server_name, server_port=args.server_port, allowed_paths=allowed_paths or None, theme=MESHFLOW_THEME, css=CUSTOM_CSS, head=APP_HEAD, ) if __name__ == "__main__": main()

{_BRAND_LETTERS}{APP_VENUE}