From 82ca446812b43d1d9cb07a64e6a3188fd84b1984 Mon Sep 17 00:00:00 2001
From: Arison591 <2234819536@qq.com>
Date: Fri, 17 Jul 2026 18:03:22 +0800
Subject: [PATCH] Support Blackwell GPUs in setup

---
 AGENTS.md                     | 14 +++++--
 README.md                     |  3 +-
 setup.sh                      | 69 ++++++++++++++++++++++++++++++++--
 tests/test_setup_blackwell.sh | 70 +++++++++++++++++++++++++++++++++++
 4 files changed, 147 insertions(+), 9 deletions(-)
 create mode 100755 tests/test_setup_blackwell.sh

diff --git a/AGENTS.md b/AGENTS.md
index 2b633b6..6f71af4 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -65,7 +65,7 @@ If `CLAUDE.md` or `CODEX.md` exists, they should stay lightweight and only conta
 
 ### Prerequisites
 - **OS**: Linux only.
-- **GPU**: NVIDIA GPU with >= 18 GB VRAM. Tested on A800, RTX 3090.
+- **GPU**: NVIDIA GPU with >= 18 GB VRAM. Tested on A800, RTX 3090. The installer has a Blackwell/RTX 50-series compatibility path.
 - **CUDA Toolkit**: 11.8 or 12.2 (needed to compile extensions).
 - **Python**: 3.10+.
 
@@ -98,7 +98,7 @@ source ./setup.sh --train
 
 **Important**: `setup.sh` must be **sourced** (`. ./setup.sh` or `source ./setup.sh`), not executed directly (`bash setup.sh` will fail to activate the environment).
 
-The script auto-detects CUDA version and installs matching wheels for PyTorch, spconv, pytorch3d, and nvdiffrast. It prefers `uv` for fast installs and falls back to `pip`.
+The script auto-detects CUDA version and GPU compute capability and installs matching wheels for PyTorch, spconv, pytorch3d, and nvdiffrast. Blackwell GPUs select PyTorch 2.7 with CUDA 12.8. It prefers `uv` for fast installs and falls back to `pip`.
 
 ## Pretrained Model Weights
 
