File size: 4,055 Bytes
74f0b48 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """
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)
# TripoSR writes to output_dir/0/mesh.obj (or mesh.glb)
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'}",
)
# Output is output_dir/0/mesh.glb or mesh.obj
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).
"""
# Import here so script can run from any cwd
_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)
|