Reverb/open3dforge
0
1"""2Open3DForge — Image-to-game-ready 3D asset pipeline3====================================================4 5Milestone 1: Foundation6-----------------------7This is the scaffold. It validates that:8 - The HF Space builds with the right Gradio + ZeroGPU configuration9 - `@spaces.GPU` allocates and releases an H200 successfully10 - `gr.Model3D` renders GLB files inline11 - The 5-tab UI structure is sound12 - Workspace folder management works13 - Quota tracking persists across page reloads14 15Subsequent milestones fill in the actual pipeline stages.16 17Deployed exclusively on HF Spaces ZeroGPU — no local execution path.18"""19 20from __future__ import annotations21 22import os23import subprocess24import sys25import time26from pathlib import Path27 28# nvdiffrast is compiled against CUDA 13.0 and its .so links to libcudart.so.13.29# os.environ["LD_LIBRARY_PATH"] only affects child processes, not the running30# Python interpreter. Use ctypes.CDLL with RTLD_GLOBAL to load libcudart.so.1331# into the current process before nvdiffrast is imported.32import ctypes, glob as _glob33 34def _preload_cudart() -> None:35 cuda_lib_dir = "/cuda-image/usr/local/cuda-13.0/lib64"36 # Try versioned name first, then unversioned fallback.37 for pattern in (38 f"{cuda_lib_dir}/libcudart.so.13*",39 f"{cuda_lib_dir}/libcudart.so",40 ):41 matches = _glob.glob(pattern)42 if matches:43 try:44 ctypes.CDLL(matches[0], mode=ctypes.RTLD_GLOBAL)45 print(f"Preloaded {matches[0]}", flush=True)46 except OSError as e:47 print(f"Preload failed for {matches[0]}: {e}", flush=True)48 return49 50_preload_cudart()51 52import gradio as gr53import spaces # HF ZeroGPU54 55 56# GPU arch list used to compile nvdiffrast.57# +PTX on the last entry generates portable PTX code that the CUDA runtime58# JIT-compiles for any newer GPU (e.g. Blackwell SM 10.0) at first use.59_NVDR_ARCH = "8.0;8.9;9.0+PTX"60_NVDR_MARKER = Path("/tmp/nvdiffrast_arch.txt")61 62 63def _install_nvdiffrast() -> None:64 """Install nvdiffrast from source at startup.65 66 Must run here (not requirements.txt) because pip's build isolation hides67 torch from nvdiffrast's build system.68 69 Shim strategy: ZeroGPU machine has CUDA 13.0 toolkit but torch 2.9.1 was70 compiled with CUDA 12.8. PyTorch's _check_cuda_version raises on mismatch.71 We put a shim nvcc on PATH that reports 12.8 for --version but delegates72 real compilation to the actual CUDA 13.0 nvcc.73 74 Arch strategy: compile native kernels for SM 8.0 (A100) and 8.9 (L4/L40),75 plus PTX for SM 9.0 (H100/H200). The +PTX entry lets the CUDA 13.0 runtime76 JIT-compile for any newer architecture (e.g. SM 10.0 Blackwell) at first use.77 78 Marker file: /tmp/nvdiffrast_arch.txt stores the arch string used. If it79 matches _NVDR_ARCH, skip rebuild on warm container restarts.80 """81 import os82 import shutil83 import tempfile84 import torch85 86 # Check if already installed with the current arch — skip rebuild if so.87 try:88 import nvdiffrast # noqa: F40189 if _NVDR_MARKER.exists() and _NVDR_MARKER.read_text().strip() == _NVDR_ARCH:90 print("nvdiffrast: already installed with correct arch.", flush=True)91 return92 print("nvdiffrast: arch mismatch or missing marker — forcing reinstall.", flush=True)93 subprocess.run(94 [sys.executable, "-m", "pip", "uninstall", "nvdiffrast", "-y"],95 check=False, capture_output=True,96 )97 except (ImportError, OSError):98 pass99 100 torch_cuda = torch.version.cuda or "12.8"101 print(f"nvdiffrast: torch.version.cuda={torch_cuda}", flush=True)102 103 # Locate real nvcc.104 real_nvcc: str | None = None105 real_cuda_home: str | None = None106 for candidate in [107 os.environ.get("CUDA_HOME", ""),108 "/cuda-image/usr/local/cuda-13.0",109 "/cuda-image/usr/local/cuda-12.8",110 "/cuda-image/usr/local/cuda",111 "/usr/local/cuda-13.0",112 "/usr/local/cuda",113 ]:114 if candidate and os.path.isfile(os.path.join(candidate, "bin", "nvcc")):115 real_cuda_home = candidate116 real_nvcc = os.path.join(candidate, "bin", "nvcc")117 break118 if real_nvcc is None:119 real_nvcc = shutil.which("nvcc")120 print(f"nvdiffrast: real_nvcc={real_nvcc}", flush=True)121 122 env = os.environ.copy()123 124 if real_nvcc:125 fake_cuda = Path(tempfile.mkdtemp(prefix="cuda_shim_"))126 (fake_cuda / "bin").mkdir()127 if real_cuda_home:128 for sub in ("include", "lib", "lib64"):129 src = Path(real_cuda_home) / sub130 if src.exists():131 (fake_cuda / sub).symlink_to(src)132 shim = fake_cuda / "bin" / "nvcc"133 shim.write_text(134 "#!/bin/sh\n"135 'case "$1" in\n'136 ' --version)\n'137 f' echo "Cuda compilation tools, release {torch_cuda}, V{torch_cuda}.0"\n'138 ' ;;\n'139 ' *)\n'140 f' exec {real_nvcc} "$@"\n'141 ' ;;\n'142 'esac\n'143 )144 shim.chmod(0o755)145 env["CUDA_HOME"] = str(fake_cuda)146 env["CUDA_PATH"] = str(fake_cuda)147 print(f"nvdiffrast: shim CUDA_HOME={fake_cuda}", flush=True)148 149 # PTX in the arch list lets the runtime JIT-compile for any newer GPU150 # (e.g. Blackwell SM 10.0) that lacks a pre-compiled cubin.151 env["TORCH_CUDA_ARCH_LIST"] = _NVDR_ARCH152 153 print(f"Building nvdiffrast (arch={_NVDR_ARCH}, ~2 min)...", flush=True)154 result = subprocess.run(155 [sys.executable, "-m", "pip", "install", "--no-build-isolation",156 "git+https://github.com/NVlabs/nvdiffrast.git"],157 env=env, check=False, capture_output=True, text=True,158 )159 print(result.stdout[-3000:] if result.stdout else "", flush=True)160 if result.returncode != 0:161 print(f"nvdiffrast build FAILED (rc={result.returncode}):", flush=True)162 print(result.stderr[-2000:] if result.stderr else "", flush=True)163 else:164 print("nvdiffrast build succeeded.", flush=True)165 _NVDR_MARKER.write_text(_NVDR_ARCH)166 167 168_install_nvdiffrast()169 170from src import quota, ui_helpers, workspace171from src.stages.stage1_generate import generate_trellis, generate_hunyuan172 173# ---------------------------------------------------------------------------174# ZeroGPU test function — proves that GPU allocation works.175# Replaced in Milestone 2 with the actual TRELLIS.2 generator.176# ---------------------------------------------------------------------------177 178@spaces.GPU(duration=15)179def zerogpu_smoke_test() -> str:180 """Allocate a GPU, run a trivial torch op, release. Reports timing."""181 start = time.time()182 try:183 import torch184 if not torch.cuda.is_available():185 return "❌ GPU not available inside @spaces.GPU. Something is wrong."186 187 device = torch.device("cuda")188 # Trivial GPU work189 a = torch.randn(1024, 1024, device=device)190 b = torch.randn(1024, 1024, device=device)191 c = a @ b192 torch.cuda.synchronize()193 result_sum = float(c.sum().item())194 195 gpu_name = torch.cuda.get_device_name(0)196 vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9197 elapsed = time.time() - start198 199 quota.record_usage("smoke_test", elapsed)200 201 return (202 f"✅ **GPU allocation successful**\n\n"203 f"- Device: `{gpu_name}`\n"204 f"- VRAM: `{vram_total:.1f} GB`\n"205 f"- Test op (1024×1024 matmul): `{elapsed:.2f}s`\n"206 f"- Result sum: `{result_sum:.2f}` (sanity check, non-zero = ✓)\n\n"207 f"ZeroGPU integration is working. Ready for Milestone 2."208 )209 except Exception as e:210 elapsed = time.time() - start211 return f"❌ GPU test failed after {elapsed:.2f}s:\n```\n{type(e).__name__}: {e}\n```"212 213 214# ---------------------------------------------------------------------------215# Stage stubs — return placeholder messages until milestones implement them.216# ---------------------------------------------------------------------------217 218def _stub(stage: str) -> str:219 return (220 f"🚧 **{stage}** — not implemented yet.\n\n"221 f"This stub will be replaced in a future milestone. "222 f"See `PLAN.md` for the full pipeline spec."223 )224 225 226def handle_generate(images, model, quality, seed, _steps, _octree, tex_size, _symmetry, do_rembg):227 """Dispatch to the correct generation backend.228 229 Yields (status_markdown, viewer_path_or_None) tuples for streaming.230 The viewer path is yielded only in the final tuple so Gradio can serve231 the GLB directly without a separate state lookup.232 """233 if "TRELLIS" in model:234 yield f"⚙️ **TRELLIS.2** · {quality}\n\nContacting remote Space...", None235 result = generate_trellis(images, quality, int(seed), int(tex_size))236 viewer_path = ui_helpers.get_viewer_model_path()237 yield result, viewer_path238 return239 yield from generate_hunyuan(images, quality, int(seed), int(tex_size), do_rembg=bool(do_rembg))240 241 242@spaces.GPU(duration=600)243def run_post_process(244 do_repair, do_cleanup, do_decimate, target_faces,245 do_symmetry, do_unwrap, do_normal_bake, normal_format,246 do_albedo, do_material, do_ao, ao_quality,247 do_inpaint, do_lods, do_collision, pivot, scale_m,248):249 """Run post-processing pipeline. Yields cumulative status markdown for streaming."""250 state = workspace.get_state()251 current_glb = state.raw_gen_glb or state.high_poly_glb252 if not current_glb or not current_glb.exists():253 yield "❌ No generated asset found. Run Stage 1 (Generate) first."254 return255 256 log = []257 258 def _emit(line: str):259 log.append(line)260 return "\n".join(log)261 262 if do_repair:263 yield _emit("⏳ Repairing mesh (pymeshfix)...")264 try:265 from src.stages.stage2_repair import repair_mesh266 current_glb, msg = repair_mesh(current_glb)267 yield _emit(f"✅ {msg}")268 except Exception as e:269 yield _emit(f"⚠️ Repair error: {e}")270 271 if do_cleanup:272 yield _emit("⏳ Cleaning geometry (trimesh)...")273 try:274 from src.stages.stage2_cleanup import cleanup_mesh275 current_glb, msg = cleanup_mesh(current_glb)276 yield _emit(f"✅ {msg}")277 except Exception as e:278 yield _emit(f"⚠️ Cleanup error: {e}")279 280 if do_decimate:281 yield _emit(f"⏳ Decimating to {int(target_faces):,} faces...")282 try:283 from src.stages.stage2_decimate import decimate_mesh_final284 current_glb, msg = decimate_mesh_final(current_glb, int(target_faces))285 yield _emit(f"✅ {msg}")286 except Exception as e:287 yield _emit(f"⚠️ Decimation error: {e}")288 289 if do_symmetry:290 yield _emit("⏳ Enforcing bilateral symmetry...")291 try:292 from src.stages.stage2_symmetry import apply_symmetry293 current_glb, msg = apply_symmetry(current_glb, "bilateral-X")294 yield _emit(f"✅ {msg}")295 except Exception as e:296 yield _emit(f"⚠️ Symmetry error: {e}")297 298 if do_unwrap:299 yield _emit("⏳ UV unwrapping (xatlas)...")300 try:301 from src.stages.stage2_uv import unwrap_uvs302 current_glb, msg = unwrap_uvs(current_glb)303 yield _emit(f"✅ {msg}")304 except Exception as e:305 yield _emit(f"⚠️ UV unwrap error: {e}")306 307 if do_normal_bake:308 yield _emit("⏳ Baking normal map (GPU)...")309 try:310 from src.stages.stage2_bake_normal import bake_normal_map311 st = workspace.get_state()312 hp = st.high_poly_glb313 lo = st.unwrapped_glb or current_glb314 if not hp or not hp.exists():315 yield _emit("⚠️ Normal bake: no high-poly GLB. Generate first.")316 else:317 _gl, _dx, msg = bake_normal_map(hp, lo, map_size=2048, dx_format=True)318 yield _emit(f"✅ {msg}")319 except Exception as e:320 yield _emit(f"⚠️ Normal bake error: {e}")321 322 st = workspace.get_state()323 hp = st.high_poly_glb324 lo = st.unwrapped_glb or current_glb325 326 if do_albedo:327 yield _emit("⏳ Baking albedo map (GPU)...")328 try:329 from src.stages.stage2_bake_albedo import bake_albedo330 if not hp or not hp.exists():331 yield _emit("⚠️ Albedo bake: no high-poly. Generate first.")332 else:333 _, msg = bake_albedo(hp, lo, map_size=2048)334 yield _emit(f"✅ {msg}")335 except Exception as e:336 yield _emit(f"⚠️ Albedo bake error: {e}")337 338 if do_material:339 yield _emit("⏳ Baking material maps (GPU)...")340 try:341 from src.stages.stage2_bake_albedo import bake_material342 if not hp or not hp.exists():343 yield _emit("⚠️ Material bake: no high-poly. Generate first.")344 else:345 _, _, msg = bake_material(hp, lo, map_size=2048)346 yield _emit(f"✅ {msg}")347 except Exception as e:348 yield _emit(f"⚠️ Material bake error: {e}")349 350 if do_ao:351 yield _emit(f"⏳ Baking AO ({ao_quality}, ray casting)...")352 try:353 from src.stages.stage2_bake_ao import bake_ao354 _, msg = bake_ao(current_glb, lo, map_size=2048, quality=ao_quality)355 yield _emit(f"✅ {msg}")356 except Exception as e:357 yield _emit(f"⚠️ AO bake error: {e}")358 359 st2 = workspace.get_state()360 361 if do_albedo or do_material or do_ao:362 yield _emit("⏳ Packing ORM texture...")363 try:364 from src.stages.stage2_finalize import pack_orm365 _, msg = pack_orm(st2.ao_png, st2.roughness_png, st2.metallic_png)366 yield _emit(f"✅ {msg}")367 except Exception as e:368 yield _emit(f"⚠️ ORM pack error: {e}")369 370 if do_lods:371 yield _emit("⏳ Generating LODs...")372 try:373 from src.stages.stage2_finalize import generate_lods374 lod_src = st2.final_glb or st2.low_poly_glb or current_glb375 _, msg = generate_lods(lod_src)376 yield _emit(f"✅ {msg}")377 except Exception as e:378 yield _emit(f"⚠️ LOD error: {e}")379 380 if do_collision:381 yield _emit("⏳ Generating collision mesh (CoACD)...")382 try:383 from src.stages.stage2_finalize import generate_collision384 col_src = st2.low_poly_glb or current_glb385 _, msg = generate_collision(col_src)386 yield _emit(f"✅ {msg}")387 except Exception as e:388 yield _emit(f"⚠️ Collision error: {e}")389 390 yield _emit("⏳ Setting pivot and scale...")391 try:392 from src.stages.stage2_finalize import set_pivot, validate_scale393 piv_src = st2.low_poly_glb or current_glb394 piv_src, msg = set_pivot(piv_src, pivot)395 yield _emit(f"✅ {msg}")396 _, msg = validate_scale(piv_src, float(scale_m))397 yield _emit(f"✅ {msg}")398 except Exception as e:399 yield _emit(f"⚠️ Pivot/scale error: {e}")400 401 if do_inpaint:402 yield _emit("🚧 SDXL inpaint — not yet implemented")403 404 final = workspace.get_state()405 out = final.final_glb or final.low_poly_glb or final.cleaned_glb or final.repaired_glb406 if out and out.exists():407 yield _emit(f"\n**Output:** `{out.name}` · ready for Stage 3 or Export.")408 409 410def handle_auto_rig(rig_type, seed, spring, fmt):411 from src.stages.stage3_rig import auto_rig412 _, msg = auto_rig(rig_type=rig_type, seed=int(seed))413 return msg414 415 416def handle_export(engine, asset_name, asset_type, include_lods, include_collision):417 from src.stages.stage4_export import export_ue5, _checklist418 state = workspace.get_state()419 420 if engine != "UE5":421 return f"🚧 {engine} export — not implemented. UE5 is the only supported engine.", None422 423 # Game-ready checklist424 issues = _checklist(state)425 checklist_md = ""426 if issues:427 checklist_md = "\n\n⚠️ **Checklist warnings:**\n" + "\n".join(f"- {i}" for i in issues)428 429 try:430 zip_path, msg, _ = export_ue5(431 asset_name=asset_name.strip() or "Asset",432 asset_type=asset_type,433 include_lods=include_lods,434 include_collision=include_collision,435 )436 return msg + checklist_md, str(zip_path)437 except Exception as e:438 return f"❌ Export failed: {e}{checklist_md}", None439 440 441# ---------------------------------------------------------------------------442# Presets tab logic (M11 — real save/load wired)443# ---------------------------------------------------------------------------444 445def ui_refresh_presets() -> gr.Dropdown:446 return gr.Dropdown(choices=workspace.list_presets(), label="Saved presets")447 448 449def ui_save_preset(450 name: str,451 gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,452 pp_repair, pp_cleanup, pp_decimate, pp_target_faces,453 pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,454 pp_albedo, pp_material, pp_ao, pp_ao_quality,455 pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m,456 rig_type, rig_seed, rig_format,457 ex_engine, ex_type,458) -> tuple[str, gr.Dropdown]:459 if not name or not name.strip():460 return "❌ Preset name cannot be empty.", ui_refresh_presets()461 config = {462 "name": name, "version": 1, "created_at": time.time(),463 "stage1": {"model": gen_model, "quality": gen_quality, "seed": int(gen_seed), "tex_size": int(gen_tex_size), "rembg": bool(gen_rembg)},464 "stage2": {465 "repair": pp_repair, "cleanup": pp_cleanup, "decimate": pp_decimate,466 "target_faces": int(pp_target_faces), "symmetry": pp_symmetry,467 "unwrap": pp_unwrap, "normal_bake": pp_normal_bake, "normal_format": pp_normal_format,468 "albedo_bake": pp_albedo, "material_bake": pp_material, "ao": pp_ao, "ao_quality": pp_ao_quality,469 "inpaint": pp_inpaint, "lods": pp_lods, "collision": pp_collision,470 "pivot": pp_pivot, "scale_m": float(pp_scale_m),471 },472 "stage3": {"rig_type": rig_type, "seed": int(rig_seed), "format": rig_format},473 "stage4": {"engine": ex_engine, "asset_type": ex_type},474 }475 workspace.save_preset(name, config)476 return f"✅ Saved preset: `{name}`", ui_refresh_presets()477 478 479def ui_load_preset(name: str):480 """Apply a saved preset to all UI components."""481 _no_op = gr.update()482 n = 27 # number of component outputs after pr_status483 if not name:484 return ("❌ Select a preset first.",) + (_no_op,) * n485 try:486 cfg = workspace.load_preset(name)487 except FileNotFoundError:488 return (f"❌ Preset '{name}' not found.",) + (_no_op,) * n489 490 s1 = cfg.get("stage1", {})491 s2 = cfg.get("stage2", {})492 s3 = cfg.get("stage3", {})493 s4 = cfg.get("stage4", {})494 495 return (496 f"✅ Loaded preset: `{name}`",497 # Stage 1498 gr.update(value=s1.get("model", "Hunyuan3D-2.1 (Organic / Characters)")),499 gr.update(value=s1.get("quality", "Balanced (~60s)")),500 gr.update(value=s1.get("seed", 42)),501 gr.update(value=s1.get("tex_size", 2048)),502 gr.update(value=s1.get("rembg", True)),503 # Stage 2504 gr.update(value=s2.get("repair", True)),505 gr.update(value=s2.get("cleanup", True)),506 gr.update(value=s2.get("decimate", True)),507 gr.update(value=s2.get("target_faces", 25000)),508 gr.update(value=s2.get("symmetry", False)),509 gr.update(value=s2.get("unwrap", True)),510 gr.update(value=s2.get("normal_bake", True)),511 gr.update(value=s2.get("normal_format", "DirectX (UE5)")),512 gr.update(value=s2.get("albedo_bake", True)),513 gr.update(value=s2.get("material_bake", True)),514 gr.update(value=s2.get("ao", True)),515 gr.update(value=s2.get("ao_quality", "Standard")),516 gr.update(value=s2.get("inpaint", False)),517 gr.update(value=s2.get("lods", True)),518 gr.update(value=s2.get("collision", True)),519 gr.update(value=s2.get("pivot", "bottom_center")),520 gr.update(value=s2.get("scale_m", 1.8)),521 # Stage 3522 gr.update(value=s3.get("rig_type", "Humanoid")),523 gr.update(value=s3.get("seed", 0)),524 gr.update(value=s3.get("format", "FBX (UE5 recommended)")),525 # Stage 4526 gr.update(value=s4.get("engine", "UE5")),527 gr.update(value=s4.get("asset_type", "Prop (SM_)")),528 )529 530 531def ui_delete_preset(name: str) -> tuple[str, gr.Dropdown]:532 if not name:533 return "❌ Select a preset to delete.", ui_refresh_presets()534 if workspace.delete_preset(name):535 return f"🗑️ Deleted preset: `{name}`", ui_refresh_presets()536 return f"❌ Preset not found: `{name}`", ui_refresh_presets()537 538 539# ---------------------------------------------------------------------------540# Layout541# ---------------------------------------------------------------------------542 543CUSTOM_CSS = """544.status-bar {545 font-size: 0.85em;546 color: #888;547 padding: 8px 12px;548 border-top: 1px solid #333;549 margin-top: 12px;550}551.asset-summary {552 font-size: 0.9em;553 background: rgba(255,255,255,0.03);554 padding: 12px;555 border-radius: 6px;556 border: 1px solid rgba(255,255,255,0.08);557}558"""559 560 561def build_ui() -> gr.Blocks:562 with gr.Blocks(563 title="Open3DForge",564 ) as demo:565 566 # --- Header --------------------------------------------------------567 gr.Markdown(568 "# 🛠️ Open3DForge\n"569 "*Personal image-to-game-ready 3D asset pipeline · UE5-first · "570 "Built on HF ZeroGPU*"571 )572 573 # --- Tabs ----------------------------------------------------------574 with gr.Tabs() as tabs:575 576 # ============ Tab 1: Generate =================================577 with gr.Tab("1. Generate", id=1):578 gr.Markdown(579 "### Stage 1 — Image to 3D\n"580 "Upload 1–4 reference images. Multi-view dramatically "581 "improves quality for characters (front / 3-quarter / side / back)."582 )583 with gr.Row():584 with gr.Column(scale=1):585 gen_images = gr.File(586 label="Reference images (1–4)",587 file_count="multiple",588 file_types=["image"],589 )590 gen_model = gr.Radio(591 choices=[592 "Hunyuan3D-2.1 (Organic / Characters)",593 "TRELLIS.2 (Hard Surface)",594 ],595 value="Hunyuan3D-2.1 (Organic / Characters)",596 label="Generation model",597 )598 gen_quality = gr.Radio(599 choices=["Fast (~30s)", "Balanced (~60s)", "Hero (~90s)"],600 value="Balanced (~60s)",601 label="Quality preset",602 )603 gen_rembg = gr.Checkbox(604 value=True,605 label="Remove background automatically",606 info="Uses rembg IS-Net to strip background before generation. "607 "Disable if your image already has a transparent background.",608 )609 with gr.Accordion("Advanced", open=False):610 gen_seed = gr.Number(value=42, label="Seed", precision=0)611 gen_steps = gr.Slider(612 20, 50, value=35, step=5,613 label="Inference steps",614 )615 gen_octree = gr.Dropdown(616 choices=[256, 384, 512], value=384,617 label="Octree resolution",618 )619 gen_tex_size = gr.Dropdown(620 choices=[1024, 2048, 4096], value=2048,621 label="Texture size",622 )623 gen_symmetry = gr.Radio(624 choices=["off", "bilateral", "radial"],625 value="off",626 label="Symmetry hint",627 )628 gen_btn = gr.Button("Generate", variant="primary")629 with gr.Column(scale=1):630 gen_status = gr.Markdown("*Awaiting input.*")631 # Click wired below, after viewer component is defined.632 633 # ============ Tab 2: Post-Process =============================634 with gr.Tab("2. Post-Process", id=2):635 gr.Markdown(636 "### Stage 2 — Mesh cleanup, UV unwrap, texture baking\n"637 "Toggle steps on/off. Decimation has a live preview, the "638 "rest run on confirm."639 )640 with gr.Row():641 with gr.Column(scale=1):642 pp_repair = gr.Checkbox(value=True, label="Mesh repair (pymeshfix)")643 pp_cleanup = gr.Checkbox(value=True, label="Geometry cleanup (PyMeshLab)")644 pp_decimate = gr.Checkbox(value=True, label="Decimation")645 pp_target_faces = gr.Slider(646 1000, 200000, value=25000, step=1000,647 label="Target faces",648 )649 pp_symmetry = gr.Checkbox(value=False, label="Enforce bilateral symmetry")650 pp_unwrap = gr.Checkbox(value=True, label="UV unwrap (xatlas)")651 pp_normal_bake = gr.Checkbox(value=True, label="Normal bake (nvdiffrast)")652 pp_normal_format = gr.Radio(653 choices=["DirectX (UE5)", "OpenGL (Unity/Godot)"],654 value="DirectX (UE5)",655 label="Normal format",656 )657 pp_albedo_bake = gr.Checkbox(value=True, label="Albedo bake")658 pp_material_bake = gr.Checkbox(value=True, label="Material bake (TRELLIS.2 attrs)")659 pp_ao = gr.Checkbox(value=True, label="AO bake")660 pp_ao_quality = gr.Radio(661 choices=["Fast", "Standard", "High"],662 value="Standard",663 label="AO quality",664 )665 pp_inpaint = gr.Checkbox(value=False, label="SDXL inpaint hidden UVs (~30s GPU)")666 pp_lods = gr.Checkbox(value=True, label="Generate LODs (LOD0/1/2)")667 pp_collision = gr.Checkbox(value=True, label="Collision mesh (CoACD)")668 pp_pivot = gr.Radio(669 choices=["bottom_center", "geometric_center", "custom"],670 value="bottom_center",671 label="Pivot point",672 )673 pp_scale_m = gr.Number(value=1.8, label="Real-world height (meters)")674 pp_btn = gr.Button("Run Post-Processing", variant="primary")675 with gr.Column(scale=1):676 pp_face_preview = gr.Markdown("*Face count preview: move the slider.*")677 pp_status = gr.Markdown("*No asset to process. Generate one first.*")678 679 def _decimate_preview(target_faces):680 state = workspace.get_state()681 src = state.low_poly_glb or state.cleaned_glb or state.repaired_glb or state.raw_gen_glb682 if not src or not src.exists():683 return "*Generate an asset first.*"684 try:685 from src.stages.stage2_decimate import decimate_preview686 fc, vc = decimate_preview(src, int(target_faces))687 return f"**Preview:** ~{fc:,} faces · ~{vc:,} vertices at target {int(target_faces):,}"688 except Exception as e:689 return f"Preview error: {e}"690 691 pp_target_faces.change(692 fn=_decimate_preview,693 inputs=pp_target_faces,694 outputs=pp_face_preview,695 )696 _pp_event = pp_btn.click(697 fn=run_post_process,698 inputs=[pp_repair, pp_cleanup, pp_decimate, pp_target_faces,699 pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,700 pp_albedo_bake, pp_material_bake, pp_ao, pp_ao_quality,701 pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m],702 outputs=pp_status,703 )704 705 # ============ Tab 3: Auto-Rig =================================706 with gr.Tab("3. Auto-Rig", id=3):707 gr.Markdown(708 "### Stage 3 — Auto-rigging (optional)\n"709 "Uses UniRig (VAST-AI). For characters and creatures. "710 "After rigging, drop the FBX into [Mixamo](https://mixamo.com) "711 "for free animation presets."712 )713 with gr.Row():714 with gr.Column(scale=1):715 rig_type = gr.Dropdown(716 choices=["Humanoid", "Quadruped", "Bird", "Insect", "Custom"],717 value="Humanoid",718 label="Character type",719 )720 rig_seed = gr.Number(value=0, label="Skeleton seed", precision=0)721 rig_spring = gr.Checkbox(value=False, label="Spring bones (hair/cloth/tail)")722 rig_format = gr.Radio(723 choices=["FBX (UE5 recommended)", "GLB"],724 value="FBX (UE5 recommended)",725 label="Export format",726 )727 rig_btn = gr.Button("Auto-Rig", variant="primary")728 with gr.Column(scale=1):729 rig_status = gr.Markdown("*Process an asset in Stage 2 first.*")730 _rig_event = rig_btn.click(731 fn=handle_auto_rig,732 inputs=[rig_type, rig_seed, rig_spring, rig_format],733 outputs=rig_status,734 )735 736 # ============ Tab 4: Export ===================================737 with gr.Tab("4. Export", id=4):738 gr.Markdown(739 "### Stage 4 — Engine-ready export\n"740 "UE5 default: FBX with DirectX normals + ORM-packed textures."741 )742 with gr.Row():743 with gr.Column(scale=1):744 ex_engine = gr.Dropdown(745 choices=["UE5", "Unity (HDRP)", "Godot 4", "Blender", "Web (Three.js)"],746 value="UE5",747 label="Target engine",748 )749 ex_name = gr.Textbox(value="Asset_01", label="Asset name")750 ex_type = gr.Radio(751 choices=["Character (SK_)", "Prop (SM_)", "Environment (SM_)"],752 value="Prop (SM_)",753 label="Asset type",754 )755 ex_include_lods = gr.Checkbox(value=True, label="Include LODs")756 ex_include_collision = gr.Checkbox(value=True, label="Include collision mesh")757 ex_btn = gr.Button("Export", variant="primary")758 with gr.Column(scale=1):759 ex_status = gr.Markdown("*Nothing to export yet.*")760 ex_file = gr.File(label="Download", visible=True)761 _ex_event = ex_btn.click(762 fn=handle_export,763 inputs=[ex_engine, ex_name, ex_type, ex_include_lods, ex_include_collision],764 outputs=[ex_status, ex_file],765 )766 767 # ============ Tab 5: Presets ==================================768 with gr.Tab("5. Presets", id=5):769 gr.Markdown(770 "### Saved configurations\n"771 "Save the current settings across all tabs as a named preset. "772 "Five defaults ship with the app: `character_UE5_hero`, "773 "`character_UE5_npc`, `prop_UE5_hero`, `prop_UE5_standard`, "774 "`environment_UE5_background`."775 )776 with gr.Row():777 with gr.Column():778 pr_list = gr.Dropdown(779 choices=workspace.list_presets(),780 label="Saved presets",781 )782 pr_refresh = gr.Button("Refresh list", size="sm")783 with gr.Row():784 pr_name = gr.Textbox(label="New preset name", scale=2)785 pr_save = gr.Button("Save current settings", variant="primary", scale=1)786 with gr.Row():787 pr_load = gr.Button("Load selected", variant="secondary", scale=1)788 pr_delete = gr.Button("Delete selected", variant="stop", scale=1)789 pr_status = gr.Markdown()790 pr_refresh.click(fn=ui_refresh_presets, outputs=pr_list)791 pr_load.click(792 fn=ui_load_preset,793 inputs=pr_list,794 outputs=[795 pr_status,796 gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,797 pp_repair, pp_cleanup, pp_decimate, pp_target_faces,798 pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,799 pp_albedo_bake, pp_material_bake, pp_ao, pp_ao_quality,800 pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m,801 rig_type, rig_seed, rig_format,802 ex_engine, ex_type,803 ],804 )805 pr_save.click(806 fn=ui_save_preset,807 inputs=[808 pr_name,809 gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,810 pp_repair, pp_cleanup, pp_decimate, pp_target_faces,811 pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,812 pp_albedo_bake, pp_material_bake, pp_ao, pp_ao_quality,813 pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m,814 rig_type, rig_seed, rig_format,815 ex_engine, ex_type,816 ],817 outputs=[pr_status, pr_list],818 )819 pr_delete.click(fn=ui_delete_preset, inputs=pr_list, outputs=[pr_status, pr_list])820 821 # ============ Tab 6: Diagnostics (hidden in prod, useful now) =822 with gr.Tab("Diagnostics", id=99):823 gr.Markdown(824 "### Milestone 1 — Foundation check\n"825 "Verify the Space environment is working correctly before "826 "building out the pipeline."827 )828 with gr.Row():829 with gr.Column():830 diag_btn = gr.Button("🧪 Run GPU smoke test", variant="primary")831 diag_out = gr.Markdown()832 with gr.Column():833 gr.Markdown("**Workspace state:**")834 diag_state = gr.JSON(value=workspace.get_state().to_dict())835 diag_refresh = gr.Button("Refresh state", size="sm")836 diag_btn.click(fn=zerogpu_smoke_test, outputs=diag_out)837 diag_refresh.click(838 fn=lambda: workspace.get_state().to_dict(),839 outputs=diag_state,840 )841 842 # --- Persistent right-side viewer + asset summary ------------------843 gr.Markdown("---")844 with gr.Row():845 with gr.Column(scale=2):846 viewer = gr.Model3D(847 label="3D viewer",848 value=ui_helpers.get_viewer_model_path(),849 clear_color=[0.1, 0.1, 0.12, 1.0],850 height=500,851 )852 with gr.Column(scale=1):853 summary = gr.Markdown(854 value=ui_helpers.get_asset_summary(),855 elem_classes=["asset-summary"],856 )857 refresh_summary = gr.Button("🔄 Refresh viewer", size="sm")858 refresh_summary.click(859 fn=lambda: (860 ui_helpers.get_viewer_model_path(),861 ui_helpers.get_asset_summary(),862 ),863 outputs=[viewer, summary],864 )865 866 # --- Global status bar --------------------------------------------867 status_bar = gr.Markdown(868 value=ui_helpers.get_status_bar(),869 elem_classes=["status-bar"],870 )871 872 # --- Wire generate button now that viewer is in scope ---------------873 # handle_generate yields (status_str, glb_path_or_None) tuples so that874 # the viewer updates the instant the final mesh is ready — no separate875 # state lookup needed.876 _gen_event = gen_btn.click(877 fn=handle_generate,878 inputs=[gen_images, gen_model, gen_quality, gen_seed,879 gen_steps, gen_octree, gen_tex_size, gen_symmetry,880 gen_rembg],881 outputs=[gen_status, viewer],882 )883 884 # --- Global refresh: every pipeline action updates summary + status bar.885 # For gen_btn the viewer is already updated directly above; for the rest886 # we include the viewer in _global_refresh so post-process / rig / export887 # results also appear.888 def _global_refresh():889 return (890 ui_helpers.get_viewer_model_path(),891 ui_helpers.get_asset_summary(),892 ui_helpers.get_status_bar(),893 )894 895 def _summary_refresh():896 return ui_helpers.get_asset_summary(), ui_helpers.get_status_bar()897 898 # gen_btn: viewer already updated by direct output; only refresh metadata.899 _gen_event.then(fn=_summary_refresh, outputs=[summary, status_bar])900 901 # Other processing buttons: refresh viewer + metadata after completion.902 for _ev in (_pp_event, _rig_event, _ex_event):903 _ev.then(fn=_global_refresh, outputs=[viewer, summary, status_bar])904 905 # Utility buttons: refresh immediately on click.906 for _btn in (diag_btn, diag_refresh, refresh_summary):907 _btn.click(fn=_global_refresh, outputs=[viewer, summary, status_bar])908 909 return demo910 911 912# ---------------------------------------------------------------------------913# Entrypoint914# ---------------------------------------------------------------------------915# On HF Spaces, app.py is executed directly. We construct the demo and call916# .launch() at module level. The HF Space runtime handles all networking;917# we just need to bind to 0.0.0.0:7860.918 919workspace.ensure_dirs()920demo = build_ui()921demo.queue(default_concurrency_limit=1).launch(922 server_name="0.0.0.0",923 server_port=7860,924 show_error=True,925 # Allow Gradio to serve files from the workspace directory so that926 # gr.Model3D can display generated GLBs without a 403 error.927 allowed_paths=[str(workspace.WORKSPACE)],928 # Gradio 6: show_api removed. Use footer_links instead.929 footer_links=["gradio", "settings"],930 theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),931 css=CUSTOM_CSS,932)933