add torch2.7 and 2.8 dockers

Dockerfile → docker/Dockerfile

@@ -27,26 +27,27 @@ RUN 'tzdata tzdata/Areas select America' | debconf-set-selections \
     && apt update -y \
     && apt install software-properties-common -y \
     && add-apt-repository ppa:deadsnakes/ppa -y \
-    && apt update
-
-RUN apt install python${PYTHON_VERSION} python${PYTHON_VERSION}-dev g++-11 gcc-11 -y \
-    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
-    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} && apt install python${PYTHON_VERSION}-distutils -y \
-    && update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 \
-    && update-alternatives --set python /usr/bin/python${PYTHON_VERSION}
-
-RUN apt install curl git sudo libibverbs-dev -y \
-    && apt install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \
-    && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py \
-    && python3 --version \
-    && python3 -m pip --version \
-    && rm -rf /var/lib/apt/lists/* \
-    && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 1 && update-alternatives --set gcc /usr/bin/gcc-11 \
-    && update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-11 1 && update-alternatives --set g++ /usr/bin/g++-11 \
-    && apt clean
+    && apt update && apt install wget git -y && apt clean
+
+# Install Miniconda
+ENV CONDA_DIR=/opt/conda
+RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh \
+    && bash /tmp/miniconda.sh -b -p ${CONDA_DIR} \
+    && rm /tmp/miniconda.sh \
+    && ${CONDA_DIR}/bin/conda clean -afy
+
+ENV PATH=${CONDA_DIR}/bin:$PATH
+
+RUN conda init bash
+
+RUN conda create -y -n nunchaku python=${PYTHON_VERSION} \
+    && conda install -y -n nunchaku -c conda-forge gxx=11 gcc=11 \
+    && conda clean -afy
+
+SHELL ["conda", "run", "-n", "nunchaku", "/bin/bash", "-c"]
 
 # Install building dependencies
-RUN pip install torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} torchaudio==${TORCHAUDIO_VERSION} --index-url https://download.pytorch.org/whl/cu${CUDA_SHORT_VERSION}
+RUN pip install torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} torchaudio==${TORCHAUDIO_VERSION} --index-url https://download.pytorch.org/whl/cu124
 RUN pip install ninja wheel diffusers transformers accelerate sentencepiece protobuf huggingface_hub comfy-cli
 
 # Start building
@@ -57,6 +58,7 @@ RUN git clone https://github.com/mit-han-lab/nunchaku.git \
     && NUNCHAKU_INSTALL_MODE=ALL python setup.py develop
 
 RUN cd .. && git clone https://github.com/comfyanonymous/ComfyUI \
-    && cd ComfyUI/custom_nodes && git clone https://github.com/ltdrdata/ComfyUI-Manager comfyui-manager \
-    && git clone https://github.com/mit-han-lab/ComfyUI-nunchaku.git nunchaku_nodes \
-    && cd .. && mkdir -p user/default/workflows/ && cp custom_nodes/nunchaku_nodes/workflows/* user/default/workflows/
+    && cd ComfyUI && pip install -r requirements.txt \
+    && cd custom_nodes && git clone https://github.com/ltdrdata/ComfyUI-Manager comfyui-manager \
+    && git clone https://github.com/mit-han-lab/ComfyUI-nunchaku.git \
+    && cd .. && mkdir -p user/default/workflows/ && cp -r custom_nodes/ComfyUI-nunchaku/workflows/ user/default/workflows/nunchaku_examples

docker/Dockerfile.torch27

+# Use an NVIDIA base image with CUDA support
+
+ARG CUDA_IMAGE="12.8.1-devel-ubuntu24.04"
+
+FROM nvidia/cuda:${CUDA_IMAGE}
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+ARG PYTHON_VERSION=3.11
+ARG TORCH_VERSION=2.7
+ARG TORCHVISION_VERSION=0.21
+ARG TORCHAUDIO_VERSION=2.6
+ARG CUDA_SHORT_VERSION=12.8
+
+# Set working directory
+WORKDIR /
+
+RUN echo PYTHON_VERSION=${PYTHON_VERSION} \
+    && echo CUDA_SHORT_VERSION=${CUDA_SHORT_VERSION} \
+    && echo TORCH_VERSION=${TORCH_VERSION} \
+    && echo TORCHVISION_VERSION=${TORCHVISION_VERSION} \
+    && echo TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION}
+
+# Setup timezone and install system dependencies
+RUN 'tzdata tzdata/Areas select America' | debconf-set-selections \
+    && echo 'tzdata tzdata/Zones/America select New_York' | debconf-set-selections \
+    && apt update -y \
+    && apt install software-properties-common -y \
+    && add-apt-repository ppa:deadsnakes/ppa -y \
+    && apt update && apt install wget git -y && apt clean
+
+# Install Miniconda
+ENV CONDA_DIR=/opt/conda
+RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh \
+    && bash /tmp/miniconda.sh -b -p ${CONDA_DIR} \
+    && rm /tmp/miniconda.sh \
+    && ${CONDA_DIR}/bin/conda clean -afy
+
+ENV PATH=${CONDA_DIR}/bin:$PATH
+
+RUN conda init bash
+
+RUN conda create -y -n nunchaku python=${PYTHON_VERSION} \
+    && conda install -y -n nunchaku -c conda-forge gxx=11 gcc=11 \
+    && conda clean -afy
+
+SHELL ["conda", "run", "-n", "nunchaku", "/bin/bash", "-c"]
+
+# Install building dependencies
+RUN pip install --pre torch==2.7.0.dev20250307+cu128 torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+RUN pip install ninja wheel diffusers transformers accelerate sentencepiece protobuf huggingface_hub comfy-cli
+
+# Start building
+RUN git clone https://github.com/mit-han-lab/nunchaku.git \
+    && cd nunchaku \
+    && git submodule init \
+    && git submodule update \
+    && NUNCHAKU_INSTALL_MODE=ALL python setup.py develop
+
+RUN cd .. && git clone https://github.com/comfyanonymous/ComfyUI \
+    && cd ComfyUI && pip install -r requirements.txt \
+    && cd custom_nodes && git clone https://github.com/ltdrdata/ComfyUI-Manager comfyui-manager \
+    && git clone https://github.com/mit-han-lab/ComfyUI-nunchaku.git \
+    && cd .. && mkdir -p user/default/workflows/ && cp -r custom_nodes/ComfyUI-nunchaku/workflows/ user/default/workflows/nunchaku_examples

docker/Dockerfile.torch28

+# Use an NVIDIA base image with CUDA support
+
+ARG CUDA_IMAGE="12.8.1-devel-ubuntu24.04"
+
+FROM nvidia/cuda:${CUDA_IMAGE}
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+ARG PYTHON_VERSION=3.11
+ARG TORCH_VERSION=2.8
+ARG TORCHVISION_VERSION=0.21
+ARG TORCHAUDIO_VERSION=2.6
+ARG CUDA_SHORT_VERSION=12.8
+
+# Set working directory
+WORKDIR /
+
+RUN echo PYTHON_VERSION=${PYTHON_VERSION} \
+    && echo CUDA_SHORT_VERSION=${CUDA_SHORT_VERSION} \
+    && echo TORCH_VERSION=${TORCH_VERSION} \
+    && echo TORCHVISION_VERSION=${TORCHVISION_VERSION} \
+    && echo TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION}
+
+# Setup timezone and install system dependencies
+RUN 'tzdata tzdata/Areas select America' | debconf-set-selections \
+    && echo 'tzdata tzdata/Zones/America select New_York' | debconf-set-selections \
+    && apt update -y \
+    && apt install software-properties-common -y \
+    && add-apt-repository ppa:deadsnakes/ppa -y \
+    && apt update && apt install wget git -y && apt clean
+
+# Install Miniconda
+ENV CONDA_DIR=/opt/conda
+RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh \
+    && bash /tmp/miniconda.sh -b -p ${CONDA_DIR} \
+    && rm /tmp/miniconda.sh \
+    && ${CONDA_DIR}/bin/conda clean -afy
+
+ENV PATH=${CONDA_DIR}/bin:$PATH
+
+RUN conda init bash
+
+RUN conda create -y -n nunchaku python=${PYTHON_VERSION} \
+    && conda install -y -n nunchaku -c conda-forge gxx=11 gcc=11 \
+    && conda clean -afy
+
+SHELL ["conda", "run", "-n", "nunchaku", "/bin/bash", "-c"]
+
+# Install building dependencies
+RUN pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+RUN pip install ninja wheel diffusers transformers accelerate sentencepiece protobuf huggingface_hub comfy-cli
+
+# Start building
+RUN git clone https://github.com/mit-han-lab/nunchaku.git \
+    && cd nunchaku \
+    && git submodule init \
+    && git submodule update \
+    && NUNCHAKU_INSTALL_MODE=ALL python setup.py develop
+
+RUN cd .. && git clone https://github.com/comfyanonymous/ComfyUI \
+    && cd ComfyUI && pip install -r requirements.txt \
+    && cd custom_nodes && git clone https://github.com/ltdrdata/ComfyUI-Manager comfyui-manager \
+    && git clone https://github.com/mit-han-lab/ComfyUI-nunchaku.git \
+    && cd .. && mkdir -p user/default/workflows/ && cp -r custom_nodes/ComfyUI-nunchaku/workflows/ user/default/workflows/nunchaku_examples

scripts/build_all_dockers.sh

+#!/bin/bash
+# Define the versions for Python, Torch, and CUDA
+NUNCHAKU_VERSION=$1
+python_versions=("3.10" "3.11" "3.12")
+torch_versions=("2.5" "2.6")
+cuda_versions=("12.4" "12.8")
+
+# Loop through all combinations of Python, Torch, and CUDA versions
+for python_version in "${python_versions[@]}"; do
+  for torch_version in "${torch_versions[@]}"; do
+    # Skip building for Python 3.13 and PyTorch 2.5
+    if [[ "$python_version" == "3.13" && "$torch_version" == "2.5" ]]; then
+      echo "Skipping Python 3.13 with PyTorch 2.5"
+      continue
+    fi
+    for cuda_version in "${cuda_versions[@]}"; do
+      bash scripts/build_docker.sh "$python_version" "$torch_version" "$cuda_version" "$NUNCHAKU_VERSION"
+    done
+  done
+done
+
+for python_version in "${python_versions[@]}"; do
+  for cuda_version in "${cuda_versions[@]}"; do
+    bash scripts/build_docker_torch27.sh "$python_version" "2.7" "$cuda_version" "$NUNCHAKU_VERSION"
+    bash scripts/build_docker_torch28.sh "$python_version" "2.8" "$cuda_version" "$NUNCHAKU_VERSION"
+  done
+done

scripts/build_docker.sh

@@ -29,12 +29,14 @@ else
   exit 2
 fi
 
-
-docker build --no-cache \
+docker build -f docker/Dockerfile --no-cache \
 --build-arg PYTHON_VERSION=${PYTHON_VERSION} \
 --build-arg CUDA_SHORT_VERSION=${CUDA_VERSION//.} \
 --build-arg CUDA_IMAGE=${CUDA_IMAGE} \
 --build-arg TORCH_VERSION=${TORCH_VERSION} \
 --build-arg TORCHVISION_VERSION=${TORCHVISION_VERSION} \
 --build-arg TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION} \
--t nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION} .
+-t lmxyy/nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION} .
+
+docker push lmxyy/nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION}
+docker rmi lmxyy/nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION}
\ No newline at end of file

scripts/build_docker_torch27.sh

+#!/bin/bash
+PYTHON_VERSION=$1
+TORCH_VERSION=$2
+CUDA_VERSION=$3
+NUNCHAKU_VERSION=$4
+TORCHVISION_VERSION=""
+TORCHAUDIO_VERSION=""
+
+if [ "$CUDA_VERSION" == "12.8" ]; then
+  CUDA_IMAGE="12.8.1-devel-ubuntu24.04"
+  echo "CUDA_VERSION is 12.8, setting CUDA_IMAGE to $CUDA_IMAGE"
+elif [ "$CUDA_VERSION" == "12.4" ]; then
+  CUDA_IMAGE="12.4.1-devel-ubuntu22.04"
+  echo "CUDA_VERSION is 12.4, setting CUDA_IMAGE to $CUDA_IMAGE"
+else
+  echo "CUDA_VERSION is not 12.8 or 12.4. Exit."
+  exit 2
+fi
+
+docker build -f docker/Dockerfile.torch27 --no-cache \
+--build-arg PYTHON_VERSION=${PYTHON_VERSION} \
+--build-arg CUDA_SHORT_VERSION=${CUDA_VERSION//.} \
+--build-arg CUDA_IMAGE=${CUDA_IMAGE} \
+--build-arg TORCH_VERSION=${TORCH_VERSION} \
+--build-arg TORCHVISION_VERSION=${TORCHVISION_VERSION} \
+--build-arg TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION} \
+-t lmxyy/nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION} .
+
+docker push lmxyy/nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION}
\ No newline at end of file

scripts/build_docker_torch28.sh

+#!/bin/bash
+PYTHON_VERSION=$1
+TORCH_VERSION=$2
+CUDA_VERSION=$3
+NUNCHAKU_VERSION=$4
+
+# Check if TORCH_VERSION is 2.5 or 2.6 and set the corresponding versions for TORCHVISION and TORCHAUDIO
+if [ "$TORCH_VERSION" == "2.5" ]; then
+  TORCHVISION_VERSION="0.20"
+  TORCHAUDIO_VERSION="2.5"
+  echo "TORCH_VERSION is 2.5, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
+elif [ "$TORCH_VERSION" == "2.6" ]; then
+  TORCHVISION_VERSION="0.21"
+  TORCHAUDIO_VERSION="2.6"
+  echo "TORCH_VERSION is 2.6, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
+else
+  echo "TORCH_VERSION is not 2.5 or 2.6. Exit."
+  exit 2
+fi
+
+if [ "$CUDA_VERSION" == "12.8" ]; then
+  CUDA_IMAGE="12.8.1-devel-ubuntu24.04"
+  echo "CUDA_VERSION is 12.8, setting CUDA_IMAGE to $CUDA_IMAGE"
+elif [ "$CUDA_VERSION" == "12.4" ]; then
+  CUDA_IMAGE="12.4.1-devel-ubuntu22.04"
+  echo "CUDA_VERSION is 12.4, setting CUDA_IMAGE to $CUDA_IMAGE"
+else
+  echo "CUDA_VERSION is not 12.8 or 12.4. Exit."
+  exit 2
+fi
+
+docker build -f docker/Dockerfile.torch28 --no-cache \
+--build-arg PYTHON_VERSION=${PYTHON_VERSION} \
+--build-arg CUDA_SHORT_VERSION=${CUDA_VERSION//.} \
+--build-arg CUDA_IMAGE=${CUDA_IMAGE} \
+--build-arg TORCH_VERSION=${TORCH_VERSION} \
+--build-arg TORCHVISION_VERSION=${TORCHVISION_VERSION} \
+--build-arg TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION} \
+-t lmxyy/nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION} .
+
+docker push lmxyy/nunchaku:${NUNCHAKU_VERSION}-py${PYTHON_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION}
\ No newline at end of file