@@ -192,7 +192,7 @@ Sample data: [VAST-AI/AniGen_sample_data](https://huggingface.co/datasets/VAST-A
 
 | Dependency | Notes |
 |---|---|
-| PyTorch 2.4–2.5 | Auto-selected by `setup.sh` based on Python version |
+| PyTorch 2.4–2.5; 2.7 on Blackwell | Auto-selected by `setup.sh` based on Python version and GPU compute capability |
 | spconv | CUDA-version-matched (`spconv-cu118` or `spconv-cu121`) |
 | pytorch3d | Pre-built wheel or source build; requires `--no-build-isolation` |
 | nvdiffrast | Differentiable rasterization; compiled extension |
@@ -218,7 +218,13 @@ Common issues:
 
 ## Testing
 
-No dedicated test suite (research codebase). Verify correctness with:
+The research codebase has one lightweight installer regression test:
+
+```bash
+tests/test_setup_blackwell.sh
+```
+
+Verify model behavior with:
 
 1. **Primary smoke test** — launch the Gradio web demo:
    ```bash
diff --git a/README.md b/README.md
index 5cd93f1..b502eaa 100755
--- a/README.md
+++ b/README.md
@@ -94,7 +94,7 @@ Full guide: [`video_editing/README.md`](video_editing/README.md).
 
 ### Prerequisites
 - **System**: The code is currently tested only on **Linux**.
-- **Hardware**: An NVIDIA GPU with at least 18GB of memory is necessary. The code has been verified on NVIDIA A800 and RTX3090 GPUs.
+- **Hardware**: An NVIDIA GPU with at least 18GB of memory is necessary. The code has been verified on NVIDIA A800 and RTX3090 GPUs. The setup script also detects Blackwell GPUs such as the RTX 50 series and selects a compatible PyTorch build.
 - **Software**:
   - The [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit-archive) is needed to compile certain submodules. The code has been tested with CUDA versions 11.8 and 12.2.
   - [Conda](https://docs.anaconda.com/miniconda/install/#quick-command-line-install) is recommended for managing dependencies.
@@ -128,6 +128,7 @@ Full guide: [`video_editing/README.md`](video_editing/README.md).
 
     > [!NOTE]
     > The setup script auto-detects your CUDA version and installs matching wheels for PyTorch, spconv, pytorch3d, and nvdiffrast. [DSINE](https://github.com/baegwangbin/DSINE) (used for surface normal estimation) is loaded at runtime via `torch.hub` and does not require separate installation.
+    > On NVIDIA Blackwell GPUs, it selects PyTorch 2.7 with CUDA 12.8 and exports the detected compute capability for local extension builds. Set `ANIGEN_CUDA_CAPABILITY` only when GPU detection is unavailable, for example `ANIGEN_CUDA_CAPABILITY=12.0` for RTX 50-series GPUs.
 
 
 <!-- Pretrained Models -->
diff --git a/setup.sh b/setup.sh
index 9a1d82c..03a65a2 100755
--- a/setup.sh
+++ b/setup.sh
@@ -56,7 +56,9 @@ if [ "$HELP" = true ] || [ "$#" -eq 0 -a "$ALL" = false -a "$BASIC" = false -a "
     echo
     echo "Environment variables:"
     echo "  ANIGEN_PYTHON=/path/to/python   Use a specific Python interpreter"
-    echo "  TORCH_VERSION=2.4.0             Override PyTorch version (default: 2.4.0 for Python <=3.12, 2.5.0 for 3.13+)"
+    echo "  ANIGEN_CUDA_CAPABILITY=12.0     Override GPU compute capability detection"
+    echo "  TORCH_VERSION=2.4.0             Override PyTorch version (Blackwell requires >=2.7)"
+    echo "  ANIGEN_TORCH_INDEX_URL=URL      Override the PyTorch wheel index"
     return 0
 fi
 
@@ -125,6 +127,53 @@ _detect_python_version() {
     "$PYTHON_BIN" -c "import sys; print(f'{sys.version_info.major}{sys.version_info.minor}')" 2>/dev/null
 }
 
+_detect_compute_capability() {
+    if [ -n "${ANIGEN_CUDA_CAPABILITY:-}" ]; then
+        printf '%s\n' "$ANIGEN_CUDA_CAPABILITY"
+        return
+    fi
+    local capability=""
+    if command -v nvidia-smi >/dev/null 2>&1; then
+        capability=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \
+            | tr -d ' ' | sort -V | tail -1)
+    fi
+    if [ -z "$capability" ]; then
+        capability=$("$PYTHON_BIN" -c "import torch; caps = [torch.cuda.get_device_capability(i) for i in range(torch.cuda.device_count())]; print(f'{max(caps)[0]}.{max(caps)[1]}' if caps else '')" 2>/dev/null || true)
+    fi
+    printf '%s\n' "$capability"
+}
+
+CUDA_CAPABILITY=$(_detect_compute_capability)
+CUDA_CAPABILITY_MAJOR=$(echo "${CUDA_CAPABILITY}" | cut -d'.' -f1)
+BLACKWELL_GPU=false
+case "$CUDA_CAPABILITY_MAJOR" in
+    ''|*[!0-9]*) ;;
+    *)
+        if [ "$CUDA_CAPABILITY_MAJOR" -ge 10 ]; then
+            BLACKWELL_GPU=true
+        fi
+        ;;
+esac
+
+if [ "$BLACKWELL_GPU" = true ]; then
+    export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-$CUDA_CAPABILITY}"
+    FLASH_ARCH=$(echo "$CUDA_CAPABILITY" | tr -d '.')
+    export FLASH_ATTN_CUDA_ARCHS="${FLASH_ATTN_CUDA_ARCHS:-$FLASH_ARCH}"
+    echo "[SETUP] Detected NVIDIA Blackwell compute capability ${CUDA_CAPABILITY}; enabling SM${FLASH_ARCH} extension builds"
+fi
+
+_require_blackwell_torch() {
+    if [ "$BLACKWELL_GPU" != true ]; then
+        return 0
+    fi
+    case "$1" in
+        ''|0.*|1.*|2.[0-6]|2.[0-6].*)
+            echo "[ERROR] NVIDIA Blackwell (compute capability ${CUDA_CAPABILITY}) requires PyTorch >=2.7 with CUDA 12.8 support."
+            return 1
+            ;;
+    esac
+}
+
 # ─── Tsinghua mirror ────────────────────────────────────────────────────────
 
 if [ "$TSINGHUA" = true ]; then
@@ -175,10 +224,12 @@ if [ "$TORCH" = true ]; then
     fi
     CUDA_MAJOR=$(echo "${CUDA_VER}" | cut -d'.' -f1)
 
-    # Auto-select default PyTorch version based on Python version
+    # Auto-select a Blackwell-capable build before applying legacy defaults.
     PYVER_NUM=$("$PYTHON_BIN" -c "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "10")
     if [ -z "${TORCH_VERSION:-}" ]; then
-        if [ "$PYVER_NUM" -ge 13 ]; then
+        if [ "$BLACKWELL_GPU" = true ]; then
+            TORCH_VER="2.7.0"
+        elif [ "$PYVER_NUM" -ge 13 ]; then
             TORCH_VER="2.5.0"
         else
             TORCH_VER="2.4.0"
@@ -187,7 +238,13 @@ if [ "$TORCH" = true ]; then
         TORCH_VER="$TORCH_VERSION"
     fi
 
-    if [ "${CUDA_MAJOR}" = "12" ]; then
+    _require_blackwell_torch "$TORCH_VER" || return 1
+
+    if [ -n "${ANIGEN_TORCH_INDEX_URL:-}" ]; then
+        TORCH_INDEX="$ANIGEN_TORCH_INDEX_URL"
+    elif [ "$BLACKWELL_GPU" = true ]; then
+        TORCH_INDEX=https://download.pytorch.org/whl/cu128
+    elif [ "${CUDA_MAJOR}" = "12" ]; then
         TORCH_INDEX=https://download.pytorch.org/whl/cu121
     elif [ "${CUDA_MAJOR}" = "11" ]; then
         TORCH_INDEX=https://download.pytorch.org/whl/cu118
@@ -201,6 +258,10 @@ if [ "$TORCH" = true ]; then
     _pip_install "torch==${TORCH_VER}" "torchvision" --index-url "$TORCH_INDEX"
 fi
 
+if [ "$BASIC" = true ] || [ "$FLASH_ATTN" = true ] || [ "$XFORMERS" = true ] || [ "$TRAIN" = true ]; then
+    _require_blackwell_torch "$(_detect_torch_version)" || return 1
+fi
+
 # ─── Step 2: Install base requirements ──────────────────────────────────────
 
 if [ "$BASIC" = true ]; then
diff --git a/tests/test_setup_blackwell.sh b/tests/test_setup_blackwell.sh
new file mode 100755
index 0000000..8b67e3c
--- /dev/null
+++ b/tests/test_setup_blackwell.sh
@@ -0,0 +1,70 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
+TEST_ROOT=$(mktemp -d)
+trap 'rm -rf "$TEST_ROOT"' EXIT
+
+PYTHON_STUB="$TEST_ROOT/python"
+cat >"$PYTHON_STUB" <<'PYTHON'
+#!/bin/sh
+if [ "$1" = "-c" ]; then
+    case "$2" in
+        *"sys.version_info.major"*) printf '%s\n' "312" ;;
+        *"sys.version_info.minor"*) printf '%s\n' "12" ;;
+        *"torch.version.cuda"*) printf '%s\n' "12.8" ;;
+        *"torch.__version__"*) printf '%s\n' "${STUB_TORCH_VERSION:-2.7.0}" ;;
+    esac
+elif [ "$1" = "-m" ] && [ "$2" = "pip" ]; then
+    shift 2
+    printf 'PIP_INSTALL %s\n' "$*"
+    printf 'ARCHS %s %s\n' "${TORCH_CUDA_ARCH_LIST:-}" "${FLASH_ATTN_CUDA_ARCHS:-}"
+elif [ "$1" = "--version" ]; then
+    printf '%s\n' "Python 3.12.0"
+fi
+PYTHON
+chmod +x "$PYTHON_STUB"
+
+run_setup() {
+    local capability=$1
+    shift
+    # shellcheck disable=SC2016
+    env \
+        -u ANIGEN_TORCH_INDEX_URL \
+        -u CUDA_HOME \
+        -u FLASH_ATTN_CUDA_ARCHS \
+        -u TORCH_CUDA_ARCH_LIST \
+        -u TORCH_VERSION \
+        -u VIRTUAL_ENV \
+        ANIGEN_CUDA_CAPABILITY="$capability" \
+        ANIGEN_PYTHON="$PYTHON_STUB" \
+        PATH=/usr/bin:/bin \
+        bash -c 'cd "$1"; shift; . ./setup.sh "$@"' _ "$REPO_ROOT" "$@"
+}
+
+blackwell_output=$(run_setup 12.0 --torch)
+grep -Fq "Installing PyTorch 2.7.0 (CUDA index: https://download.pytorch.org/whl/cu128)" <<<"$blackwell_output"
+grep -Fq "ARCHS 12.0 120" <<<"$blackwell_output"
+
+legacy_output=$(run_setup 8.6 --torch)
+grep -Fq "Installing PyTorch 2.4.0 (CUDA index: https://download.pytorch.org/whl/cu121)" <<<"$legacy_output"
+
+# shellcheck disable=SC2016
+if old_torch_output=$(env \
+    -u ANIGEN_TORCH_INDEX_URL \
+    -u CUDA_HOME \
+    -u FLASH_ATTN_CUDA_ARCHS \
+    -u TORCH_CUDA_ARCH_LIST \
+    -u VIRTUAL_ENV \
+    ANIGEN_CUDA_CAPABILITY=12.0 \
+    ANIGEN_PYTHON="$PYTHON_STUB" \
+    PATH=/usr/bin:/bin \
+    TORCH_VERSION=2.6.0 \
+    bash -c 'cd "$1"; . ./setup.sh --torch' _ "$REPO_ROOT" 2>&1); then
+    echo "Expected PyTorch 2.6 to be rejected on Blackwell" >&2
+    exit 1
+fi
+grep -Fq "requires PyTorch >=2.7 with CUDA 12.8 support" <<<"$old_torch_output"
+
+echo "Blackwell setup selection tests passed"