| """
|
| Mesh generator: image → 3D mesh via TripoSR (local).
|
| Expects TripoSR repo cloned at project_root/TripoSR. Uses subprocess to run run.py.
|
| For text→mesh: first generate image with text_to_image(), then call this.
|
| """
|
|
|
| import os
|
| import subprocess
|
| import sys
|
| import time
|
| from pathlib import Path
|
|
|
|
|
| def find_triposr_root(project_root: str | None = None) -> str | None:
|
| """Locate TripoSR repo: ./TripoSR or ../TripoSR from script dir."""
|
| if project_root is None:
|
| project_root = str(Path(__file__).resolve().parent.parent)
|
| candidates = [
|
| os.path.join(project_root, "TripoSR"),
|
| os.path.join(project_root, "..", "TripoSR"),
|
| ]
|
| for p in candidates:
|
| run_py = os.path.join(p, "run.py")
|
| if os.path.isfile(run_py):
|
| return p
|
| return None
|
|
|
|
|
| def generate_mesh_from_image(
|
| image_path: str,
|
| output_dir: str = "outputs",
|
| mesh_format: str = "glb",
|
| triposr_root: str | None = None,
|
| device: str = "cuda:0",
|
| ) -> tuple[str | None, float, str]:
|
| """
|
| Run TripoSR on an image. Returns (path_to_mesh, inference_time_sec, message).
|
| If TripoSR is not found, returns (None, 0, error_message).
|
| """
|
| project_root = str(Path(__file__).resolve().parent.parent)
|
| triposr_root = triposr_root or find_triposr_root(project_root)
|
| if not triposr_root:
|
| return (
|
| None,
|
| 0.0,
|
| "TripoSR not found. Clone it: git clone https://github.com/VAST-AI-Research/TripoSR.git",
|
| )
|
|
|
| Path(output_dir).mkdir(parents=True, exist_ok=True)
|
|
|
| run_py = os.path.join(triposr_root, "run.py")
|
| cmd = [
|
| sys.executable,
|
| run_py,
|
| image_path,
|
| "--output-dir",
|
| output_dir,
|
| "--model-save-format",
|
| mesh_format,
|
| "--device",
|
| device,
|
| ]
|
|
|
| t0 = time.perf_counter()
|
| try:
|
| result = subprocess.run(
|
| cmd,
|
| cwd=triposr_root,
|
| capture_output=True,
|
| text=True,
|
| timeout=120,
|
| )
|
| t1 = time.perf_counter()
|
| if result.returncode != 0:
|
| return (
|
| None,
|
| t1 - t0,
|
| f"TripoSR failed: {result.stderr or result.stdout or 'unknown'}",
|
| )
|
|
|
| mesh_path = os.path.join(output_dir, "0", f"mesh.{mesh_format}")
|
| if not os.path.isfile(mesh_path):
|
| return (None, t1 - t0, f"TripoSR did not produce {mesh_path}")
|
| return (os.path.abspath(mesh_path), t1 - t0, "OK")
|
| except subprocess.TimeoutExpired:
|
| t1 = time.perf_counter()
|
| return (None, t1 - t0, "TripoSR timed out (120s)")
|
| except Exception as e:
|
| t1 = time.perf_counter()
|
| return (None, t1 - t0, str(e))
|
|
|
|
|
| def generate_mesh_from_text(
|
| prompt: str,
|
| output_dir: str = "outputs",
|
| mesh_format: str = "glb",
|
| seed: int | None = None,
|
| ) -> tuple[str | None, float, str]:
|
| """
|
| Text → image (SD) → mesh (TripoSR). Returns (path_to_mesh, total_time_sec, message).
|
| """
|
|
|
| _root = str(Path(__file__).resolve().parent.parent)
|
| if _root not in sys.path:
|
| sys.path.insert(0, _root)
|
| from scripts.text_to_image import text_to_image
|
|
|
| Path(output_dir).mkdir(parents=True, exist_ok=True)
|
| t0 = time.perf_counter()
|
| try:
|
| image_path, _ = text_to_image(prompt, output_dir=output_dir, seed=seed)
|
| except Exception as e:
|
| return (None, 0.0, f"Text-to-image failed: {e}")
|
|
|
| mesh_path, mesh_time, msg = generate_mesh_from_image(
|
| image_path,
|
| output_dir=os.path.join(output_dir, "mesh_run"),
|
| mesh_format=mesh_format,
|
| )
|
| total_time = time.perf_counter() - t0
|
| if mesh_path:
|
| return (mesh_path, total_time, msg)
|
| return (None, total_time, msg)
|
|
|
