scripts/sam3d_generate.py

#!/usr/bin/env python3
"""CLI wrapper around facebook/sam-3d-objects' single-image-to-3D pipeline.

Turns a single RGBA (or RGB) image crop into a textured GLB mesh, mirroring
scripts/trellis_generate.py's contract exactly so run_object_pipeline.py's
`mesh` stage can treat --mesh-source trellis and --mesh-source sam3d
identically (see run_mesh() in scripts/run_object_pipeline.py). Meant to be
invoked as a subprocess -- from the `fpgm` env or anywhere else -- using the
`sam3d-objects` conda env's interpreter, since that env holds SAM 3D
Objects' own torch build (2.5.1+cu121), kept isolated from `fpgm`'s torch
(2.10.0+cu128) and `trellis2`'s (2.6.0+cu124) exactly like TRELLIS.2 is.

Usage:
    source scripts/nvidia_lib_shim.sh   # required: fixes NVML driver mismatch
    /home/quang/miniconda3/envs/sam3d-objects/bin/python scripts/sam3d_generate.py \\
        --image crop.png --out mesh.glb [--seed 0] [--texture-size 1024] \\
        [--no-texture-baking] [--no-mesh-postprocess] [--simplify-ratio 0.95]

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 (mirrors
scripts/trellis_generate.py's convention). On failure, prints a clear message
to stderr and exits non-zero. The raw 3D Gaussian splat SAM 3D Objects
produces as a byproduct of reconstruction is also always saved next to the
mesh, as `<out-stem>_splat.ply` -- not the primary artifact this pipeline
consumes, but kept for inspection/debugging since it's nearly free to write.

--- Where the mesh actually comes from (this took real digging) --------------
SAM 3D Objects' own README/demo.py only shows `output["gs"].save_ply(...)`
i.e. the Gaussian splat -- which reads as if mesh export were unsupported.
It is NOT unsupported. `sam3d_objects.pipeline.inference_pipeline
.InferencePipeline` (the base class of the checkpoint's actual
`InferencePipelinePointMap`) decodes the structured latent to *both* a
Gaussian splat and a full triangle mesh by default
(`decode_formats=["gaussian", "mesh"]`; see `decode_slat()`:
``ret["mesh"] = self.models["slat_decoder_mesh"](slat)``, using the
checkpoint's own `slat_decoder_mesh.ckpt`/`slat_decoder_mesh.yaml`, which do
ship in the gated HF repo). `postprocess_slat_output()` then calls
`postprocessing_utils.to_glb(gaussian, mesh, ...)`, which bakes the splat's
appearance into a UV-textured (or vertex-colored) `trimesh.Trimesh` -- the
exact vertices+faces contract `fpgm.objects.types.ObjectMesh` /
`run_object_pipeline.py`'s `_save_mesh()` already expect from TRELLIS. No
Gaussian-splat-to-mesh conversion (Poisson reconstruction, marching cubes,
open3d, ...) is needed at all -- do not add one; use this path instead.

The catch, and why this script does not use the README's own quick-start
snippet: `notebook/inference.py`'s public `Inference.__call__` wrapper
hardcodes `with_mesh_postprocess=False, with_texture_baking=False` in its
call to `pipeline.run(...)`, which is why the README's own example only ever
shows `output["gs"]` -- that convenience wrapper throws the mesh path away.
This script bypasses that wrapper: it builds the pipeline the same way
(`hydra.utils.instantiate` over `checkpoints/hf/pipeline.yaml`, the exact
config `Inference.__init__` loads), calls `pipeline.run(...,
decode_formats=["mesh", "gaussian"], with_mesh_postprocess=False,
with_texture_baking=False)` itself just to get the raw decoded
`output["mesh"][0]`/`output["gaussian"][0]` cheaply, and then calls
`postprocessing_utils.to_glb(...)` a second time *directly* with this
script's own `--simplify-ratio`/`--texture-size` -- `InferencePipeline
.postprocess_slat_output()` (which is what would produce `output["glb"]` if
called through `pipeline.run()` with those flags True) hardcodes
`simplify=0.95, texture_size=1024` internally with no way to override them
through `run()`'s own arguments, so those two CLI flags would silently do
nothing if this script instead just read `output["glb"]`. It also
deliberately does NOT import `notebook/inference.py` at all: that module's
top-level `os.environ["CUDA_HOME"] = os.environ["CONDA_PREFIX"]` (a) raises
KeyError when this script is invoked via the env's interpreter directly
rather than through an activated `conda activate sam3d-objects` shell (no
CONDA_PREFIX in that case), and (b) would clobber the CUDA_HOME this script
sets deliberately below (this repo's `sam3d-objects` env is a minimal
python+pip env, not built from environments/default.yml's full CUDA-toolkit
conda env, so CONDA_PREFIX has no nvcc/headers under it).

--- TESTED end-to-end (2026-07-30) -------------------------------------------
Both the default config (mesh postprocess + texture baking on) and
--no-texture-baking were run for real against
third_party/sam-3d-objects/notebook/images/kid_box (mask 0.png merged into
alpha) on this host and produced valid, trimesh-loadable .glb output --
see logs/sam3d_setup/09_smoketest_generate.log (--no-texture-baking) and
logs/sam3d_setup/11_smoketest_textured.log (full default: 7105 vertices,
13564 faces, baked 1024x1024 PBR texture, ~123s, ~18GB peak VRAM on an
H200 NVL). Three env gaps had to be closed to get there, all now fixed in
the `sam3d-objects` conda env (see logs/sam3d_setup/ for each):
  - `sam3d_objects/__init__.py` unconditionally does `import sam3d_objects.init`,
    a submodule that does not exist in the public GitHub release (see the
    LIDRA_SKIP_INIT env var set below -- this script sets it automatically).
  - `postprocessing_utils.to_glb`'s hole-filling step (`with_mesh_postprocess=
    True`, the default) needs `nvdiffrast` for rasterization *regardless* of
    the `rendering_engine="pytorch3d"` override used for texture baking --
    installed into the env from the same `third_party/TRELLIS.2`-adjacent
    `.trellis_ext_build/nvdiffrast` clone TRELLIS.2 itself uses.
  - Texture baking's multi-view Gaussian render hardcodes the "inria"
    `diff_gaussian_rasterization` backend with a `kernel_size` argument that
    vanilla `graphdeco-inria/diff-gaussian-rasterization` does not accept
    (`TypeError: unexpected keyword argument 'kernel_size'`) -- the Mip-
    Splatting-style fork vendored at
    `autonomousvision/mip-splatting`'s `submodules/diff-gaussian-rasterization`
    (the same one `microsoft/TRELLIS`'s own setup.sh --mipgaussian flag
    installs) does have it and is what's installed in this env instead.

One real, non-obvious output-quality gotcha found during testing:
`--no-texture-baking` alone (leaving `with_mesh_postprocess=True`, the
default) produces a mesh with **no color at all** -- flat trimesh-default
gray. This is not a bug in this script: `postprocessing_utils.to_glb`'s
per-vertex-color path only fires when `with_mesh_postprocess` is *also*
False (`if not with_mesh_postprocess and not with_texture_baking and
use_vertex_color: ...`), because mesh postprocessing (decimation + hole-fill
+ mincut) changes the vertex count/topology and the original per-vertex
color array can no longer be validly indexed against it. To get a colored
mesh you must pick one of: (a) leave texture baking on (the default --
slower, ~2min, renders 100 views + a short UV-bake optimization, but colors
survive postprocessing correctly), or (b) pass both --no-texture-baking
*and* --no-mesh-postprocess (fast, but keeps the raw ~367k-vertex/735k-face
mesh un-decimated and with sharp/floating artifacts un-cleaned).

Not yet exercised: real DROID crops through the full run_object_pipeline.py
`mesh` stage (only the standalone repo's own sample image was tested here);
attention backend locks in as "sdpa" rather than "flash_attn" with this
script's current import order (harmless, just not maximally fast -- see
inline comment near the sam3d_objects import below).
"""
from __future__ import annotations

