From 19d6f6802fd5d421e112411dd7d59aa38fe8278e Mon Sep 17 00:00:00 2001 From: modelbeast Date: Thu, 16 Jul 2026 09:31:19 +1000 Subject: [PATCH] modelbeast: e2e CLI + env-tunable remesh target - generate_e2e.py: single-command image->textured GLB (Stage 1 + Stage 2), CLI flags for steps/guidance/octree/views/texture-size/seed/no-texture, stage timings + MLX peak-memory logging. Distilled from tests/test_stage1_to_stage2.py. - simplify_mesh_utils.remesh_mesh: honor HY3D_REMESH_FACES env (default 40000 unchanged) so the bake-mesh budget is tunable per node. --- .gitignore | 6 + generate_e2e.py | 160 +++++++++++++++++++++++++ hy3dpaint/utils/simplify_mesh_utils.py | 7 +- 3 files changed, 172 insertions(+), 1 deletion(-) create mode 100644 generate_e2e.py diff --git a/.gitignore b/.gitignore index 72251f2..842e3b0 100644 --- a/.gitignore +++ b/.gitignore @@ -183,3 +183,9 @@ demo_untextured.glb white_mesh_remesh.obj hy3dpaint/ckpt/ +.venv/ +weights/ +*.safetensors +*.glb +*.obj +ckpt/ diff --git a/generate_e2e.py b/generate_e2e.py new file mode 100644 index 0000000..774acfd --- /dev/null +++ b/generate_e2e.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +"""End-to-end image -> textured GLB, fully local on Apple Silicon (MLX). + +Chains Hunyuan3D 2.1 Stage 1 (shape) + Stage 2 (PBR texture) into one command. +Distilled from tests/test_stage1_to_stage2.py, plus CLI flags, stage timings, +and MLX peak-memory reporting for MODELBEAST job logs. + + python generate_e2e.py IMAGE --output STEM [flags] + +Produces (STEM given without extension): + STEM_shape.glb Stage 1 mesh (untextured) [omitted with --no-texture] + STEM.glb final PBR-textured mesh [or the shape mesh if --no-texture] +""" +import argparse +import gc +import os +import sys +import time +from pathlib import Path + +REPO = Path(__file__).resolve().parent + + +def _mlx_mem(fn_names): + import mlx.core as mx + candidates = [] + for name in fn_names: + candidates.append(getattr(mx, name, None)) + metal = getattr(mx, "metal", None) + if metal is not None: + candidates.append(getattr(metal, name, None)) + for fn in candidates: + if fn is None: + continue + try: + return fn() + except Exception: + continue + return None + + +def peak_gb(): + v = _mlx_mem(["get_peak_memory"]) + return v / 1e9 if v is not None else float("nan") + + +def reset_peak(): + _mlx_mem(["reset_peak_memory"]) + + +def main(): + ap = argparse.ArgumentParser(description="Hunyuan3D 2.1 MLX image->textured GLB") + ap.add_argument("image", help="input reference image (ideally a bg-removed cutout)") + ap.add_argument("--output", required=True, help="output path stem (extension ignored)") + # Stage 1 (shape) + ap.add_argument("--steps", type=int, default=50, help="Stage 1 denoising steps") + ap.add_argument("--guidance", type=float, default=7.5, help="Stage 1 CFG scale") + ap.add_argument("--octree-resolution", type=int, default=256, help="Stage 1 marching-cubes grid") + # Stage 2 (texture) + ap.add_argument("--max-num-view", type=int, default=6, help="Stage 2 baked views") + ap.add_argument("--view-resolution", type=int, default=512, help="Stage 2 per-view diffusion res") + ap.add_argument("--texture-size", type=int, default=2048, help="Stage 2 UV atlas res (2048|4096)") + ap.add_argument("--remesh-faces", type=int, default=40000, help="Stage 2 bake-mesh decimation target") + # Common + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--no-texture", action="store_true", help="Stage 1 only (geometry, much faster)") + ap.add_argument("--weights", default="dgrauet/hunyuan3d-2.1-mlx", help="HF repo id or local dir") + a = ap.parse_args() + + # Resolve I/O to absolute BEFORE chdir; the pipeline reads relative config + # paths (hy3dpaint/cfgs/...), so it must run from the repo root. + image_path = str(Path(a.image).resolve()) + out_stem = Path(a.output).resolve() + out_stem.parent.mkdir(parents=True, exist_ok=True) + if not os.path.exists(image_path): + print(f"ERROR: image not found: {image_path}", flush=True) + sys.exit(1) + + os.environ["HY3D_REMESH_FACES"] = str(a.remesh_faces) + os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + + os.chdir(REPO) + sys.path.insert(0, str(REPO / "hy3dshape")) + sys.path.insert(0, str(REPO / "hy3dpaint")) + + # ---- Stage 1: image -> mesh ---- + print(f"[Stage 1] shape generation from {image_path}", flush=True) + from hy3dshape.pipeline_mlx import ShapePipeline + + reset_peak() + t0 = time.time() + shape_pipe = ShapePipeline.from_pretrained(a.weights) + print(f" pipeline ready in {time.time() - t0:.1f}s", flush=True) + + t0 = time.time() + mesh = shape_pipe( + image_path, + num_inference_steps=a.steps, + guidance_scale=a.guidance, + octree_resolution=a.octree_resolution, + seed=a.seed, + ) + dt1 = time.time() - t0 + nverts, nfaces = len(mesh.vertices), len(mesh.faces) + print(f" shape in {dt1:.1f}s ({nverts} verts, {nfaces} faces), " + f"peak {peak_gb():.1f}GB", flush=True) + + if a.no_texture: + final = out_stem.with_suffix(".glb") + mesh.export(str(final)) + print(f"[done] {final} (geometry only, {dt1:.0f}s)", flush=True) + return + + shape_glb = str(out_stem.parent / (out_stem.name + "_shape.glb")) + mesh.export(shape_glb) + print(f" saved shape: {shape_glb}", flush=True) + + # Free the DiT before loading the paint stack (both are memory-heavy) + del shape_pipe, mesh + gc.collect() + reset_peak() + + # ---- Stage 2: mesh + reference -> textured GLB ---- + print("[Stage 2] PBR texture synthesis", flush=True) + from textureGenPipeline_mlx import ( + Hunyuan3DPaintConfigMLX, + Hunyuan3DPaintPipelineMLX, + ) + + cfg = Hunyuan3DPaintConfigMLX(max_num_view=a.max_num_view, resolution=a.view_resolution) + cfg.texture_size = a.texture_size + cfg.mlx_seed = a.seed + + t0 = time.time() + paint_pipe = Hunyuan3DPaintPipelineMLX(cfg) + print(f" pipeline ready in {time.time() - t0:.1f}s", flush=True) + + tex_obj = str(out_stem.with_suffix(".obj")) + t0 = time.time() + paint_pipe( + mesh_path=shape_glb, + image_path=image_path, + output_mesh_path=tex_obj, + use_remesh=True, + save_glb=True, + ) + dt2 = time.time() - t0 + + final = out_stem.with_suffix(".glb") + if not final.exists(): + print(f"ERROR: expected textured GLB not produced at {final}", flush=True) + sys.exit(2) + size_mb = final.stat().st_size / 1024 / 1024 + print(f" texture in {dt2:.1f}s, peak {peak_gb():.1f}GB", flush=True) + print(f"[done] {final} ({size_mb:.1f} MB) " + f"shape={dt1:.0f}s tex={dt2:.0f}s total={dt1 + dt2:.0f}s", flush=True) + + +if __name__ == "__main__": + main() diff --git a/hy3dpaint/utils/simplify_mesh_utils.py b/hy3dpaint/utils/simplify_mesh_utils.py index 5df3606..bb811d1 100644 --- a/hy3dpaint/utils/simplify_mesh_utils.py +++ b/hy3dpaint/utils/simplify_mesh_utils.py @@ -12,12 +12,17 @@ # fine-tuning enabling code and other elements of the foregoing made publicly available # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT. +import os import trimesh import pymeshlab def remesh_mesh(mesh_path, remesh_path): - mesh = mesh_simplify_trimesh(mesh_path, remesh_path) + # MODELBEAST: HY3D_REMESH_FACES env overrides the decimation target so the + # bake-mesh budget is tunable per node (Studio GPUs handle far more than the + # laptop-tuned 40k default). Additive — unset keeps upstream behavior. + target = int(os.environ.get("HY3D_REMESH_FACES", "40000")) + mesh = mesh_simplify_trimesh(mesh_path, remesh_path, target_count=target) def mesh_simplify_trimesh(inputpath, outputpath, target_count=40000):