scripts/trellis_generate.py

#!/usr/bin/env python3
"""CLI wrapper around microsoft/TRELLIS.2's image-to-3D pipeline.

Turns a single RGBA (or RGB) image crop into a textured GLB mesh. Meant to be
invoked as a subprocess -- from the `fpgm` env or anywhere else -- using the
`trellis2` conda env's interpreter, since that env holds TRELLIS.2's own
torch build (2.6.0+cu124) which must stay isolated from the rest of this
project's torch (2.10.0+cu128 in `fpgm`).

Usage:
    source scripts/nvidia_lib_shim.sh   # required: fixes NVML driver mismatch
    /home/quang/miniconda3/envs/trellis2/bin/python scripts/trellis_generate.py \\
        --image crop.png --out mesh.glb [--resolution 512] [--seed 0] [--simplify 1000000]

On success prints the absolute output path to stdout and exits 0. A sibling
JSON file (<out>.json) is always written on success, recording what the
calling process needs to know about what it got. On failure, prints a clear
message to stderr and exits non-zero.
"""
import os
import sys

# Must happen before torch (or anything that transitively imports it) is
# imported -- per the TRELLIS.2 model card.
os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_TRELLIS_DIR = os.path.join(_REPO_ROOT, "third_party", "TRELLIS.2")
_CONDA_ROOT = "/home/quang/miniconda3"

# TRELLIS.2's repo root has no pyproject.toml -- `trellis2`/`o_voxel`'s
# runtime companion modules are meant to be run with the repo on the path
# (exactly how example.py in the upstream repo works), not pip-installed.
if _TRELLIS_DIR not in sys.path:
    sys.path.insert(0, _TRELLIS_DIR)

# flex_gemm (a TRELLIS.2 CUDA extension) JIT-compiles Triton kernels the
# first time it's used, and Triton's build path (triton/runtime/build.py)
# looks for a C compiler via $CC or `gcc`/`clang` on $PATH -- it does not
# reuse whatever compiled flex_gemm itself at install time. Point it at the
# same buildtools gcc used by scripts/setup_trellis_env.sh.
os.environ.setdefault("CC", os.path.join(_CONDA_ROOT, "envs", "buildtools", "bin", "gcc"))
os.environ.setdefault("CXX", os.path.join(_CONDA_ROOT, "envs", "buildtools", "bin", "g++"))
os.environ["PATH"] = os.path.join(_CONDA_ROOT, "envs", "buildtools", "bin") + os.pathsep + os.environ.get("PATH", "")
os.environ.setdefault("CUDA_HOME", os.path.join(_CONDA_ROOT, "envs", "cuda128"))

# Attention backend: whichever one scripts/setup_trellis_env.sh managed to
# install is recorded in .trellis_attn_backend; respect an explicit
# SPARSE_ATTN_BACKEND/ATTN_BACKEND from the caller's environment first.
_attn_backend_file = os.path.join(_REPO_ROOT, ".trellis_attn_backend")
if "SPARSE_ATTN_BACKEND" not in os.environ and "ATTN_BACKEND" not in os.environ and os.path.exists(_attn_backend_file):
    with open(_attn_backend_file) as _f:
        _backend = _f.read().strip()
    if _backend:
        os.environ["SPARSE_ATTN_BACKEND"] = _backend
        os.environ["ATTN_BACKEND"] = _backend

import argparse
import json
import subprocess
import time
import traceback
from pathlib import Path

PIPELINE_TYPES = ("512", "1024", "1024_cascade", "1536_cascade")


