|
|
"""
|
|
|
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)
|
|
|
|
