Add USD Export Support for Meshed Objects

demo.py

@@ -13,9 +13,18 @@
 image = load_image("notebook/images/shutterstock_stylish_kidsroom_1640806567/image.png")
 mask = load_single_mask("notebook/images/shutterstock_stylish_kidsroom_1640806567", index=14)
 
+# Optional USD export path (set to a filename to enable)
+usd_path = "reconstruction.usd"  # e.g., "reconstruction.usd"
+# Adjust the scale factor to match your scene units (default 100)
+usd_scale_factor = 100.0
+
 # run model
-output = inference(image, mask, seed=42)
+output = inference(image, mask, seed=42, export_usd_path=usd_path, usd_scale_factor=usd_scale_factor)
 
 # export gaussian splat
 output["gs"].save_ply(f"splat.ply")
 print("Your reconstruction has been saved to splat.ply")
+if output.get("usd_path"):
+    print(f"USD mesh with texture saved to {output['usd_path']}")
+elif usd_path is not None:
+    print("USD export requested but failed; check logs for details.")

doc/setup.md

@@ -7,28 +7,38 @@
 
 ## 1. Setup Python Environment
 
-The following will install the default environment. If you use `conda` instead of `mamba`, replace its name in the first two lines. Note that you may have to build the environment on a compute node with GPU (e.g., you may get a `RuntimeError: Not compiled with GPU support` error when running certain parts of the code that use Pytorch3D).
+`uv` is now the default way to create the environment. The additional NVIDIA/PyTorch indices and Kaolin find-links are configured in `pyproject.toml`, so no extra exports are needed. Use Python 3.11 (the project targets 3.11 only). **CUDA toolkit 12.9 is required for building gsplat/pytorch3d**; install it from the [CUDA 12.9 archive](https://developer.nvidia.com/cuda-12-9-0-download-archive) and export the paths before running `uv sync`:
 
 ```bash
-# create sam3d-objects environment
-mamba env create -f environments/default.yml
-mamba activate sam3d-objects
+export CUDA_HOME=/usr/local/cuda-12.9
+export PATH=$CUDA_HOME/bin:$PATH
+export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
+```
+
+Adding these exports to `~/.bashrc` is recommended so the correct toolkit is always picked up.
 
-# for pytorch/cuda dependencies
-export PIP_EXTRA_INDEX_URL="https://pypi.ngc.nvidia.com https://download.pytorch.org/whl/cu121"
+Example workflow:
+
+```bash
+# install uv if you don't have it yet
+curl -LsSf https://astral.sh/uv/install.sh | sh
+source $HOME/.local/bin/env
 
-# install sam3d-objects and core dependencies
-pip install -e '.[dev]'
-pip install -e '.[p3d]' # pytorch3d dependency on pytorch is broken, this 2-step approach solves it
+# install the base set of pinned dependencies
+uv sync
+source .venv/bin/activate
 
-# for inference
-export PIP_FIND_LINKS="https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.5.1_cu121.html"
-pip install -e '.[inference]'
+# install nvdiffrast from source while the venv is active
+git clone https://github.com/NVlabs/nvdiffrast.git
+cd nvdiffrast
+uv pip install .
 
 # patch things that aren't yet in official pip packages
 ./patching/hydra # https://github.com/facebookresearch/hydra/pull/2863
 ```
 
+> If you still prefer a Conda-based workflow for GPU toolchains, you can reuse `environments/default.yml` to provision system libraries, then activate your environment and run `uv sync` inside it to install Python dependencies.
+
 ## 2. Getting Checkpoints
 
 ### From HuggingFace
@@ -54,5 +64,3 @@ hf download \
 mv checkpoints/${TAG}-download/checkpoints checkpoints/${TAG}
 rm -rf checkpoints/${TAG}-download
 ```
-
-

notebook/inference.py

@@ -2,7 +2,16 @@
 import os
 
 # not ideal to put that here
-os.environ["CUDA_HOME"] = os.environ["CONDA_PREFIX"]
+# Prefer an existing CUDA_HOME, otherwise fall back to conda prefix or torch's detected CUDA path.
+_cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX")
+if not _cuda_home:
+    try:
+        from torch.utils.cpp_extension import CUDA_HOME as _torch_cuda_home
+    except Exception:
+        _torch_cuda_home = None
+    _cuda_home = _torch_cuda_home
+if _cuda_home:
+    os.environ["CUDA_HOME"] = _cuda_home
 os.environ["LIDRA_SKIP_INIT"] = "true"
 
 import sys
@@ -104,19 +113,25 @@ def __call__(
         mask: Optional[Union[None, Image.Image, np.ndarray]],
         seed: Optional[int] = None,
         pointmap=None,
+        export_usd_path: Optional[str] = None,
+        usd_scale_factor: float = 100.0,
+        embed_textures: bool = True,
     ) -> dict:
         image = self.merge_mask_to_rgba(image, mask)
         return self._pipeline.run(
             image,
             None,
             seed,
             stage1_only=False,
-            with_mesh_postprocess=False,
-            with_texture_baking=False,
+            with_mesh_postprocess=True,
+            with_texture_baking=True,
             with_layout_postprocess=True,
             use_vertex_color=True,
             stage1_inference_steps=None,
             pointmap=pointmap,
+            usd_path=export_usd_path,
+            usd_scale_factor=usd_scale_factor,
+            embed_textures=embed_textures,
         )
 
 

pyproject.toml

@@ -1,30 +1,132 @@
 [build-system]
-requires = ["hatchling", "hatch-requirements-txt"]
+requires = ["hatchling"]
 build-backend = "hatchling.build"
 
-[tool.hatch.envs.default.env-vars]
-PIP_EXTRA_INDEX_URL = "https://pypi.ngc.nvidia.com https://download.pytorch.org/whl/cu121"
-
-[tool.hatch.metadata]
-# for git-referenced dependencies
-allow-direct-references = true
-
 [project]
 name = "sam3d_objects"
 version = "0.0.1"
-# required for "hatch-requirements-txt" to work
-dynamic = ["dependencies", "optional-dependencies"]
+requires-python = ">=3.11,<3.12"
+dependencies = [
+  "MoGe @ git+https://github.com/microsoft/MoGe.git@a8c37341bc0325ca99b9d57981cc3bb2bd3e255b",
+  "Werkzeug==3.0.6",
+  "astor==0.8.1",
+  "async-timeout==4.0.3",
+  "auto_gptq==0.7.1",
+  "autoflake==2.3.1",
+  "av==12.0.0",
+  "bitsandbytes==0.43.0",
+  "black==24.3.0",
+  "bpy==4.3.0",
+  "colorama==0.4.6",
+  "conda-pack==0.7.1",
+  "crcmod==1.7",
+  "cuda-python>=12.6.0",
+  "dataclasses==0.6",
+  "decord==0.6.0",
+  "deprecation==2.1.0",
+  "easydict==1.13",
+  "einops-exts==0.0.4",
+  "exceptiongroup==1.2.0",
+  "fastavro==1.9.4",
+  "fasteners==0.19",
+  "flake8==7.0.0",
+  "Flask==3.0.3",
+  "fqdn==1.5.1",
+  "ftfy==6.2.0",
+  "fvcore==0.1.5.post20221221",
+  "gdown==5.2.0",
+  "h5py==3.12.1",
+  "hdfs==2.7.3",
+  "httplib2==0.22.0",
+  "hydra-core==1.3.2",
+  "hydra-submitit-launcher==1.2.0",
+  "igraph==0.11.8",
+  "imath==0.0.2",
+  "isoduration==20.11.0",
+  "jsonlines==4.0.0",
+  "jsonpickle==3.0.4",
+  "jsonpointer==2.4",
+  "jupyter==1.1.1",
+  "librosa==0.10.1",
+  "lightning==2.3.3",
+  "loguru==0.7.2",
+  "mosaicml-streaming==0.7.5",
+  "nvidia-cuda-nvcc-cu12>=12.6.77",
+  "nvidia-pyindex==1.0.9",
+  "objsize==0.7.0",
+  "open3d==0.18.0",
+  "opencv-python==4.9.0.80",
+  "OpenEXR==3.3.3",
+  "optimum==1.18.1",
+  "optree==0.14.1",
+  "orjson==3.10.0",
+  "panda3d-gltf==1.2.1",
+  "pdoc3==0.10.0",
+  "peft==0.10.0",
+  "pip-system-certs==4.0",
+  "point-cloud-utils==0.29.5",
+  "polyscope==2.3.0",
+  "pycocotools==2.0.7",
+  "pydot==1.4.2",
+  "pymeshfix==0.17.0",
+  "pymongo==4.6.3",
+  "pyrender==0.1.45",
+  "PySocks==1.7.1",
+  "pytest==8.1.1",
+  "python-pycg==0.9.2",
+  "randomname==0.2.1",
+  "roma==1.5.1",
+  "rootutils==1.0.7",
+  "Rtree==1.3.0",
+  "sagemaker==2.242.0",
+  "scikit-image==0.23.1",
+  "sentence-transformers==2.6.1",
+  "simplejson==3.19.2",
+  "smplx==0.1.28",
+  "spconv-cu121==2.3.8",
+  "tensorboard==2.16.2",
+  "torch==2.8.0+cu129",
+  "timm==0.9.16",
+  "tomli==2.0.1",
+  "torchaudio==2.8.0+cu129",
+  "uri-template==1.3.0",
+  "usort==1.0.8.post1",
+  "wandb==0.20.0",
+  "webcolors==1.13",
+  "webdataset==0.2.86",
+  "xatlas==0.0.9",
+  "xformers==0.0.33+5d4b92a5.d20251029",
+  "findpydeps",
+  "lovely_tensors",
+  "pipdeptree",
+  "pytest",
+  "gradio==5.49.0",
+  "gsplat @ git+https://github.com/nerfstudio-project/gsplat.git@2323de5905d5e90e035f792fe65bad0fedd413e7",
+  "kaolin==0.18.0",
+  "seaborn==0.13.2",
+  "flash_attn==2.8.3",
+  "pytorch3d @ git+https://github.com/facebookresearch/pytorch3d.git@75ebeeaea0908c5527e7b1e305fbc7681382db47",
+]
 
 [tool.hatch.build]
 ignore-vcs = true
 include = ["**/*.py"]
 exclude = ["conftest.py", "*_test.py"]
 packages = ["sam3d_objects"]
 
-[tool.hatch.metadata.hooks.requirements_txt]
-files = ["requirements.txt"]
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.uv]
+extra-index-url = [
+  "https://pypi.ngc.nvidia.com",
+  "https://download.pytorch.org/whl/cu129",
+]
+find-links = ["https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.8.0_cu129.html"]
+index-strategy = "unsafe-best-match"
 
-[tool.hatch.metadata.hooks.requirements_txt.optional-dependencies]
-p3d = ["requirements.p3d.txt"]
-inference = ["requirements.inference.txt"]
-dev = ["requirements.dev.txt"]
+[tool.uv.extra-build-dependencies]
+flash-attn = ["torch==2.8.0+cu129"]
+gsplat = ["torch==2.8.0+cu129"]
+nvidia-pyindex = ["pip"]
+pytorch3d = ["torch==2.8.0+cu129", "numpy"]