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

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 python3-pip python3.10-dev \
    git git-lfs ffmpeg libsm6 libxext6 libgl1 libegl1 \
    build-essential ninja-build cmake rsync \
    && rm -rf /var/lib/apt/lists/* \
    && git lfs install

RUN useradd -m -u 1000 user
USER user

ENV HOME=/home/user \
    PATH=/home/user/.local/bin:$PATH \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    HF_HOME=/home/user/.cache/huggingface \
    TORCH_HOME=/home/user/.cache/torch \
    ATTN_BACKEND=xformers \
    SPARSE_ATTN_BACKEND=xformers \
    TORCH_CUDA_ARCH_LIST="7.5;8.6;8.9"

WORKDIR $HOME/app

COPY --chown=user:user . $HOME/app

RUN python3.10 -m pip install --upgrade pip setuptools wheel

RUN python3.10 -m pip install \
    torch==2.4.0 torchvision==0.19.0 \
    --index-url https://download.pytorch.org/whl/cu121

RUN python3.10 -m pip install \
    pillow imageio imageio-ffmpeg tqdm easydict scipy ninja psutil safetensors \
    scikit-learn opencv-python-headless rembg onnxruntime \
    trimesh xatlas pyvista pymeshfix igraph pygltflib geffnet \
    transformers \
    gradio==4.44.1 gradio_litmodel3d==0.0.1 "huggingface_hub<0.25"

RUN python3.10 -m pip install \
    git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8

RUN python3.10 -m pip install \
    "pytorch3d==0.7.8" \
    --find-links https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt240/download.html

RUN python3.10 -m pip install \
    xformers==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu121

RUN python3.10 -m pip install spconv-cu121

RUN python3.10 -m pip install \
    kaolin -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.4.0_cu121.html

RUN python3.10 -m pip install \
    "git+https://github.com/NVlabs/nvdiffrast.git" --no-build-isolation

EXPOSE 7860

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