113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Install ComfyUI-Trellis2Apple dependencies into the current Python env."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import platform
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
CONFIG_FILE = ROOT / "config.json"
|
|
|
|
PYPI_DEPS = [
|
|
"torch>=2.11.0",
|
|
"torchvision>=0.26.0",
|
|
"transformers>=4.40.0,<5",
|
|
"safetensors",
|
|
"huggingface_hub",
|
|
"pillow",
|
|
"numpy",
|
|
"trimesh",
|
|
"xatlas",
|
|
"fast-simplification",
|
|
"pygltflib",
|
|
"opencv-python-headless",
|
|
"mlx>=0.31.0",
|
|
"tqdm",
|
|
"imageio",
|
|
"imageio-ffmpeg",
|
|
"scipy",
|
|
"einops",
|
|
]
|
|
|
|
METAL_DEPS = [
|
|
"https://github.com/pedronaugusto/mtldiffrast/archive/main.tar.gz",
|
|
"https://github.com/pedronaugusto/mtlmesh/archive/main.tar.gz",
|
|
"https://github.com/pedronaugusto/mtlgemm/archive/main.tar.gz",
|
|
]
|
|
|
|
|
|
def run(cmd: list[str], env: dict[str, str] | None = None) -> None:
|
|
print("+", " ".join(cmd), flush=True)
|
|
subprocess.check_call(cmd, env=env)
|
|
|
|
|
|
def looks_like_trellis2_apple(path: Path) -> bool:
|
|
return (
|
|
(path / "mlx_backend" / "pipeline.py").is_file()
|
|
and (path / "trellis2" / "pipelines" / "trellis2_image_to_3d.py").is_file()
|
|
and (path / "o-voxel" / "setup.py").is_file()
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Install ComfyUI-Trellis2Apple dependencies into the current Python environment."
|
|
)
|
|
parser.add_argument("--trellis2-apple", required=True, help="Path to pedronaugusto/trellis2-apple.")
|
|
parser.add_argument("--skip-metal", action="store_true", help="Skip Metal packages.")
|
|
parser.add_argument(
|
|
"--download-weights",
|
|
action="store_true",
|
|
help="Download microsoft/TRELLIS.2-4B into trellis2-apple/weights/TRELLIS.2-4B.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
trellis2_apple = Path(args.trellis2_apple).expanduser().resolve()
|
|
if not looks_like_trellis2_apple(trellis2_apple):
|
|
raise SystemExit(f"{trellis2_apple} does not look like trellis2-apple.")
|
|
|
|
print(f"Python: {sys.executable}")
|
|
print(f"trellis2-apple: {trellis2_apple}")
|
|
if platform.machine() != "arm64":
|
|
print("WARNING: this node pack targets Apple Silicon (arm64).")
|
|
|
|
run([sys.executable, "-m", "pip", "install", *PYPI_DEPS])
|
|
run([sys.executable, "-m", "pip", "install", str(trellis2_apple / "o-voxel")])
|
|
|
|
if not args.skip_metal:
|
|
run([sys.executable, "-m", "pip", "install", "setuptools", "wheel", "pybind11"])
|
|
env = os.environ.copy()
|
|
env.setdefault("MACOSX_DEPLOYMENT_TARGET", "12.0")
|
|
run([sys.executable, "-m", "pip", "install", "--no-build-isolation", *METAL_DEPS], env=env)
|
|
|
|
if args.download_weights:
|
|
run(
|
|
[
|
|
sys.executable,
|
|
str(trellis2_apple / "scripts" / "download_weights.py"),
|
|
"--output-dir",
|
|
str(trellis2_apple / "weights" / "TRELLIS.2-4B"),
|
|
]
|
|
)
|
|
|
|
CONFIG_FILE.write_text(
|
|
json.dumps({"trellis2_apple_path": str(trellis2_apple)}, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
print("\nInstall complete.")
|
|
print(f"Wrote {CONFIG_FILE}")
|
|
print(f"Run: {sys.executable} {ROOT / 'doctor.py'}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|