| |
| |
| |
| |
| |
|
|
| 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'<span class="mf-brand-letter" style="--m-i:{i}">{ch}</span>' |
| for i, ch in enumerate(APP_TITLE) |
| ) |
| _HOWTO_LI = "".join(f"<li><strong>{a}</strong> — {b}</li>" for a, b in HOW_TO_STEPS) |
| _AUTHORS = "".join( |
| f'<a href="{u}" target="_blank" rel="noopener noreferrer">{n}</a>' |
| for u, n in PAPER_AUTHORS |
| ) |
| _HEADER_LINKS = " ".join( |
| f'<a class="mf-link" href="{u}" target="_blank" rel="noopener noreferrer">{t}</a>' |
| for u, t in ( |
| (PROJECT_PAGE_URL, "Project"), |
| (GITHUB_URL, "GitHub"), |
| (HF_MODEL_URL, "Model"), |
| (ARXIV_URL, "Paper"), |
| ) |
| ) |
| APP_HEADER_HTML = f""" |
| <div class="mf-app-header"> |
| <div class="mf-app-header-top"> |
| <div class="mf-app-title-wrap"> |
| <h1 class="mf-app-title"><span class="mf-brand-word" aria-label="{APP_TITLE}">{_BRAND_LETTERS}</span><span class="mf-venue-badge">{APP_VENUE}</span></h1> |
| <p class="mf-app-tagline">{APP_TAGLINE}</p> |
| <p class="mf-authors">{_AUTHORS}</p> |
| </div> |
| <div class="mf-app-links">{_HEADER_LINKS}</div> |
| </div> |
| <details class="mf-howto-details" open> |
| <summary>How to use</summary> |
| <ol class="mf-howto-steps">{_HOWTO_LI}</ol> |
| </details> |
| </div> |
| """ |
|
|
| _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 = ( |
| "<script>(function(){function f(){document.documentElement.classList.remove('dark');" |
| "document.documentElement.style.colorScheme='light';" |
| "document.querySelectorAll('.dark').forEach(function(e){e.classList.remove('dark');});}" |
| "f();new MutationObserver(function(){requestAnimationFrame(f);}).observe(" |
| "document.documentElement,{attributes:true,attributeFilter:['class']});})();</script>" |
| ) |
|
|
| GALLERY_PAGINATION_SCROLL_JS = ( |
| "<script>(function(){document.addEventListener('click',function(e){" |
| "var btn=e.target.closest('.mf-gallery .paginate button');" |
| "if(!btn)return;var y=window.scrollY;" |
| "function fix(){window.scrollTo(0,y);}" |
| "[0,1,16,48,120,240].forEach(function(d){setTimeout(fix,d);});" |
| "},{capture:true});})();</script>" |
| ) |
|
|
| GALLERY_REVEAL_JS = ( |
| "<script>(function(){var revealed=false;" |
| "function reveal(){if(revealed)return;var g=document.querySelector('.mf-gallery');" |
| "if(!g||!g.querySelector('.gallery-item'))return;revealed=true;" |
| "g.classList.add('mf-gallery-ready');}" |
| "new MutationObserver(function(){requestAnimationFrame(reveal);})" |
| ".observe(document.body,{childList:true,subtree:true});})();</script>" |
| ) |
|
|
| 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 |
|
|
|
|
| |
| _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 |
|
|
| |
| 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/<checkpoint_bundle>/ or download from " |
| f"Hugging Face ({CHECKPOINT_REPO_ID}/<checkpoint_bundle>/)." |
| ), |
| ) |
| 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() |
|
|
