# Final stage Dockerfile - optimized for HuggingFace
FROM gpue/foundationpose-base-l2:latest

# FoundationPose configuration
ENV FOUNDATIONPOSE_MODEL_REPO=gpue/foundationpose-weights
ENV USE_REAL_MODEL=true

# Ensure NumPy 1.x for CUDA extension compatibility and install SAM/pytorch3d
RUN pip install --no-cache-dir "numpy<2" transformers==4.41.2 \
    && pip install --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu118_pyt210/download.html

# Set MAX_JOBS=1 BEFORE any CUDA compilation to limit memory usage
ENV MAX_JOBS=1

# Clone FoundationPose source
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 CPU-only C++ code
WORKDIR /app/FoundationPose
RUN cd mycpp && mkdir -p build && cd build && cmake .. && make

# Download model weights (246MB)
WORKDIR /app
COPY download_weights.py ./download_weights.py
RUN python3 download_weights.py

# Install nvdiffrast (CUDA rasterizer) - needs GPU, build here
RUN git clone --depth 1 https://github.com/NVlabs/nvdiffrast.git /tmp/nvdiffrast \
    && cd /tmp/nvdiffrast \
    && python3 setup.py build_ext --inplace
RUN 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]"
RUN 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('')"
RUN python3 -c "import nvdiffrast.torch"
RUN rm -rf /tmp/nvdiffrast

# Build CUDA extensions (requires GPU - only part that needs HuggingFace GPU)
WORKDIR /app/FoundationPose
RUN cd bundlesdf/mycuda && pip install . --no-build-isolation

# Copy application files
WORKDIR /app
COPY app.py client.py estimator.py masks.py .

CMD ["python3", "app.py"]