Add FoundationPose dependencies, build nvdiffrast on HuggingFace with GPU

Dockerfile

@@ -6,41 +6,19 @@ FROM gpue/foundationpose-base:latest
 ENV FOUNDATIONPOSE_MODEL_REPO=gpue/foundationpose-weights
 ENV USE_REAL_MODEL=true
 
-# Downgrade NumPy to 1.x for CUDA extension compatibility
+# Ensure NumPy 1.x for CUDA extension compatibility
 RUN pip install --no-cache-dir "numpy<2"
 
-# Install build tools temporarily (CUDA compiler already in devel base)
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    cmake \
-    build-essential \
-    ninja-build \
-    libeigen3-dev \
-    python3.10-dev \
-    libboost-system-dev \
-    libboost-program-options-dev \
-    pybind11-dev \
-    && rm -rf /var/lib/apt/lists/*
+# Install nvdiffrast (CUDA rasterizer) - needs GPU, build here
+RUN pip install --no-cache-dir --no-build-isolation git+https://github.com/NVlabs/nvdiffrast.git
 
-# Build FoundationPose C++ extensions (requires GPU)
+# Build CUDA extensions (requires GPU - only part that needs HuggingFace GPU)
 # MAX_JOBS=1 to reduce memory usage during compilation
 WORKDIR /app/FoundationPose
 ENV MAX_JOBS=1
 RUN cd bundlesdf/mycuda && pip install . --no-build-isolation
-RUN cd mycpp && mkdir -p build && cd build && cmake .. && make
 
-# Remove build tools to save space (keep CUDA compiler)
-RUN apt-get purge -y cmake ninja-build libeigen3-dev python3.10-dev \
-    libboost-system-dev libboost-program-options-dev pybind11-dev && \
-    apt-get autoremove -y && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
-
-# Download model weights
+# Note: mycpp build, weights download, and build deps are already in base image
 WORKDIR /app
-RUN python3 -c "from huggingface_hub import snapshot_download; \
-snapshot_download(repo_id='gpue/foundationpose-weights', local_dir='weights', repo_type='model')"
-
-# Clean pip cache
-RUN pip cache purge
 
 CMD ["python3", "app.py"]

Dockerfile.base

@@ -47,15 +47,48 @@ RUN pip install --no-cache-dir \
     huggingface-hub>=0.20.0 \
     && pip cache purge
 
-# Clone FoundationPose (but don't build yet - that's done in final stage)
+# Install build dependencies (keep them for faster HuggingFace builds)
+# Install BEFORE nvdiffrast because it needs python3.10-dev
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    cmake \
+    build-essential \
+    ninja-build \
+    libeigen3-dev \
+    python3.10-dev \
+    libboost-system-dev \
+    libboost-program-options-dev \
+    pybind11-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install FoundationPose dependencies
+RUN pip install --no-cache-dir \
+    trimesh==4.2.2 \
+    scipy==1.12.0 \
+    scikit-image==0.22.0 \
+    kornia==0.7.2 \
+    einops==0.7.0 \
+    timm==0.9.16 \
+    transformations==2024.6.1 \
+    pyyaml==6.0.1 \
+    && pip cache purge
+
+# Note: nvdiffrast will be built in final Dockerfile on HuggingFace (needs GPU)
+
+# Clone FoundationPose
 RUN git clone --depth 1 https://github.com/NVlabs/FoundationPose.git /app/FoundationPose && \
     cd /app/FoundationPose/bundlesdf/mycuda && \
     sed -i 's/-std=c++14/-std=c++17/g' setup.py
 
+# Build mycpp (non-GPU C++ code - can be built without GPU)
+WORKDIR /app/FoundationPose
+RUN cd mycpp && mkdir -p build && cd build && cmake .. && make
+
+# Download model weights (246MB)
+WORKDIR /app
+RUN python3 -c "from huggingface_hub import snapshot_download; \
+snapshot_download(repo_id='gpue/foundationpose-weights', local_dir='weights', repo_type='model')"
+
 # Copy application files
 COPY app.py client.py estimator.py ./
 
-# Create weights directory
-RUN mkdir -p weights
-
 EXPOSE 7860