FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV PIP_NO_CACHE_DIR=1

# System deps (compiler toolchain + git + minimal libs for OpenCV headless)
RUN apt-get update && apt-get install -y \
    python3.10 python3.10-dev python3-pip \
    git git-lfs \
    build-essential cmake ninja-build \
    libglib2.0-0 libsm6 libxrender1 libxext6 \
    && rm -rf /var/lib/apt/lists/*

RUN git lfs install

# Upgrade pip tooling
RUN python3.10 -m pip install --upgrade pip setuptools wheel

# Pin numpy first (your requirement)
RUN pip install numpy==1.26.4
RUN printf "numpy==1.26.4\n" > /tmp/constraints.txt

# Install PyTorch cu118 pinned versions (your requirement)
RUN pip uninstall -y torch torchvision torchaudio || true

RUN pip install -c /tmp/constraints.txt torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 \
    --index-url https://download.pytorch.org/whl/cu118
    
# Clone ProFuse with submodules
RUN git clone --recurse-submodules https://github.com/chiou1203/ProFuse.git /workspace/ProFuse

# Install python deps (match your colab steps)
WORKDIR /workspace/ProFuse

WORKDIR /workspace/ProFuse/feature_registration

RUN pip install -c /tmp/constraints.txt tqdm plyfile open_clip_torch
RUN pip install -c /tmp/constraints.txt gradio requests huggingface_hub opencv-python-headless
RUN pip install -c /tmp/constraints.txt ninja kmeans_pytorch faiss-cpu

# Install submodules + extra deps you used
# sanity check: torch must exist before building extensions
RUN python3.10 -c "import torch; print('torch', torch.__version__); print('cuda', torch.version.cuda)"

# pure python deps first
#RUN pip install -q ninja kmeans_pytorch faiss-cpu

# build CUDA/C++ extensions in the main env (avoid PEP517 isolated env missing torch)
# Build-time CUDA settings (HF build has no GPU -> must set arch list)
ENV CUDA_HOME=/usr/local/cuda
ENV FORCE_CUDA=1

# IMPORTANT: keep this small to reduce compile time + RAM.
# If you plan to run on A10G, 8.6 is usually enough.
ENV TORCH_CUDA_ARCH_LIST="8.6"

# Limit compilation parallelism to avoid OOM (exit 137)
ENV MAX_JOBS=1
ENV CMAKE_BUILD_PARALLEL_LEVEL=1
ENV MAKEFLAGS="-j1"

RUN pip install -c /tmp/constraints.txt -q --no-build-isolation --no-deps submodules/simple-knn
RUN pip install -q --no-build-isolation submodules/langsplat-rasterization
RUN pip install -q --no-build-isolation submodules/segment-anything-langsplat

#RUN pip uninstall -y gradio || true
#RUN pip install -U "gradio>=4.0.0" "gradio_client>=0.8.0"

# App/runtime deps
#RUN pip install gradio requests huggingface_hub opencv-python-headless

# Copy app
WORKDIR /app
COPY app.py /app/app.py

RUN python3.10 -m pip uninstall -y gradio gradio_client || true
RUN printf "numpy==1.26.4\n" > /tmp/constraints.txt
RUN python3.10 -m pip install -c /tmp/constraints.txt --no-cache-dir --force-reinstall "gradio>=4.44,<6" "gradio_client>=0.8"
RUN python3.10 -c "import gradio as gr; print('gradio in image =', gr.__version__)"

EXPOSE 7860
CMD ["python3.10", "/app/app.py"]