Use torch CUDA 12.6 as default

Comfy-Org · Jan 30, 2025 · 2be01a8 · 2be01a8
1 parent 98bfef7
commit 2be01a8
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 6 deletions.
diff --git a/.github/workflows/run-on-gpu.yml b/.github/workflows/run-on-gpu.yml
@@ -61,7 +61,7 @@ jobs:
         env:
           PYTHONPATH: ${{ github.workspace }}
           TEST_E2E: true
-          TEST_E2E_COMFY_INSTALL_FLAGS: --nvidia --cuda-version 12.4
+          TEST_E2E_COMFY_INSTALL_FLAGS: --nvidia --cuda-version 12.6
           TEST_E2E_COMFY_LAUNCH_FLAGS_EXTRA: ""
         run: |
           pytest tests/e2e

diff --git a/.github/workflows/test-windows.yml b/.github/workflows/test-windows.yml
@@ -29,7 +29,7 @@ jobs:
           python -m pip install --upgrade pip
           pip install pytest
           pip install -e .
-          comfy --skip-prompt --workspace ./ComfyUI install --fast-deps --nvidia --cuda-version 12.4 --skip-manager
+          comfy --skip-prompt --workspace ./ComfyUI install --fast-deps --nvidia --cuda-version 12.6 --skip-manager
           comfy --workspace ./ComfyUI standalone  --platform windows --proc x86_64
           ls
           comfy standalone --rehydrate --platform windows --proc x86_64

diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py
@@ -198,7 +198,7 @@ def install(
             callback=g_gpu_exclusivity.validate,
         ),
     ] = None,
-    cuda_version: Annotated[CUDAVersion, typer.Option(show_default=True)] = CUDAVersion.v12_4,
+    cuda_version: Annotated[CUDAVersion, typer.Option(show_default=True)] = CUDAVersion.v12_6,
     amd: Annotated[
         Optional[bool],
         typer.Option(

diff --git a/comfy_cli/command/install.py b/comfy_cli/command/install.py
@@ -68,7 +68,13 @@ def pip_install_comfyui_dependencies(
                 "torchvision",
                 "torchaudio",
             ]
-            if plat == constants.OS.WINDOWS and cuda_version == constants.CUDAVersion.v12_4:
+
+            if plat == constants.OS.WINDOWS and cuda_version == constants.CUDAVersion.v12_6:
+                base_command += [
+                    "--extra-index-url",
+                    "https://download.pytorch.org/whl/cu126",
+                ]
+            elif plat == constants.OS.WINDOWS and cuda_version == constants.CUDAVersion.v12_4:
                 base_command += [
                     "--extra-index-url",
                     "https://download.pytorch.org/whl/cu124",
@@ -164,7 +170,7 @@ def execute(
     commit: Optional[str] = None,
     manager_commit: Optional[str] = None,
     gpu: constants.GPU_OPTION = None,
-    cuda_version: constants.CUDAVersion = constants.CUDAVersion.v12_4,
+    cuda_version: constants.CUDAVersion = constants.CUDAVersion.v12_6,
     plat: constants.OS = None,
     skip_torch_or_directml: bool = False,
     skip_requirement: bool = False,

diff --git a/comfy_cli/constants.py b/comfy_cli/constants.py
@@ -60,6 +60,7 @@ class PROC(str, Enum):
 
 
 class CUDAVersion(str, Enum):
+    v12_6 = "12.6"
     v12_4 = "12.4"
     v12_1 = "12.1"
     v11_8 = "11.8"

diff --git a/comfy_cli/uv.py b/comfy_cli/uv.py
@@ -65,7 +65,7 @@ def parse_req_file(rf: PathLike, skips: Optional[list[str]] = None):
 
 class DependencyCompiler:
     rocmPytorchUrl = "https://download.pytorch.org/whl/rocm6.1"
-    nvidiaPytorchUrl = "https://download.pytorch.org/whl/cu121"
+    nvidiaPytorchUrl = "https://download.pytorch.org/whl/cu126"
 
     overrideGpu = dedent(
         """