def pick_gpu() -> int:
    """Return the index of the GPU with the most free memory right now.

    The 4x H200s on this host are shared and free VRAM fluctuates 11-44GB,
    so the choice is made fresh on every invocation rather than hardcoded.
    If the caller already pinned CUDA_VISIBLE_DEVICES, that choice is
    respected instead.
    """
    if "CUDA_VISIBLE_DEVICES" in os.environ:
        return 0
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
            text=True,
            timeout=15,
        )
        free = [int(x) for x in out.strip().splitlines() if x.strip()]
        if not free:
            raise ValueError("nvidia-smi returned no GPUs")
        return max(range(len(free)), key=lambda i: free[i])
    except Exception as e:
        print(
            f"WARNING: could not query nvidia-smi to pick a GPU ({e}); "
            "defaulting to device 0. If this fails with an NVML error, "
            "did you `source scripts/nvidia_lib_shim.sh`?",
            file=sys.stderr,
        )
        return 0


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Turn an RGBA image crop into a textured GLB mesh with TRELLIS.2-4B.",
    )
    p.add_argument("--image", required=True, type=Path, help="Path to an RGBA/RGB PNG image.")
    p.add_argument("--out", required=True, type=Path, help="Output .glb path.")
    p.add_argument(
        "--resolution",
        default="512",
        choices=PIPELINE_TYPES,
        help="TRELLIS.2 pipeline_type (shape/texture SLat resolution). "
        "Default 512, the lowest-VRAM option, since GPUs here are shared/contended.",
    )
    p.add_argument("--seed", type=int, default=0, help="Random seed.")
    p.add_argument(
        "--simplify",
        type=int,
        default=50_000,
        help="Target vertex count for final mesh decimation (o_voxel.postprocess.to_glb's decimation_target). "
        "The model card example uses 1_000_000, but xatlas UV-chart computation over that many vertices "
        "was observed to exhaust host RAM (silent SIGKILL, no traceback) on this shared machine for "
        "irregular/noisy silhouettes -- 50_000 is a much safer default; raise it explicitly if you have "
        "headroom and need finer geometry.",
    )
    p.add_argument(
        "--texture-size",
        type=int,
        default=1024,
        help="Baked texture resolution (px). Model card default is 4096; lowered further here (from an "
        "earlier 2048 default) alongside --simplify to avoid the same OOM risk during UV baking.",
    )
    return p.parse_args()


