TRELLIS.2 image-to-3D (H100)

cd examples/trellis2-image-to-3d
super create --sku gpu_1x_h100        # → s-xxxx (keep host while building)
super run . --script run.sh --on s-xxxx              # default: the repo's example image
# your own image (upload it first, then pass its workspace path):
super upload s-xxxx ./my.png my.png
super run . --script run.sh --on s-xxxx -- my.png

super download s-xxxx /workspace/out/sample.glb ./sample.glb
super stop s-xxxx

#!/usr/bin/env python3
"""Image-to-3D inference with microsoft/TRELLIS.2-4B.

Runs INSIDE a supercomputer session after run.sh has built the extensions (see
README.md). Loads a single image and writes a textured GLB with PBR materials.
Adapted from the TRELLIS.2 README inference snippet.
"""
from __future__ import annotations

import argparse
import glob
import os
import shutil
import time

os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import cv2
import torch
import transformers
from PIL import Image

# transformers>=5 sets `all_tied_weights_keys` on every model during __init__, but
# RMBG-2.0's trust_remote_code BiRefNet (4.x-era) skips it, raising AttributeError
# when the pipeline loads its background remover. Provide a class-level default
# (empty = no tied weights) so older remote models load under transformers 5.x.
if not isinstance(getattr(transformers.PreTrainedModel, "all_tied_weights_keys", None), dict):
    transformers.PreTrainedModel.all_tied_weights_keys = {}

from trellis2.pipelines import Trellis2ImageTo3DPipeline
from trellis2.renderers import EnvMap
import o_voxel

MODEL = "microsoft/TRELLIS.2-4B"
# Soak knob (stress tier): keep converting images until this many seconds have
# elapsed; the pass-signal --output is written only at the end. 0 = one image.
STRESS_DURATION_S = int(os.environ.get("SUPERCOMPUTER_STRESS_DURATION_S", "0") or "0")


def generate(pipeline, image_path: str, out_path: str) -> None:
    """One full image->3D round: run the pipeline and export a textured GLB."""
    mesh = pipeline.run(Image.open(image_path))[0]
    mesh.simplify(16777216)
    glb = o_voxel.postprocess.to_glb(
        vertices=mesh.vertices, faces=mesh.faces,
        attr_volume=mesh.attrs, coords=mesh.coords,
        attr_layout=mesh.layout, voxel_size=mesh.voxel_size,
        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
        decimation_target=1000000, texture_size=4096,
        remesh=True, verbose=True,
    )
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    # PNG textures, not WebP: trimesh's WebP texture export calls
    # PIL._webp.HAVE_WEBPANIM, which Pillow 12 removed. PNG always works.
    glb.export(out_path)
    print("wrote", out_path, flush=True)


