chore: support torch2.10 and more cuda versions (#890)

lmxyy · web-flow · commit a0c4bfc9df47 · 2026-01-24T17:14:31.000-08:00
* add torch2.10, eliminate 2.7

* add torch2.10, eliminate 2.7

* done

* bump the version

* debugging
diff --git a/.github/workflows/nightly-build.yaml b/.github/workflows/nightly-build.yaml
@@ -76,8 +76,15 @@ jobs:
     if: needs.tag.outputs.need_build == 'true' && github.repository == 'nunchaku-ai/nunchaku'
     strategy:
       matrix:
-        python: ["3.10", "3.11", "3.12", "3.13"]
-        torch: ["2.7", "2.8", "2.9", "pre"]
+        # python: ["3.10", "3.11", "3.12", "3.13"]
+        # torch: ["2.8", "2.9", "2.10", "pre"]
+        # cuda: ["12.8", "13.0"]
+        python: ["3.10"]
+        torch: ["2.8"]
+        cuda: ["12.8", "13.0"]
+        exclude:
+          - torch: "2.8"
+            cuda: "13.0"
     steps:
       - name: Checkout to the tag
         uses: actions/checkout@v4
@@ -89,17 +96,10 @@ jobs:
         run: git log -1 --oneline
       - name: Build wheels
         run: |
-          if [[ "${{ matrix.torch }}" == "2.7" || "${{ matrix.torch }}" == "2.8" ]]; then
-            cuda_version="12.8"
-          elif [[ "${{ matrix.torch }}" == "2.9" ]]; then
-            cuda_version="13.0"
-          else
-            cuda_version="12.4"
-          fi
           if [[ "${{ matrix.torch }}" == "pre" ]]; then
-            bash scripts/build_linux_wheel_torch_nightly.sh ${{ matrix.python }} ${{ matrix.torch }} 13.0
+            bash scripts/build_linux_wheel_torch_nightly.sh ${{ matrix.python }} ${{ matrix.torch }} ${{ matrix.cuda }}
           else
-            bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
+            bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} ${{ matrix.cuda }}
           fi
       - name: Upload wheels to GitHub Release
         uses: softprops/action-gh-release@v2
@@ -120,8 +120,15 @@ jobs:
     if: needs.tag.outputs.need_build == 'true' && github.repository == 'nunchaku-tech/nunchaku'
     strategy:
       matrix:
-        python: ["3.10", "3.11", "3.12", "3.13"]
-        torch: ["2.9", "pre"]
+        # python: ["3.10", "3.11", "3.12", "3.13"]
+        # torch: ["2.8", "2.9", "2.10", "pre"]
+        # cuda: ["12.8", "13.0"]
+        python: ["3.10"]
+        torch: ["2.8"]
+        cuda: ["12.8", "13.0"]
+        exclude:
+          - torch: "2.8"
+            cuda: "13.0"
     steps:
       - name: Checkout to the tag
         uses: actions/checkout@v4
@@ -134,24 +141,11 @@ jobs:
       - name: Build wheels
         shell: cmd
         run: |
