dockerfile

# syntax=docker/dockerfile:1

ARG CUDA_VERSION=12.8.1
ARG TORCH_CUDA_ARCH_LIST="8.6 8.9 12.0+PTX"


# =============================================================================
# Stage 1: Build Conda environment and compile CUDA extensions
# =============================================================================
FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu22.04 AS builder

ARG DEBIAN_FRONTEND=noninteractive
ARG TORCH_CUDA_ARCH_LIST

ENV CF3D_ENV=/opt/conda/envs/controlfix3d

ENV PATH=/usr/local/cuda/bin:${CF3D_ENV}/bin:/opt/conda/bin:${PATH} \
    CUDA_HOME=/usr/local/cuda \
    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    PIP_DISABLE_PIP_VERSION_CHECK=1

# Only builder dependencies are installed here.
# They will not be copied into the final image.
RUN apt-get update \
    && apt-get install -y --no-install-recommends \
        build-essential \
        ca-certificates \
        git \
        ninja-build \
        wget \
    && rm -rf /var/lib/apt/lists/*

# Install Miniforge.
RUN wget -qO /tmp/miniforge.sh \
        https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh \
    && bash /tmp/miniforge.sh -b -p /opt/conda \
    && rm -f /tmp/miniforge.sh \
    && conda clean -afy

WORKDIR /build/ControlFix3D

# Copy dependency definitions and CUDA extension sources first,
# preserving the Docker cache when application source changes.
COPY environment.yml ./
COPY submodules/ ./submodules/

# Remove:
#   1. The environment name because --prefix is used.
#   2. Local CUDA extensions because they must be installed after PyTorch.
#
# conda clean is performed in the same layer so the package cache does not
# remain inside Docker image history.
RUN sed \
        -e '/^name:/d' \
        -e '\|submodules/diff-gaussian-rasterization|d' \
        -e '\|submodules/simple-knn|d' \
        -e '\|submodules/fused-ssim|d' \
        environment.yml > /tmp/environment.docker.yml \
    && conda env create \
        --prefix "${CF3D_ENV}" \
        --file /tmp/environment.docker.yml \
    && rm -f /tmp/environment.docker.yml \
    && conda clean -afy \
    && rm -rf /root/.cache /tmp/*

# Verify that the expected PyTorch CUDA build was installed.
RUN python -c \
        "import torch; \
        print('PyTorch:', torch.__version__); \
        print('PyTorch CUDA:', torch.version.cuda)" \
    && nvcc --version

# Compile all CUDA extensions against:
#   - PyTorch cu128
#   - CUDA Toolkit 12.8
#   - sm_86, sm_89 and sm_120
RUN python -m pip install \
        --no-cache-dir \
        --no-build-isolation \
        ./submodules/diff-gaussian-rasterization \
        ./submodules/simple-knn \
        ./submodules/fused-ssim \
    && rm -rf /root/.cache /tmp/*

# The Streamlit UI.  Deliberately NOT in environment.yml: touching that file
# invalidates the conda-env and CUDA-extension layers above, an hour of rebuild for a
# pure-python dependency.  `iproute2` supplies `ss`, which run_app.sh uses for its
# --status / --stop / --daemon paths; without it --daemon always reports "failed to
# start" even when the server is up.
RUN apt-get update && apt-get install -y --no-install-recommends iproute2 \
    && rm -rf /var/lib/apt/lists/* \
    && conda run --prefix "${CF3D_ENV}" python -m pip install streamlit==1.62.0 \
    && conda run --prefix "${CF3D_ENV}" python -c "import streamlit; print('streamlit:', streamlit.__version__)"

# Import verification does not require a GPU at build time.
RUN python -c \
    "import torch; \
     import streamlit; \
     import diff_gaussian_rasterization; \
     import simple_knn; \
     import fused_ssim; \
     print('PyTorch:', torch.__version__); \
     print('CUDA:', torch.version.cuda); \
     print('Streamlit:', streamlit.__version__); \
     print('CUDA extensions: OK')" \
    && huggingface-cli --help > /dev/null


# =============================================================================
# Stage 2: Runtime image
#
# Contains CUDA runtime and cuDNN, but does not contain:
#   - nvcc
#   - CUDA headers
#   - build-essential
#   - ninja
#   - git
#   - Miniforge base environment
#   - Conda package cache
# =============================================================================
FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu22.04 AS runtime

ARG DEBIAN_FRONTEND=noninteractive
ARG TORCH_CUDA_ARCH_LIST

ENV CF3D_ENV=/opt/conda/envs/controlfix3d

ENV PATH=${CF3D_ENV}/bin:${PATH} \
    CUDA_HOME=/usr/local/cuda \
    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    HF_HOME=/models/huggingface \
    TORCH_HOME=/models/torch \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    STREAMLIT_SERVER_PORT=8531 \
    STREAMLIT_SERVER_ADDRESS=0.0.0.0 \
    STREAMLIT_BROWSER_GATHER_USAGE_STATS=false

# Runtime-only system dependencies.
RUN apt-get update \
    && apt-get install -y --no-install-recommends \
        ca-certificates \
        ffmpeg \
        iproute2 \
        libgl1 \
        libglib2.0-0 \
        libsm6 \
        libxext6 \
        libxrender1 \
    && rm -rf /var/lib/apt/lists/* \
    && rm -rf /root/.cache /tmp/*

# Copy only the completed Conda environment.
#
# The path is kept identical to the builder path because Conda environments
# and executable shebangs can contain absolute paths.
COPY --from=builder \
    /opt/conda/envs/controlfix3d \
    /opt/conda/envs/controlfix3d

WORKDIR /app/ControlFix3D

# Application code is copied last so source changes do not rebuild dependencies.
COPY . /app/ControlFix3D

RUN chmod +x \
        /app/ControlFix3D/run_app.sh \
        /app/ControlFix3D/docker-entrypoint.sh

# Verify imports in the actual runtime image.
# This only checks loading/linking; it does not require a GPU during docker build.
RUN python -c \
    "import torch; \
     import streamlit; \
     import diff_gaussian_rasterization; \
     import simple_knn; \
     import fused_ssim; \
     print('Runtime PyTorch:', torch.__version__); \
     print('Runtime CUDA:', torch.version.cuda); \
     print('Runtime Streamlit:', streamlit.__version__); \
     print('Runtime CUDA extensions: OK')"

RUN mkdir -p \
    /models/huggingface \
    /models/torch

EXPOSE 8531

# Uncomment if docker-entrypoint.sh performs required initialization.
# ENTRYPOINT ["./docker-entrypoint.sh"]
# CMD ["./run_app.sh", "--lan"]