def main() -> None:
    ap = argparse.ArgumentParser(description="TRELLIS.2 image-to-3D")
    ap.add_argument("--image", default="assets/example_image/T.png")
    ap.add_argument("--hdri", default="assets/hdri/forest.exr")
    ap.add_argument("--output", default="/workspace/out/sample.glb")
    args = ap.parse_args()

    # Lighting environment for PBR material baking.
    envmap = EnvMap(torch.tensor(
        cv2.cvtColor(cv2.imread(args.hdri, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
        dtype=torch.float32, device="cuda",
    ))

    pipeline = Trellis2ImageTo3DPipeline.from_pretrained(MODEL)
    pipeline.cuda()

    if not STRESS_DURATION_S:
        generate(pipeline, args.image, args.output)
        return

    # Soak: keep converting the repo's example images with the loaded pipeline
    # until the window closes. Every round overwrites the same scratch GLB
    # (bounded disk); the pass-signal args.output is the final round's copy.
    deadline = time.monotonic() + STRESS_DURATION_S
    images = sorted(glob.glob("assets/example_image/*.png")) or [args.image]
    scratch = os.path.join(os.path.dirname(args.output), "soak.glb")
    rounds = 0
    while True:
        generate(pipeline, images[rounds % len(images)], scratch)
        rounds += 1
        if time.monotonic() >= deadline:
            break
    shutil.copy(scratch, args.output)
    print(f"soak: {rounds} generations; wrote {args.output}", flush=True)


if __name__ == "__main__":
    main()

#!/usr/bin/env bash
# Build microsoft/TRELLIS.2 and run image-to-3D inference INSIDE a supercomputer session.
# Run this in the session, not on your laptop (see README.md).
#
# First run clones the repo and compiles its CUDA extensions (flash-attn,
# nvdiffrast, cumesh, o_voxel, flexgemm) — slow, but it lands under /workspace
# (mounted volume) so `super stop` and resume keep it. Weights cache there too.
set -euo pipefail

REPO="${REPO:-/workspace/TRELLIS.2}"
IMAGE_IN="${1:-assets/example_image/T.png}"   # a repo asset by default
OUT="${OUT:-/workspace/out/sample.glb}"
EIGEN_MIRROR="${EIGEN_MIRROR:-/workspace/.cache/eigen.git}"

# From some egress ranges github.com serves git's ref advertisement (a GET)
# normally, then answers the POST /git-upload-pack that actually fetches the
# pack with a spurious `401 WWW-Authenticate: Basic realm="GitHub"` — and only
# over HTTP/2. Git reads the 401 as a credential challenge and asks for a
# username; `super run` runs this script under a PTY with no writer on the other
# end (the CLI never forwards stdin), so that read blocks forever: the
# 2026-09-02 prod smoke sat on `Username for 'https://github.com': ` until its
# whole 3600s budget expired, holding an H100 the entire time. The identical
# clone succeeds over HTTP/1.1, so pin the version. Global, because submodules,
# setup.sh's own recursive clones and pip's `git+https://` installs each shell
# out to their own git.
git config --global http.version HTTP/1.1
# Defence in depth: should a credential prompt ever be raised anyway, fail
# loudly instead of hanging the budget.
export GIT_TERMINAL_PROMPT=0

# Run a git step up to 3 times with backoff, then fail with `why`. This covers
# genuinely transient trouble — a dropped connection, a clone interrupted
# mid-transfer, gitlab's per-IP 403. It is deliberately NOT what rescues the
# github clone: that 401 is deterministic (measured 6/6 refusals over 60s on an
# affected host), so only the HTTP/1.1 pin above clears it and a retry would
# just spend the budget more slowly.
retry_git() {
    local why="$1"; shift
    local backoff
    for backoff in 0 30 60; do
        [ "$backoff" -eq 0 ] || sleep "$backoff"
        if "$@"; then
            return 0
        fi
    done
    echo "$why" >&2
    return 1
}

mirror_eigen() {
    rm -rf "$EIGEN_MIRROR"   # a 403 mid-clone still leaves the directory
    git clone --mirror https://gitlab.com/libeigen/eigen.git "$EIGEN_MIRROR"
}

clone_trellis() {
    rm -rf "$REPO"   # a refused or interrupted clone still leaves the directory
    git clone -b main https://github.com/microsoft/TRELLIS.2.git --recursive "$REPO"
}

# Eigen is the one dependency in this build that does not come from GitHub: it is
# a gitlab.com submodule in TWO places — TRELLIS's own o-voxel, and cubvh, nested
# under the CuMesh extension that setup.sh clones `--recursive`. gitlab.com serves
# unauthenticated clones from some cloud IP ranges an HTTP 403, which aborts the
# clone (and with it the run) wherever the session happens to land.
#
# So touch gitlab at most once per volume: mirror Eigen onto the cache volume and
# rewrite both submodule URLs at that mirror. Reruns on a warm volume never reach
# gitlab at all, and the pins still resolve to the exact upstream commits.
#
# Validate the mirror, not just the dir: this volume outlives the run, and the
# 2026-08-23 wedge restored mirrors that were saved HEALTHY on the preceding
# passing runs (HEAD, packed refs, ~132MB packfiles) yet came back unusable —
# volume restoration can corrupt or drop objects, not just interrupted clones.
# `rev-parse --verify HEAD` is not enough (it passes with the HEAD commit
# object deleted), and `cat-file -e` alone is not either (it passes with the
# commit's tree deleted): require the commit AND a connected object graph.
# fsck --connectivity-only on this packfile is seconds against the 3600s budget.
if ! git --git-dir="$EIGEN_MIRROR" cat-file -e 'HEAD^{commit}' 2>/dev/null \
        || ! git --git-dir="$EIGEN_MIRROR" fsck --connectivity-only --no-dangling \
            >/dev/null 2>&1; then
    retry_git "gitlab.com refused the Eigen clone 3x — this host's IP looks blocked" \
        mirror_eigen
else
    # A warm mirror can predate an upstream pin bump; refresh it when we can, but
    # never fail on it — a stale mirror still serves every pin fetched so far.
    git -C "$EIGEN_MIRROR" fetch -q --prune || true
fi

# insteadOf is a longest-prefix rewrite, so both spellings are needed: with only
# the bare form, ".../eigen.git" would rewrite to "$EIGEN_MIRROR.git". --unset-all
# first so a rerun replaces the entries instead of appending duplicates.
# protocol.file.allow: git refuses submodule clones over the local-path transport
# by default (CVE-2022-39253); the path here is one we just created ourselves.
git config --global protocol.file.allow always
git config --global --unset-all url."$EIGEN_MIRROR".insteadOf || true
git config --global --add url."$EIGEN_MIRROR".insteadOf https://gitlab.com/libeigen/eigen.git
git config --global --add url."$EIGEN_MIRROR".insteadOf https://gitlab.com/libeigen/eigen

# Validate the clone, not just the dir: a run killed mid-clone (preemption)
# leaves a half-fetched tree that satisfies -d but breaks every rerun.
if ! git -C "$REPO" rev-parse --git-dir >/dev/null 2>&1 \
        || [ ! -f "$REPO/trellis2/modules/image_feature_extractor.py" ]; then
    retry_git "github.com refused the TRELLIS.2 clone 3x — check this host's egress to github.com" \
        clone_trellis
fi
# Submodules can be half-fetched even when the top-level tree validates
# (preemption mid `--recursive`); this repairs them and no-ops when complete.
# Retried too: it is the same forge over the same egress as the clone above.
retry_git "github.com refused the TRELLIS.2 submodules 3x — check this host's egress to github.com" \
    git -C "$REPO" submodule update --init --recursive
cd "$REPO"

# transformers refactored DINOv3ViTModel to nest its transformer blocks under a
# `.model` submodule; TRELLIS's feature extractor still reads `.layer` directly.
# Point it at the nested path (idempotent — the pattern won't re-match once fixed).
sed -i 's/enumerate(self\.model\.layer)/enumerate(self.model.model.layer)/' \
    trellis2/modules/image_feature_extractor.py

# Pin torch to the base image's CUDA 12.4. TRELLIS's app requirements otherwise
# pull a cu13 torch, which mismatches the system nvcc (12.4) and aborts
# flash-attn's source build with a CUDA version-mismatch error.
# Also pin opencv < 5: OpenCV 5.0.0 wheels no longer read EXR images even with
# OPENCV_IO_ENABLE_OPENEXR=1, so imread(assets/hdri/*.exr) returns empty and
# cvtColor asserts. 4.x reads them fine.
printf 'torch==2.6.0\ntorchvision==0.21.0\nopencv-python<5\nopencv-python-headless<5\n' \
    > "$REPO/torch-constraints.txt"
export PIP_CONSTRAINT="$REPO/torch-constraints.txt"

# nvdiffrast/nvdiffrec link against the CUDA driver (-lcuda). The real libcuda
# is injected by the host at runtime, but the build-time linker needs the
# toolkit's stub, which lives in lib64/stubs (not on the default search path).
export LIBRARY_PATH="/usr/local/cuda/lib64/stubs:${LIBRARY_PATH:-}"

# Build the extensions into the image's system python (--basic, no new env).
# flash_attn is the build marker — skip the (slow) compile only when it's already
# present (a file marker would wrongly skip after a container recycle, since the
# compiled packages live in the container, not the persisted /workspace).
if ! python -c "import flash_attn" 2>/dev/null; then
    . ./setup.sh --basic --flash-attn --nvdiffrast --nvdiffrec --cumesh --o-voxel --flexgemm
fi

# Python deps setup.sh doesn't install usably, fixed up after the build:
#   plyfile  — imported by the o_voxel extension (not pulled in)
#   utils3d  — setup builds it from git as an unimportable "UNKNOWN" wheel
pip install -q plyfile utils3d

# nvdiffrast also builds as a nameless "UNKNOWN" (its pyproject derives the
# version from git tags, which the shallow clone lacks), so its __init__'s
# importlib.metadata.version('nvdiffrast') raises at import. The module files are
# fine — drop them on the path and synthesize minimal metadata. It JIT-compiles
# its CUDA kernels at first use (nvcc is in the image), so no build step here.
python -c "import nvdiffrast.torch" 2>/dev/null || python - <<'PY'
import site, shutil, pathlib
sp = pathlib.Path(site.getsitepackages()[0])
src = pathlib.Path("/tmp/extensions/nvdiffrast/nvdiffrast")
if src.exists() and not (sp / "nvdiffrast").exists():
    shutil.copytree(src, sp / "nvdiffrast")
info = sp / "nvdiffrast-0.3.3.dist-info"
info.mkdir(exist_ok=True)
(info / "METADATA").write_text("Metadata-Version: 2.1\nName: nvdiffrast\nVersion: 0.3.3\n")
PY

# TRELLIS's pipeline pulls two GATED models (facebook/dinov3-vitl16 and
# briaai/RMBG-2.0). Pre-download the pipeline so we can repoint those at ungated
# community mirrors (identical weights) in its cached config before it loads
# them — avoids having to accept the HF gates. The repo id also appears in the
# repo's own configs (dinov3), so patch both.
python -c "from huggingface_hub import snapshot_download; snapshot_download('microsoft/TRELLIS.2-4B')"
_hub="${HF_HOME:-$HOME/.cache/huggingface}/hub/models--microsoft--TRELLIS.2-4B"
for tgt in "$_hub" "$REPO/configs"; do
    # `|| true`: grep exits 1 when a repo id isn't in a given tree (e.g. RMBG is
    # only in the model cache, not the repo configs) — without it, set -o
    # pipefail + set -e would abort the whole run before inference.
    grep -rl "facebook/dinov3-vitl16-pretrain-lvd1689m" "$tgt" 2>/dev/null \
        | xargs -r sed -i 's#facebook/dinov3-vitl16-pretrain-lvd1689m#camenduru/dinov3-vitl16-pretrain-lvd1689m#g' || true
    grep -rl "briaai/RMBG-2.0" "$tgt" 2>/dev/null \
        | xargs -r sed -i 's#briaai/RMBG-2.0#camenduru/RMBG-2.0#g' || true
done

# infer.py is uploaded alongside this script (in the same dir); run it from the
# repo root so the default `assets/...` paths resolve.
cp "$(dirname "$0")/infer.py" "$REPO/infer.py"
mkdir -p "$(dirname "$OUT")"
python infer.py --image "$IMAGE_IN" --output "$OUT"
echo "wrote $OUT"

{
    "sku": "gpu_1x_h100",
    "image": {
        "base": "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel",
        "apt": ["git", "wget", "ca-certificates", "build-essential", "ninja-build", "libgl1", "libglib2.0-0", "sudo", "libjpeg-dev", "python3-dev", "libgl1-mesa-dev", "libegl1-mesa-dev"],
        "env": {
            "PYTHONUNBUFFERED": "1",
            "HF_HOME": "/workspace/.cache/huggingface",
            "HF_HUB_CACHE": "/workspace/.cache/huggingface/hub",
            "OPENCV_IO_ENABLE_OPENEXR": "1",
            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
            "SSL_CERT_FILE": "/etc/ssl/certs/ca-certificates.crt",
            "PIP_CERT": "/etc/ssl/certs/ca-certificates.crt",
            "REQUESTS_CA_BUNDLE": "/etc/ssl/certs/ca-certificates.crt"
        }
    },
    "volumes": [
        {"name": "trellis2-cache", "mount": "/workspace/.cache"}
    ]
}