# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

# SPDX-License-Identifier: Apache-2.0

# CUDA 13 x86 — Blackwell (GB200, sm_100 / sm_121) with full sm_121 support

FROM nvcr.io/nvidia/pytorch:25.04-py3

ENV DEBIAN_FRONTEND=noninteractive

SHELL ["/bin/bash", "-o", "pipefail", "-c"]

ARG FLASH_ATTN_MAX_JOBS=16

ARG FLASH_ATTN_NVCC_THREADS=2

# Install CUDA 13 toolkit (base image ships CUDA 12.9; torch cu130 needs CUDA 13 nvcc)

RUN wget -qO /tmp/cuda-keyring.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb \

&& dpkg -i /tmp/cuda-keyring.deb && rm /tmp/cuda-keyring.deb

RUN apt-get update && apt-get install -y -qq --no-install-recommends \

wget git curl build-essential gcc-11 g++-11 libgl1-mesa-dev libglib2.0-0 cuda-toolkit-13-0 \

&& rm -rf /var/lib/apt/lists/*

ENV CUDA_HOME=/usr/local/cuda-13.0

ENV PATH=${CUDA_HOME}/bin:${PATH}

RUN truncate -s 0 /etc/pip/constraint.txt

RUN pip install torch==2.11.0 torchvision --index-url https://download.pytorch.org/whl/cu130

RUN pip uninstall -y flash-attn opencv-python || true

RUN rm -rf /usr/local/lib/python3.12/dist-packages/cv2*

WORKDIR /workspace

COPY ./thirdparty/ ./thirdparty/

WORKDIR ./thirdparty/3DGRUT-ArtiFixer/

RUN pip install -r requirements.txt

RUN bash scripts/install_slangc.sh /usr/local

RUN pip install -e .

# FA3 — Hopper kernels (warp specialization, async TMA, FP8)

# Filter ptxas output to stay under Docker's 2MB log limit per step

RUN git clone --depth 1 https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention && \

cd /tmp/flash-attention/hopper && \

MAX_JOBS=${FLASH_ATTN_MAX_JOBS} \

NVCC_THREADS=${FLASH_ATTN_NVCC_THREADS} \

pip install . 2>&1 | awk '!/^ptxas info/' && \

cd / && rm -rf /tmp/flash-attention

# FA4 — CuTeDSL kernels with [cu13] extra for sm_121 (GB200)

RUN pip install --pre "flash-attn-4[cu13]"

# flash-attn-4 upgrades cuda-python to 13.x namespace stub missing cuda.bindings.driver on x86;

# downgrade to monolithic 12.6.2 which bundles the compiled bindings

RUN pip install --force-reinstall --no-deps "cuda-python==12.6.2.post1"

RUN pip install accelerate==1.13.0 diffusers==0.37.1 transformers==5.5.0 ftfy

WORKDIR /workspace

RUN pip install einops scipy wandb tqdm Pillow matplotlib opencv-python-headless \

pyyaml torchmetrics imageio-ffmpeg h5py av torch-fidelity \

git+https://github.com/microsoft/MoGe.git

COPY tests/test_flash_attn.py /workspace/test_flash_attn.py

RUN python -c "import flash_attn_interface; print('FA3 ok')"

RUN python -c "import flash_attn_3; print('FA3 package ok')"

RUN pip show flash-attn-4 | grep Version && echo 'FA4 ok'

RUN python -c "import cv2; print('cv2 ok:', cv2.__version__)"

RUN python -c "from moge.model.v2 import MoGeModel; print('MoGe ok')"