-          SET TORCH_VERSION=${{ matrix.torch }}
-          SET PYTHON_VERSION=${{ matrix.python }}
-          IF "%TORCH_VERSION%"=="2.7" (
-            SET CUDA_VERSION=12.8
-          ) ELSE IF "%TORCH_VERSION%"=="2.8" (
-            SET CUDA_VERSION=12.8
-          ) ELSE IF "%TORCH_VERSION%"=="2.9" (
-            SET CUDA_VERSION=13.0
-          ) ELSE IF "%TORCH_VERSION%"=="pre" (
-            SET CUDA_VERSION=13.0
-          ) ELSE (
-            SET CUDA_VERSION=12.4
-          )
-          call C:\Users\muyangl\miniconda3\condabin\activate.bat activate
-          IF "%TORCH_VERSION%"=="pre" (
-            call scripts\build_windows_wheel_torch_nightly.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
+          call C:\Users\nunchaku\miniconda3\condabin\activate.bat activate
+          IF "${{ matrix.torch }}"=="pre" (
+            call scripts\build_windows_wheel_torch_nightly.cmd ${{ matrix.python }} ${{ matrix.torch }} ${{ matrix.cuda }}
           ) ELSE (
-            call scripts\build_windows_wheel.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
+            call scripts\build_windows_wheel.cmd ${{ matrix.python }} ${{ matrix.torch }} ${{ matrix.cuda }}
           )
       - name: Upload wheels to GitHub Release
         uses: softprops/action-gh-release@v2
diff --git a/.github/workflows/release-build.yaml b/.github/workflows/release-build.yaml
@@ -51,7 +51,11 @@ jobs:
     strategy:
       matrix:
         python: ["3.10", "3.11", "3.12", "3.13"]
-        torch: ["2.7", "2.8", "2.9", "pre"]
+        torch: ["2.8", "2.9", "2.10", "pre"]
+        cuda: ["12.8", "13.0"]
+        exclude:
+          - torch: "2.8"
+            cuda: "13.0"
     steps:
       - name: Checkout to the tag
         uses: actions/checkout@v4
@@ -63,17 +67,10 @@ jobs:
         run: git log -1 --oneline
       - name: Build wheels
         run: |
-          if [[ "${{ matrix.torch }}" == "2.7" || "${{ matrix.torch }}" == "2.8" ]]; then
-            cuda_version="12.8"
-          elif [[ "${{ matrix.torch }}" == "2.9" ]]; then
-            cuda_version="13.0"
-          else
-            cuda_version="12.4"
-          fi
           if [[ "${{ matrix.torch }}" == "pre" ]]; then
-            bash scripts/build_linux_wheel_torch_nightly.sh ${{ matrix.python }} ${{ matrix.torch }} 13.0
+            bash scripts/build_linux_wheel_torch_nightly.sh ${{ matrix.python }} ${{ matrix.torch }} ${{ matrix.cuda }}
           else
-            bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} $cuda_version
+            bash scripts/build_linux_wheel.sh ${{ matrix.python }} ${{ matrix.torch }} ${{ matrix.cuda }}
           fi
       - name: Upload wheels to GitHub Release
         uses: softprops/action-gh-release@v2
@@ -106,7 +103,11 @@ jobs:
     strategy:
       matrix:
         python: ["3.10", "3.11", "3.12", "3.13"]
-        torch: ["2.9", "pre"]
+        torch: ["2.8", "2.9", "2.10", "pre"]
+        cuda: ["12.8", "13.0"]
+        exclude:
+          - torch: "2.8"
+            cuda: "13.0"
     steps:
       - name: Checkout to the tag
         uses: actions/checkout@v4
@@ -119,24 +120,11 @@ jobs:
       - name: Build wheels
         shell: cmd
         run: |
-          SET TORCH_VERSION=${{ matrix.torch }}
-          SET PYTHON_VERSION=${{ matrix.python }}
-          IF "%TORCH_VERSION%"=="2.7" (
-            SET CUDA_VERSION=12.8
-          ) ELSE IF "%TORCH_VERSION%"=="2.8" (
-            SET CUDA_VERSION=12.8
-          ) ELSE IF "%TORCH_VERSION%"=="2.9" (
-            SET CUDA_VERSION=13.0
-          ) ELSE IF "%TORCH_VERSION%"=="pre" (
-            SET CUDA_VERSION=13.0
-          ) ELSE (
-            SET CUDA_VERSION=12.4
-          )
-          call C:\Users\muyangl\miniconda3\condabin\activate.bat activate
-          IF "%TORCH_VERSION%"=="pre" (
-            call scripts\build_windows_wheel_torch_nightly.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
+          call C:\Users\nunchaku\miniconda3\condabin\activate.bat activate
+          IF "${{ matrix.torch }}"=="pre" (
+            call scripts\build_windows_wheel_torch_nightly.cmd ${{ matrix.python }} ${{ matrix.torch }} ${{ matrix.cuda }}
           ) ELSE (
-            call scripts\build_windows_wheel.cmd %PYTHON_VERSION% %TORCH_VERSION% %CUDA_VERSION%
+            call scripts\build_windows_wheel.cmd ${{ matrix.python }} ${{ matrix.torch }} ${{ matrix.cuda }}
           )
       - name: Upload wheels to GitHub Release
         uses: softprops/action-gh-release@v2
@@ -152,7 +140,7 @@ jobs:
     #   env:
     #     HF_TOKEN: ${{ secrets.HF_WHEEL_UPLOAD_TOKEN }}
     #   run: |
-    #     call C:\Users\muyangl\miniconda3\condabin\activate.bat activate
+    #     call C:\Users\nunchaku\miniconda3\condabin\activate.bat activate
     #     pip install huggingface_hub
     #     for %%w in (dist\*.whl) do (
     #       hf upload nunchaku-tech/nunchaku "%%w"
diff --git a/nunchaku/__version__.py b/nunchaku/__version__.py
@@ -1 +1 @@
-__version__ = "1.2.0"
+__version__ = "1.2.1dev"
diff --git a/scripts/build_linux_wheel.sh b/scripts/build_linux_wheel.sh
@@ -7,30 +7,21 @@ CUDA_VERSION=$3
 MAX_JOBS=${4:-} # optional
 PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}
 
