modelbeast/server/operators/trellis2cpp/run.py

"""TRELLIS.2 via the C++/ggml port, on Metal. No Python, no venv.

Why this lane exists alongside trellis2_mlx: that one runs a Python venv per node, and has
repeatedly failed when the farm routed a job to a node whose venv was missing. This is a
self-contained binary + GGUF files, so "operator registered, venv absent" cannot happen — if
the install dir is there it runs, and it says so loudly if it is not.

The fine/cascade path is only reachable through the bundled Go server (the CLI examples cover
the coarse path plus mesh utilities), so we keep ONE resident server per node and POST to it.
Resident means the ~15 GB of GGUFs stay cached between jobs; -unload-idle frees them when the
node goes quiet, so a shared box is not permanently down 15 GB.

Measured on the same image (trellis2-bench/anatomy1.jpeg), 512 fine:
  m3ultra 42.3s · m2max 99.5s · m4probook 118.0s   (peak RSS 9-12 GB)
  1024 cascade: m3ultra 108.6s · m4probook 300.7s
vs the trellis2_mlx lane's 156s for 392k tris — this is ~3.7x faster and produces more geometry.
"""
import argparse
import json
import mimetypes
import os
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
PORT = int(os.environ.get("TRELLIS2CPP_PORT", "8743"))
API = f"http://127.0.0.1:{PORT}"

# Install-dir resolution, same spirit as trellis2_mlx: vendored first, then the conventional
# per-box locations, so an existing manual install keeps working without re-vendoring.
CANDIDATES = [
    ROOT / "vendor" / "trellis2cpp",
    Path.home() / "trellis2cpp",
    Path.home() / "trellis2cpp-test" / "src",
]


def resolve_install():
    for c in CANDIDATES:
        lib = c / "build-shared" / "libtrellis2.dylib"
        srv = c / "server" / "trellis2-server"
        if lib.exists() and srv.exists() and (c / "ggufs").is_dir():
            return c
    return None


ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)
outdir = Path(a.outdir)

if not a.input:
    print("ERROR: no input image")
    sys.exit(1)

INSTALL = resolve_install()
if INSTALL is None:
    print("ERROR: trellis2cpp not installed on this node. Expected one of:")
    for c in CANDIDATES:
        print(f"  {c}/  (needs build-shared/libtrellis2.dylib, server/trellis2-server, ggufs/)")
    print("  Deploy with wardrobegod tools/deploy_trellis2cpp.sh <user@host>")
    sys.exit(1)

quality = str(p.get("quality", "1024"))
if quality not in ("coarse", "512", "1024", "1536"):
    print(f"ERROR: quality must be coarse|512|1024|1536, got {quality!r}")
    sys.exit(1)


def alive():
    try:
        urllib.request.urlopen(f"{API}/api/info", timeout=5)
        return True
    except Exception:
        return False


def ensure_server():
    if alive():
        print("trellis2cpp: server already resident (models stay cached)", flush=True)
        return
    print(f"trellis2cpp: starting resident server on :{PORT} ...", flush=True)
    log = open("/tmp/trellis2cpp-server.log", "ab")
    subprocess.Popen(
        [str(INSTALL / "server" / "trellis2-server"),
         "-lib", str(INSTALL / "build-shared" / "libtrellis2.dylib"),
         "-ggufs", str(INSTALL / "ggufs"),
         "-store", str(Path.home() / "trellis2cpp" / "generations"),
         "-unload-idle",                       # free ~15 GB when the node goes quiet
         "-addr", f"127.0.0.1:{PORT}"],
        cwd=str(INSTALL / "server"), stdout=log, stderr=log, start_new_session=True)
    for _ in range(120):                       # model load is slow on a cold start
        if alive():
            print("trellis2cpp: server up", flush=True)
            return
        time.sleep(2)
    print("ERROR: server did not come up — see /tmp/trellis2cpp-server.log")
    sys.exit(1)


def post_image(path):
    boundary = uuid.uuid4().hex
    ctype = mimetypes.guess_type(path)[0] or "application/octet-stream"
    body = (
        f"--{boundary}\r\nContent-Disposition: form-data; name=\"quality\"\r\n\r\n{quality}\r\n"
        f"--{boundary}\r\nContent-Disposition: form-data; name=\"image\"; "
        f"filename=\"{os.path.basename(path)}\"\r\nContent-Type: {ctype}\r\n\r\n"
    ).encode() + open(path, "rb").read() + f"\r\n--{boundary}--\r\n".encode()
    req = urllib.request.Request(
        f"{API}/api/generate", data=body,
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"})
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.load(r)["job"]


ensure_server()
src = str(Path(a.input[0]).resolve())
print(f"+ trellis2cpp quality={quality} <- {src}", flush=True)
job = post_image(src)
print(f"trellis2cpp: job {job}", flush=True)

t0 = time.time()
last = ""
state, info = "running", {}
while time.time() - t0 < 3600:
    time.sleep(5)
    try:
        with urllib.request.urlopen(f"{API}/api/job/{job}", timeout=30) as r:
            info = json.load(r)
    except Exception:
        continue
    state = info.get("state", "?")
    stages = info.get("stageTimings") or []
    if stages and stages[-1]["stage"] != last:
        last = stages[-1]["stage"]
        print(f"  ... {last}", flush=True)
    if state in ("done", "error", "failed", "cancelled"):
        break

if state != "done":
    print(f"ERROR: job {job} ended {state}: {str(info.get('error'))[:300]}")
    sys.exit(1)

outdir.mkdir(parents=True, exist_ok=True)
stem = Path(a.input[0]).stem
dest = outdir / f"trellis2cpp_{stem}_{quality}.glb"
with urllib.request.urlopen(f"{API}/api/glb/{job}", timeout=1800) as r:
    data = r.read()
if len(data) < 1024 or data[:4] != b"glTF":
    print(f"ERROR: server returned {len(data)} bytes that are not a GLB")
    sys.exit(1)
dest.write_bytes(data)

secs = (info.get("durationMs") or 0) / 1000.0
print(f"trellis2cpp: {secs:.1f}s on {info.get('device')} -> {dest} ({len(data)/1e6:.1f} MB)",
      flush=True)
outputs = [{"path": str(dest), "name": dest.name,
            "meta": {"tool": "trellis2cpp", "quality": quality,
                     "device": info.get("device"), "seconds": round(secs, 1),
                     "stages": {t["stage"]: t["milliseconds"] for t in (info.get("stageTimings") or [])}}}]
(outdir / "result.json").write_text(json.dumps({"outputs": outputs}))
print("done:", dest, flush=True)