import os
import sys

# Must happen before torch (or anything that transitively imports it) is
# imported.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

# sam3d_objects/__init__.py unconditionally does `import sam3d_objects.init`
# unless this is set (see notebook/inference.py, which sets it for the same
# reason) -- that submodule does not exist in the public GitHub release (an
# internal-only module was stripped but the __init__.py guard against it
# wasn't cleaned up), so without this every `import sam3d_objects` raises
# `ModuleNotFoundError: No module named 'sam3d_objects.init'`. Confirmed by
# hand against this repo's actual checkout; not a guess.
os.environ.setdefault("LIDRA_SKIP_INIT", "true")

_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_SAM3D_DIR = os.path.join(_REPO_ROOT, "third_party", "sam-3d-objects")
_CONDA_ROOT = "/home/quang/miniconda3"

# sam3d_objects itself is `pip install -e`d into the sam3d-objects env (see
# doc/setup.md), so it does not need a sys.path insertion the way TRELLIS.2
# does -- but guard for a not-yet-installed package with a clear error rather
# than a confusing ImportError further down.
if _SAM3D_DIR not in sys.path:
    sys.path.insert(0, _SAM3D_DIR)

# Reuse the same buildtools/cuda128 conda envs TRELLIS.2 uses (see
# scripts/setup_trellis_env.sh) for anything that JIT-compiles CUDA/C++ at
# import or first-call time (xatlas, pymeshfix, gsplat's rasterizer, ...).
# Must be set before any such import, and deliberately NOT derived from
# CONDA_PREFIX (see the module docstring's UNTESTED section).
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"))

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


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

    Identical policy to scripts/trellis_generate.py's pick_gpu() -- see that
    docstring for why (shared, contended GPUs on this host).
    """
    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 facebook/sam-3d-objects.",
    )
    p.add_argument("--image", required=True, type=Path, help="Path to an RGBA/RGB PNG image (alpha = object mask, matching scripts/trellis_generate.py's convention).")
    p.add_argument("--out", required=True, type=Path, help="Output .glb path.")
    p.add_argument("--seed", type=int, default=0, help="Random seed.")
    p.add_argument(
        "--checkpoint-tag",
        default="hf",
        help="Checkpoint subdirectory under third_party/sam-3d-objects/checkpoints/ "
        "(the README's ${TAG} convention; 'hf' is what `hf download` + the "
        "mv-into-place step in doc/setup.md produces).",
    )
    p.add_argument(
        "--texture-size", type=int, default=1024,
        help="Baked texture resolution (px), passed through to postprocessing_utils.to_glb. Ignored with --no-texture-baking.",
    )
    p.add_argument(
        "--simplify-ratio", type=float, default=0.95,
        help="Fraction of triangles to REMOVE during mesh simplification (to_glb's own "
        "`simplify` parameter/semantics -- NOT a target vertex count like TRELLIS's "
        "--simplify). 0 disables simplification. Ignored with --no-mesh-postprocess.",
    )
    p.add_argument(
        "--no-mesh-postprocess", action="store_true",
        help="Skip hole-filling/simplification of the raw decoded mesh (postprocessing_utils.to_glb's with_mesh_postprocess=False). "
        "Forces vertex-color output (no UV texture bake) since untouched raw geometry isn't a good UV-parametrization target.",
    )
    p.add_argument(
        "--no-texture-baking", action="store_true",
        help="Skip the multi-view texture bake (renders 100 views through pytorch3d) and use per-vertex color instead -- "
        "much faster/lower VRAM, matching the README demo's own use_vertex_color=True default. Try this first if a run is slow or OOMs.",
    )
    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)

    checkpoint_config = Path(_SAM3D_DIR) / "checkpoints" / args.checkpoint_tag / "pipeline.yaml"
    if not checkpoint_config.exists():
        print(
            f"ERROR: checkpoint config not found at {checkpoint_config}. "
            "Follow third_party/sam-3d-objects/doc/setup.md's 'Getting Checkpoints' "
            "section (hf download facebook/sam-3d-objects, then move "
            "checkpoints/<tag>-download/checkpoints -> checkpoints/<tag>).",
            file=sys.stderr,
        )
        sys.exit(1)

    with_mesh_postprocess = not args.no_mesh_postprocess
    with_texture_baking = not args.no_texture_baking and with_mesh_postprocess
    use_vertex_color = not with_texture_baking

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

    t0 = time.time()

    try:
        import torch
        import numpy as np
        from PIL import Image
        from omegaconf import OmegaConf
        from hydra.utils import instantiate
        # NOTE: importing sam3d_objects.model.backbone.tdfy_dit.utils.postprocessing_utils
        # here (next line) transitively imports the `sparse` attention module
        # before hydra.utils.instantiate() below reaches
        # sam3d_objects.pipeline.inference_pipeline (whose module-level
        # set_attention_backend() forces ATTN_BACKEND=flash_attn on
        # A100/H100/H200 GPUs). Because `sparse`'s own backend selection reads
        # the env var once at ITS import time, it locks in "sdpa" instead --
        # confirmed in logs/sam3d_setup/09_smoketest_generate.log. Harmless
        # (sdpa works, just not the fastest option); fixing it would mean
        # importing inference_pipeline first, which isn't worth the added
        # coupling for a one-shot CLI script.
        import sam3d_objects  # noqa: F401 -- side-effect import, see notebook/inference.py's own "do not remove"
        from sam3d_objects.model.backbone.tdfy_dit.utils import postprocessing_utils
    except Exception as e:
        print(f"ERROR: failed to import the sam3d-objects stack: {e}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(1)

    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)
        image_np = np.array(image).astype("uint8")
    except Exception as e:
        print(f"ERROR: could not open image {args.image}: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        # Same config-loading recipe as notebook/inference.py's Inference.__init__
        # (config_path -> OmegaConf.load -> override 3 fields -> hydra instantiate),
        # deliberately reimplemented here rather than importing that module -- see
        # the module docstring's "UNTESTED end-to-end" section for why.
        config = OmegaConf.load(str(checkpoint_config))
        config.rendering_engine = "pytorch3d"  # avoids needing a working nvdiffrast/OpenGL build
        config.compile_model = False  # torch.compile adds a slow first-call warmup we don't want for one-shot CLI runs
        config.workspace_dir = str(checkpoint_config.parent)
        pipeline = instantiate(config)

        # image_np already carries the object mask in its alpha channel
        # (this script's contract, matching scripts/trellis_generate.py); pass
        # mask=None so InferencePipelinePointMap.merge_image_and_mask() uses it
        # as-is instead of trying to merge a second mask in.
        #
        # with_mesh_postprocess/with_texture_baking=False here on purpose: this
        # only decodes the slat and gets the internal (cheap, no-op) to_glb
        # fallback path so we can read raw output["mesh"][0]/["gaussian"][0];
        # the real postprocessing with this script's own --simplify-ratio/
        # --texture-size happens in the explicit to_glb call below (see the
        # module docstring's "catch" paragraph for why).
        output = pipeline.run(
            image_np,
            None,
            seed=args.seed,
            stage1_only=False,
            with_mesh_postprocess=False,
            with_texture_baking=False,
            with_layout_postprocess=False,
            use_vertex_color=True,
            decode_formats=["mesh", "gaussian"],
        )

        raw_mesh = output.get("mesh")
        raw_gaussian = output.get("gaussian")
        if not raw_mesh or not raw_gaussian:
            raise RuntimeError(
                "pipeline.run() returned no 'mesh'/'gaussian' outputs -- decode_slat() "
                "only populates these when decode_formats includes them (passed above) "
                "and the corresponding slat_decoder_{mesh,gs} checkpoints loaded "
                "correctly. Check the pipeline construction logs above."
            )

        glb = postprocessing_utils.to_glb(
            raw_gaussian[0],
            raw_mesh[0],
            simplify=args.simplify_ratio,
            texture_size=args.texture_size,
            verbose=True,
            with_mesh_postprocess=with_mesh_postprocess,
            with_texture_baking=with_texture_baking,
            use_vertex_color=use_vertex_color,
            rendering_engine="pytorch3d",
        )

        args.out.parent.mkdir(parents=True, exist_ok=True)
        glb.export(str(args.out))

        splat_path = args.out.with_name(args.out.stem + "_splat.ply")
        gs = output.get("gs")
        if gs is not None:
            gs.save_ply(str(splat_path))
        else:
            splat_path = None
    except Exception as e:
        print(f"ERROR: SAM 3D Objects 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()),
        "output_splat_ply": str(splat_path.resolve()) if splat_path else None,
        "checkpoint_tag": args.checkpoint_tag,
        "seed": args.seed,
        "with_mesh_postprocess": with_mesh_postprocess,
        "with_texture_baking": with_texture_baking,
        "use_vertex_color": use_vertex_color,
        "simplify_ratio": args.simplify_ratio if with_mesh_postprocess else None,
        "texture_size": args.texture_size if with_texture_baking else None,
        "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()