-# 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"
-elif [ "$TORCH_VERSION" == "2.7" ]; then
-  TORCHVISION_VERSION="0.22"
-  TORCHAUDIO_VERSION="2.7"
-  echo "TORCH_VERSION is 2.7, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
-elif [ "$TORCH_VERSION" == "2.8" ]; then
+# Set the corresponding versions for TORCHVISION and TORCHAUDIO
+if [ "$TORCH_VERSION" == "2.8" ]; then
   TORCHVISION_VERSION="0.23"
   TORCHAUDIO_VERSION="2.8"
-  echo "TORCH_VERSION is 2.8, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
 elif [ "$TORCH_VERSION" == "2.9" ]; then
   TORCHVISION_VERSION="0.24"
   TORCHAUDIO_VERSION="2.9"
-  echo "TORCH_VERSION is 2.9, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
+elif [ "$TORCH_VERSION" == "2.10" ]; then
+  TORCHVISION_VERSION="0.25"
+  TORCHAUDIO_VERSION="2.10"
 else
-  echo "TORCH_VERSION is not 2.5, 2.6, 2.7, 2.8 or 2.9, no changes to versions."
+  echo "Unsupported TORCH_VERSION: $TORCH_VERSION"
+  exit 1
 fi
+echo "TORCH_VERSION is $TORCH_VERSION, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
 
 docker run --rm \
     -v "$(pwd)":/nunchaku \
@@ -39,8 +30,9 @@ docker run --rm \
     cd /nunchaku && \
     rm -rf build && \
     gcc --version && g++ --version && \
-    ${PYTHON_ROOT_PATH}/bin/pip install --no-cache-dir torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} torchaudio==${TORCHAUDIO_VERSION} --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//.} && \
-    ${PYTHON_ROOT_PATH}/bin/pip install build ninja wheel setuptools && \
+    ${PYTHON_ROOT_PATH}/bin/pip install uv && \
+    ${PYTHON_ROOT_PATH}/bin/uv pip install --no-cache-dir torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} torchaudio==${TORCHAUDIO_VERSION} --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//.} && \
+    ${PYTHON_ROOT_PATH}/bin/uv pip install build ninja wheel setuptools && \
     export NUNCHAKU_INSTALL_MODE=ALL && \
     export NUNCHAKU_BUILD_WHEELS=1 && \
     export MAX_JOBS=${MAX_JOBS} && \
diff --git a/scripts/build_linux_wheel_torch_nightly.sh b/scripts/build_linux_wheel_torch_nightly.sh
@@ -22,8 +22,9 @@ docker run --rm \
     cd /nunchaku && \
     rm -rf build && \
     gcc --version && g++ --version && \
