Add docker setup

.dockerignore

@@ -0,0 +1,19 @@
+.git
+.github
+.venv
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+build/
+dist/
+checkpoints/
+notebook/.ipynb_checkpoints/
+.mypy_cache/
+.pytest_cache/
+node_modules/
+pip-wheel-metadata/
+*.egg-info/
+patching/__pycache__/
+patching/hydra.cpython-*
+uv.lock.bak

doc/setup.md

@@ -1,34 +1,40 @@
 # Setup
 
+
 ## Prerequisites
 
 * A linux 64-bits architecture (i.e. `linux-64` platform in `mamba info`).
 * A NVIDIA GPU with at least 32 Gb of VRAM.
 
 ## 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
+```
 
-# for pytorch/cuda dependencies
-export PIP_EXTRA_INDEX_URL="https://pypi.ngc.nvidia.com https://download.pytorch.org/whl/cu121"
+Adding these exports to `~/.bashrc` is recommended so the correct toolkit is always picked up.
 
-# 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
+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
 
-# 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 the base set of pinned dependencies
+uv sync
+source .venv/bin/activate
 
 # 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 +60,3 @@ hf download \
 mv checkpoints/${TAG}-download/checkpoints checkpoints/${TAG}
 rm -rf checkpoints/${TAG}-download
 ```
-
-

docker-compose.yml

@@ -0,0 +1,31 @@
+version: "3.9"
+
+services:
+  sam3d:
+    build:
+      context: .
+      dockerfile: docker/Dockerfile
+      args:
+        USERNAME: ${LOCAL_USER:-developer}
+        USER_UID: ${LOCAL_UID:-1000}
+        USER_GID: ${LOCAL_GID:-1000}
+        EXTRA_APT_PACKAGES: ${EXTRA_APT_PACKAGES:-}
+        TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST:-sm_80}
+    image: sam3d-dev
+    working_dir: /workspace
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=all
+      - NVIDIA_DRIVER_CAPABILITIES=compute,utility
+    volumes:
+      - .:/workspace
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities:
+                - gpu
+    shm_size: 16g
+    tty: true
+    stdin_open: true

docker/Dockerfile

@@ -0,0 +1,112 @@
+# syntax=docker/dockerfile:1.7
+FROM nvidia/cuda:12.9.1-cudnn-devel-ubuntu24.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+ARG USERNAME=developer
+ARG USER_UID=1000
+ARG USER_GID=1000
+ARG EXTRA_APT_PACKAGES=""
+ARG TORCH_CUDA_ARCH_LIST="86"
+
+# Base OS setup
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends \
+    bash \
+    build-essential \
+    ca-certificates \
+    cmake \
+    curl \
+    ffmpeg \
+    git \
+    git-lfs \
+    libffi-dev \
+    libglib2.0-0 \
+    libgl1 \
+    libjpeg-dev \
+    libomp-dev \
+    libopenexr-dev \
+    libpng-dev \
+    libsm6 \
+    libssl-dev \
+    libtiff-dev \
+    libx11-6 \
+    libxext6 \
+    libxrender1 \
+    locales \
+    ninja-build \
+    pciutils \
+    pkg-config \
+    python3 \
+    python3-dev \
+    python3-venv \
+    sudo \
+    unzip \
+    vim \
+    wget \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN if [ -n "${EXTRA_APT_PACKAGES}" ]; then \
+        apt-get update && \
+        apt-get install -y --no-install-recommends ${EXTRA_APT_PACKAGES} && \
+        rm -rf /var/lib/apt/lists/*; \
+    fi
+
+# Configure UTF-8 locale
+RUN locale-gen en_US.UTF-8 && update-locale LANG=en_US.UTF-8
+ENV LANG=en_US.UTF-8 \
+    LC_ALL=en_US.UTF-8
+
+# CUDA 12.9 toolchain paths are required during dependency builds
+ENV CUDA_HOME=/usr/local/cuda-12.9 \
+    PATH=${CUDA_HOME}/bin:${PATH} \
+    LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
+
+ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
+
+# Install uv (https://github.com/astral-sh/uv) system-wide
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
+    && install -Dm755 /root/.local/bin/uv /usr/local/bin/uv
+
+# Create a non-root user that matches the host UID/GID for seamless volume access
+RUN set -eux; \
+    if ! getent group "${USER_GID}" >/dev/null; then \
+    groupadd --gid "${USER_GID}" "${USERNAME}"; \
+    fi; \
+    EXISTING_USER="$(getent passwd "${USER_UID}" | cut -d: -f1 || true)"; \
+    if [ -n "${EXISTING_USER}" ] && [ "${EXISTING_USER}" != "${USERNAME}" ]; then \
+    usermod -l "${USERNAME}" "${EXISTING_USER}"; \
+    fi; \
+    if id -u "${USERNAME}" >/dev/null 2>&1; then \
+    usermod -u "${USER_UID}" -g "${USER_GID}" "${USERNAME}"; \
+    else \
+    useradd --uid "${USER_UID}" --gid "${USER_GID}" -m "${USERNAME}"; \
+    fi; \
+    usermod -d /home/${USERNAME} -m "${USERNAME}"; \
+    usermod -aG sudo "${USERNAME}"; \
+    echo "%${USERNAME} ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/${USERNAME}; \
+    chmod 440 /etc/sudoers.d/${USERNAME}; \
+    mkdir -p /workspace /opt/sam3d /opt/uv-python; \
+    chown -R ${USERNAME}:${USER_GID} /workspace /opt/sam3d /opt/uv-python
+
+USER ${USERNAME}
+WORKDIR /opt/sam3d
+
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv-python
+
+# Copy the minimum project files required to materialize the environment
+COPY --chown=${USERNAME}:${USER_GID} pyproject.toml uv.lock README.md ./ 
+COPY --chown=${USERNAME}:${USER_GID} patching ./patching
+COPY --chown=${USERNAME}:${USER_GID} sam3d_objects ./sam3d_objects
+
+# Provision the Python 3.11 environment with uv and patch Hydra
+RUN uv python install 3.11 \
+    && uv sync --frozen --python 3.11 --no-editable \
+    && .venv/bin/python patching/hydra
+
+ENV VIRTUAL_ENV=/opt/sam3d/.venv
+ENV PATH=${VIRTUAL_ENV}/bin:${PATH}
+
+WORKDIR /workspace
+
+# Default shell
+CMD ["bash"]

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

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"]