troutmoose/foundationpose

1# Base image with FoundationPose dependencies split into CPU (L1) and GPU (L2)2 3# Stage 1: CPU-only base with Python deps4FROM docker.io/ubuntu:22.04 AS foundationpose-base-l15 6ENV DEBIAN_FRONTEND=noninteractive7 8# Install system deps needed to build/run python packages9RUN apt-get update && apt-get install -y --no-install-recommends \10    ca-certificates \11    git \12    python3.10 \13    python3-pip \14    build-essential \15    cmake \16    ninja-build \17    libeigen3-dev \18    python3.10-dev \19    libboost-system-dev \20    libboost-program-options-dev \21    pybind11-dev \22    libgl1 \23    libglib2.0-0 \24    libgomp1 \25    libsm6 \26    libxext6 \27    libxrender1 \28    libxkbcommon0 \29    libx11-6 \30    libxrandr2 \31    libxi6 \32    libxinerama1 \33    libxcursor1 \34    libspatialindex-dev \35    && rm -rf /var/lib/apt/lists/*36 37# Set python as default38RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \39    update-alternatives --install /usr/bin/python python /usr/bin/python3.10 140 41RUN python3 -m pip install --no-cache-dir --upgrade pip42 43# Core python deps (CPU-safe)44# Keep NumPy <2 for extension compatibility45RUN pip install --no-cache-dir \46    "numpy<2" \47    Pillow>=10.0.0 \48    gradio>=4.0.0 \49    huggingface-hub>=0.20.0 \50    scipy==1.12.0 \51    scikit-image==0.22.0 \52    scikit-learn==1.4.1.post1 \53    kornia==0.7.2 \54    einops==0.7.0 \55    timm==0.9.16 \56    pyyaml==6.0.1 \57    ruamel.yaml==0.18.6 \58    omegaconf==2.3.0 \59    h5py==3.10.0 \60    numba==0.59.1 \61    imageio==2.34.0 \62    joblib==1.3.2 \63    psutil==6.1.1 \64    albumentations==1.4.2 \65    imgaug==0.4.0 \66    seaborn==0.13.2 \67    plotly==5.20.0 \68    bokeh==3.4.0 \69    colorama==0.4.6 \70    GPUtil==1.4.0 \71    simplejson==3.19.2 \72    openpyxl==3.1.2 \73    xlsxwriter==3.2.0 \74    nodejs==0.1.1 \75    jupyterlab==4.1.5 \76    ipywidgets==8.1.2 \77    py-spy==0.3.14 \78    videoio==0.2.8 \79    pypng==0.20220715.0 \80    roma==1.4.4 \81    transformations==2024.6.1 \82    meshcat==0.3.2 \83    webdataset==0.2.86 \84    wandb==0.16.5 \85    g4f==0.2.7.1 \86    objaverse==0.1.7 \87    opencv-python==4.9.0.80 \88    opencv-contrib-python==4.9.0.80 \89    open3d==0.18.0 \90    pyglet==1.5.28 \91    pysdf==0.1.9 \92    trimesh==4.2.2 \93    xatlas==0.0.9 \94    rtree==1.2.0 \95    pyrender==0.1.45 \96    pyOpenGL>=3.1.0 \97    pyOpenGL_accelerate>=3.1.0 \98    pybullet==3.2.6 \99    pycocotools==2.0.7 \100    Panda3D==1.10.14 \101    pin==2.7.0 \102    && pip cache purge103 104# Stage 2: GPU-enabled base105FROM docker.io/nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS foundationpose-base-l2106 107ENV DEBIAN_FRONTEND=noninteractive108ENV CUDA_HOME=/usr/local/cuda109ENV PATH=${CUDA_HOME}/bin:${PATH}110ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}111 112# Only build for T4 (7.5) - reduces compilation memory by 50%113ENV TORCH_CUDA_ARCH_LIST="7.5"114 115# Install system deps116RUN rm -f /etc/apt/sources.list.d/cuda*.list /etc/apt/sources.list.d/*.list && \117    apt-get update && apt-get install -y --no-install-recommends --allow-unauthenticated \118    ca-certificates \119    && apt-get clean && rm -rf /var/lib/apt/lists/* && \120    apt-get update && apt-get install -y --no-install-recommends \121    python3.10 \122    python3-pip \123    git \124    libgl1 \125    libglib2.0-0 \126    libgomp1 \127    libsm6 \128    libxext6 \129    libxrender1 \130    libxkbcommon0 \131    libx11-6 \132    libxrandr2 \133    libxi6 \134    libxinerama1 \135    libxcursor1 \136    libspatialindex-dev \137    && rm -rf /var/lib/apt/lists/* \138    && apt-get clean139 140# Set python as default141RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \142    update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1143 144# Upgrade pip145RUN python3 -m pip install --no-cache-dir --upgrade pip146 147# Copy CPU-only python deps from L1148COPY --from=foundationpose-base-l1 /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages149COPY --from=foundationpose-base-l1 /usr/local/bin /usr/local/bin150 151# Install PyTorch (CUDA 11.8)152RUN pip install --no-cache-dir torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118153 154# GPU/torch-dependent deps155RUN pip install --no-cache-dir \156    fvcore==0.1.5.post20221221 \157    torchnet==0.0.4 \158    ultralytics==8.0.120 \159    warp-lang==1.0.2 \160    && pip cache purge161 162# Build deps required for CUDA extensions163RUN apt-get update && apt-get install -y --no-install-recommends \164    cmake \165    build-essential \166    ninja-build \167    libeigen3-dev \168    python3.10-dev \169    libboost-system-dev \170    libboost-program-options-dev \171    pybind11-dev \172    && rm -rf /var/lib/apt/lists/*173 174WORKDIR /app175 176EXPOSE 7860177 178# ----- L2 build steps to include all installation work -----179# Install SAM + pytorch3d180RUN pip install --no-cache-dir "numpy<2" transformers==4.41.2 \181    && pip install --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu118_pyt210/download.html182 183# Set MAX_JOBS=1 BEFORE any CUDA compilation to limit memory usage184ENV MAX_JOBS=1185 186# Clone FoundationPose source187RUN git clone --depth 1 https://github.com/NVlabs/FoundationPose.git /app/FoundationPose \188    && cd /app/FoundationPose/bundlesdf/mycuda \189    && sed -i 's/-std=c++14/-std=c++17/g' setup.py190 191# Build CPU-only C++ code192WORKDIR /app/FoundationPose193RUN cd mycpp && mkdir -p build && cd build && cmake .. && make194 195# Download model weights (246MB)196WORKDIR /app197COPY download_weights.py ./download_weights.py198RUN python3 download_weights.py199 200# Install nvdiffrast (CUDA rasterizer) - needs GPU, build here201RUN git clone --depth 1 https://github.com/NVlabs/nvdiffrast.git /tmp/nvdiffrast \202    && cd /tmp/nvdiffrast \203    && python3 setup.py build_ext --inplace204RUN python3 -c "import shutil, sysconfig, glob; from pathlib import Path; site=Path(sysconfig.get_paths()['purelib']); src=Path('/tmp/nvdiffrast/nvdiffrast'); dst=site/'nvdiffrast'; shutil.rmtree(dst, ignore_errors=True); shutil.copytree(src, dst); so_files=(glob.glob('/tmp/nvdiffrast/_nvdiffrast_c*.so') + glob.glob('/tmp/nvdiffrast/nvdiffrast/_nvdiffrast_c*.so') + glob.glob('/tmp/nvdiffrast/build/lib.*/*_nvdiffrast_c*.so')); [shutil.copy2(p, site) for p in so_files]"205RUN python3 -c "import sysconfig; from pathlib import Path; site=Path(sysconfig.get_paths()['purelib']); dist=site/'nvdiffrast-0.0.0.dist-info'; dist.mkdir(exist_ok=True); (dist/'METADATA').write_text('Metadata-Version: 2.1\\nName: nvdiffrast\\nVersion: 0.0.0\\n'); (dist/'WHEEL').write_text('Wheel-Version: 1.0\\nGenerator: manual\\nRoot-Is-Purelib: false\\nTag: py3-none-any\\n'); (dist/'top_level.txt').write_text('nvdiffrast\\n'); (dist/'RECORD').write_text('')"206RUN python3 -c "import nvdiffrast.torch"207RUN rm -rf /tmp/nvdiffrast208 209# Build CUDA extensions (requires GPU)210WORKDIR /app/FoundationPose211RUN cd bundlesdf/mycuda && pip install . --no-build-isolation212