-    ${PYTHON_ROOT_PATH}/bin/pip install --pre --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu${CUDA_VERSION//.} && \
-    ${PYTHON_ROOT_PATH}/bin/pip install build ninja wheel setuptools && \
+    ${PYTHON_ROOT_PATH}/bin/pip install uv && \
+    ${PYTHON_ROOT_PATH}/bin/uv pip install --pre --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu${CUDA_VERSION//.} && \
+    ${PYTHON_ROOT_PATH}/bin/uv pip install build ninja wheel setuptools && \
     export NUNCHAKU_INSTALL_MODE=ALL && \
     export NUNCHAKU_BUILD_WHEELS=1 && \
     export MAX_JOBS=${MAX_JOBS} && \
diff --git a/scripts/build_windows_wheel.cmd b/scripts/build_windows_wheel.cmd
@@ -10,23 +10,18 @@ set CUDA_SHORT_VERSION=%CUDA_VERSION:.=%
 echo %CUDA_SHORT_VERSION%
 
 :: setup some variables
-if "%TORCH_VERSION%"=="2.5" (
-    set TORCHVISION_VERSION=0.20
-    set TORCHAUDIO_VERSION=2.5
-) else if "%TORCH_VERSION%"=="2.6" (
-    set TORCHVISION_VERSION=0.21
-    set TORCHAUDIO_VERSION=2.6
-) else if "%TORCH_VERSION%"=="2.7" (
-    set TORCHVISION_VERSION=0.22
-    set TORCHAUDIO_VERSION=2.7
-) else if "%TORCH_VERSION%"=="2.8" (
+if "%TORCH_VERSION%"=="2.8" (
     set TORCHVISION_VERSION=0.23
     set TORCHAUDIO_VERSION=2.8
 ) else if "%TORCH_VERSION%"=="2.9" (
     set TORCHVISION_VERSION=0.24
     set TORCHAUDIO_VERSION=2.9
+) else if "%TORCH_VERSION%"=="2.10" (
+    set TORCHVISION_VERSION=0.25
+    set TORCHAUDIO_VERSION=2.10
 ) else (
-    echo TORCH_VERSION is not 2.5, 2.6, 2.7, 2.8 or 2.9, no changes to versions.
+    echo Unsupported TORCH_VERSION: %TORCH_VERSION%
+    exit /b 1
 )
 echo setting TORCHVISION_VERSION to %TORCHVISION_VERSION% and TORCHAUDIO_VERSION to %TORCHAUDIO_VERSION%
 
@@ -39,19 +34,22 @@ call conda create -y -n %ENV_NAME% python=%PYTHON_VERSION%
 call conda activate %ENV_NAME%
 
 :: install dependencies
-call pip install ninja setuptools wheel build
-call pip install --no-cache-dir torch==%TORCH_VERSION% torchvision==%TORCHVISION_VERSION% torchaudio==%TORCHAUDIO_VERSION% --index-url "https://download.pytorch.org/whl/cu%CUDA_SHORT_VERSION%/"
+call pip install uv
+call uv pip install ninja setuptools wheel build
+call uv pip install --no-cache-dir torch==%TORCH_VERSION% torchvision==%TORCHVISION_VERSION% torchaudio==%TORCHAUDIO_VERSION% --index-url "https://download.pytorch.org/whl/cu%CUDA_SHORT_VERSION%/"
 
 :: set environment variables
 set NUNCHAKU_INSTALL_MODE=ALL
 set NUNCHAKU_BUILD_WHEELS=1
+set NVCC_PREPEND_FLAGS=-allow-unsupported-compiler
+set CUDA_HOME=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION%
 
 :: cd to the parent directory
 cd /d "%~dp0.."
 if exist build rd /s /q build
 
 :: set up Visual Studio compilation environment
-call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
+call "C:\Program Files (x86)\Microsoft Visual Studio\18\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
 set DISTUTILS_USE_SDK=1
 
 :: build wheels
diff --git a/scripts/build_windows_wheel_torch_nightly.cmd b/scripts/build_windows_wheel_torch_nightly.cmd
@@ -18,20 +18,22 @@ call conda create -y -n %ENV_NAME% python=%PYTHON_VERSION%
 call conda activate %ENV_NAME%
 
 :: install dependencies
-call pip install ninja setuptools wheel build
-
-call pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu%CUDA_SHORT_VERSION%
+call pip install uv
+call uv pip install ninja setuptools wheel build
+call uv pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu%CUDA_SHORT_VERSION%
 
 :: set environment variables
 set NUNCHAKU_INSTALL_MODE=ALL
 set NUNCHAKU_BUILD_WHEELS=1
+set NVCC_PREPEND_FLAGS=-allow-unsupported-compiler
+set CUDA_HOME=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION%
 
 :: cd to the parent directory
 cd /d "%~dp0.."
 if exist build rd /s /q build
 
 :: set up Visual Studio compilation environment
-call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
+call "C:\Program Files (x86)\Microsoft Visual Studio\18\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
 set DISTUTILS_USE_SDK=1
 
 :: build wheels
diff --git a/setup.py b/setup.py
@@ -70,9 +70,11 @@ def get_sm_targets() -> list[str]:
 
     torch_version = torch.__version__.split("+")[0]
     torch_major_minor_version = ".".join(torch_version.split(".")[:2])
+    cuda_version = torch.version.cuda  # e.g., "12.4"
     if "dev" in version:
-        version = version + date.today().strftime("%Y%m%d")  # data
-    version = version + "+torch" + torch_major_minor_version
+        version = version + date.today().strftime("%Y%m%d")
+    # Version format: 1.2.0+cu12.4torch2.4
+    version = f"{version}+cu{cuda_version}torch{torch_major_minor_version}"
 
     ROOT_DIR = os.path.dirname(__file__)
 

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-__version__ = "1.2.0"`
	`1`	`+__version__ = "1.2.1dev"`