def main() -> None:
    args = parse_args()

    if not args.image.exists():
        print(f"ERROR: input image not found: {args.image}", file=sys.stderr)
        sys.exit(1)

    gpu_idx = pick_gpu()
    os.environ.setdefault("CUDA_VISIBLE_DEVICES", str(gpu_idx))

    t0 = time.time()

    try:
        import torch
        from PIL import Image
        import o_voxel
        from trellis2.pipelines import Trellis2ImageTo3DPipeline
        import trellis2.pipelines.rembg as _rembg_module
    except Exception as e:
        print(f"ERROR: failed to import the TRELLIS.2 stack: {e}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(1)

    # Trellis2ImageTo3DPipeline.from_pretrained() unconditionally instantiates
    # the background-removal model (BiRefNet, backed by the gated
    # briaai/RMBG-2.0 checkpoint on HF) even when it will never be called --
    # it downloads it at pipeline-construction time, before any image is even
    # looked at. This CLI's contract is that --image is already a proper RGBA
    # crop (see preprocess_image: it only invokes rembg when the input has no
    # real alpha channel), so that model is genuinely dead weight here -- and
    # on this machine's HF account it isn't even gate-approved, which would
    # otherwise hard-fail every run. Swap in a no-op before construction
    # rather than patching TRELLIS.2's own source.
    class _NoRembg:
        def __init__(self, *_a, **_kw):
            pass

        def to(self, _device):
            pass

        def cuda(self):
            pass

        def cpu(self):
            pass

        def __call__(self, _image):
            raise RuntimeError(
                "This input image has no real alpha channel (preprocess_image "
                "wants to run background removal), but this wrapper disables "
                "TRELLIS.2's rembg model (gated briaai/RMBG-2.0 checkpoint, "
                "not needed for a pre-masked crop). Pass an --image that "
                "already carries a real (non-uniform) alpha channel."
            )

    _rembg_module.BiRefNet = _NoRembg

    # TRELLIS.2's DINOv3 image-conditioning feature extractor
    # (trellis2/modules/image_feature_extractor.py:
    # DinoV3FeatureExtractor.extract_features) reaches into
    # `self.model.layer`, which matched the `transformers` internals when
    # TRELLIS.2 was written. In the `transformers` version installed here
    # (5.14.1, current at install time -- no version was pinned upstream),
    # `DINOv3ViTModel` wraps its encoder in a `.model` submodule
    # (`DINOv3ViTEncoder`), so the real path is `self.model.model.layer`;
    # `self.model.layer` raises AttributeError. Rebind the method with the
    # corrected attribute path rather than patching TRELLIS.2's source or
    # pinning an old, unmaintained transformers release.
    import trellis2.modules.image_feature_extractor as _ife_module

    def _dinov3_extract_features_fixed(self, image: "torch.Tensor") -> "torch.Tensor":
        import torch.nn.functional as F
        image = image.to(self.model.embeddings.patch_embeddings.weight.dtype)
        hidden_states = self.model.embeddings(image, bool_masked_pos=None)
        position_embeddings = self.model.rope_embeddings(image)
        for layer_module in self.model.model.layer:  # was self.model.layer
            hidden_states = layer_module(hidden_states, position_embeddings=position_embeddings)
        return F.layer_norm(hidden_states, hidden_states.shape[-1:])

    _ife_module.DinoV3FeatureExtractor.extract_features = _dinov3_extract_features_fixed

    if not torch.cuda.is_available():
        print(
            "ERROR: torch.cuda.is_available() is False. Did you "
            "`source scripts/nvidia_lib_shim.sh` before running this script?",
            file=sys.stderr,
        )
        sys.exit(1)

    torch.cuda.reset_peak_memory_stats()

    try:
        image = Image.open(args.image)
    except Exception as e:
        print(f"ERROR: could not open image {args.image}: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        pipeline = Trellis2ImageTo3DPipeline.from_pretrained("microsoft/TRELLIS.2-4B")
        if not pipeline.low_vram:
            print("WARNING: pipeline.low_vram is False (expected True by default).", file=sys.stderr)
        pipeline.cuda()

        meshes = pipeline.run(image, seed=args.seed, pipeline_type=args.resolution)
        mesh = meshes[0]
        mesh.simplify(16_777_216)  # nvdiffrast's index limit, per the model card example

        # NOTE: render_utils.render_video is intentionally never called here --
        # it's a preview convenience the calling pipeline doesn't need and is
        # pure extra GPU/time cost.
        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=args.simplify,
            texture_size=args.texture_size,
            remesh=True,
            remesh_band=1,
            remesh_project=0,
            verbose=True,
        )
        args.out.parent.mkdir(parents=True, exist_ok=True)
        glb.export(str(args.out), extension_webp=True)
    except Exception as e:
        print(f"ERROR: TRELLIS.2 generation failed: {e}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(1)

    elapsed = time.time() - t0
    peak_vram_bytes = torch.cuda.max_memory_allocated()

    info = {
        "input_image": str(args.image.resolve()),
        "output_glb": str(args.out.resolve()),
        "resolution": args.resolution,
        "seed": args.seed,
        "num_vertices": int(glb.vertices.shape[0]),
        "num_faces": int(glb.faces.shape[0]),
        "elapsed_seconds": round(elapsed, 2),
        "gpu_index": gpu_idx,
        "gpu_name": torch.cuda.get_device_name(0),
        "peak_vram_bytes": int(peak_vram_bytes),
        "peak_vram_gb": round(peak_vram_bytes / (1024**3), 2),
    }
    json_path = args.out.with_suffix(".json")
    json_path.write_text(json.dumps(info, indent=2))

    print(str(args.out.resolve()))
    sys.exit(0)


if __name__ == "__main__":
    main()