Add tests using gpu runners.

.github/workflows/build-and-test.yml

@@ -15,6 +15,7 @@ jobs:
       PYTHONIOENCODING: "utf8"
 
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
         python-version: ["3.9"]
@@ -32,7 +33,9 @@ jobs:
         run: |
           pip install -e .
 
-      - name: Run Tests
+      - name: Run Core Functionality
         run: |
           comfy --skip-prompt --no-enable-telemetry env
           comfy --skip-prompt install --cpu
+          comfy launch --background -- --cpu
+

.github/workflows/run-on-gpu.yml

@@ -0,0 +1,47 @@
+name: "Test CLI Tool on GPU runners"
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - 'comfy_cli/**'
+  pull_request:
+    branches:
+      - main
+    paths:
+      - 'comfy_cli/**'
+
+jobs:
+  test-cli-gpu:
+    name: "Run Tests on GPU Runners"
+    runs-on: 
+      group: gpu-runners
+      labels: ${{ matrix.os }}-x64-gpu # 
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [windows, linux]
+
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      - name: Check Nvidia
+        run: |
+          nvidia-smi
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.12
+
+      - name: Install Dependencies
+        run: |
+          pip install -e .
+
+      - name: Run Core Functionality
+        run: |
+          comfy --skip-prompt --no-enable-telemetry env
+          comfy --skip-prompt install ${{ matrix.comfy-flags }} --nvidia --cuda-version 11.8
+          comfy launch --background

comfy_cli/cmdline.py

@@ -21,7 +21,7 @@
 from comfy_cli.command import install as install_inner
 from comfy_cli.command.models import models as models_command
 from comfy_cli.config_manager import ConfigManager
-from comfy_cli.constants import GPU_OPTION
+from comfy_cli.constants import GPU_OPTION, CUDAVersion
 from comfy_cli.env_checker import EnvChecker, check_comfy_server_running
 from comfy_cli.update import check_for_updates
 from comfy_cli.workspace_manager import (
@@ -179,6 +179,9 @@ def install(
             callback=gpu_exclusivity_callback,
         ),
     ] = None,
+    cuda_version: Annotated[
+        CUDAVersion, typer.Option(show_default=True)
+    ] = CUDAVersion.v12_1,
     amd: Annotated[
         bool,
         typer.Option(
@@ -312,6 +315,7 @@ def install(
         skip_manager,
         commit=commit,
         gpu=gpu,
+        cuda_version=cuda_version,
         plat=platform,
         skip_torch_or_directml=skip_torch_or_directml,
         skip_requirement=skip_requirement,

comfy_cli/command/install.py

@@ -24,6 +24,7 @@ def install_comfyui_dependencies(
     repo_dir,
     gpu: GPU_OPTION,
     plat: constants.OS,
+    cuda_version: constants.CUDAVersion,
     skip_torch_or_directml: bool,
     skip_requirement: bool,
 ):
@@ -50,18 +51,33 @@ def install_comfyui_dependencies(
 
         # install torch for NVIDIA
         if gpu == GPU_OPTION.NVIDIA:
-            pip_url = ["--extra-index-url", "https://download.pytorch.org/whl/cu121"]
-            result = subprocess.run(
-                [
-                    sys.executable,
-                    "-m",
-                    "pip",
-                    "install",
-                    "torch",
-                    "torchvision",
-                    "torchaudio",
+            base_command = [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "torch",
+                "torchvision",
+                "torchaudio",
+            ]
+            if (
+                plat == constants.OS.WINDOWS
+                and cuda_version == constants.CUDAVersion.v12_1
+            ):
+                base_command += [
+                    "--extra-index-url",
+                    "https://download.pytorch.org/whl/cu121",
                 ]
-                + pip_url,
+            elif (
+                plat == constants.OS.WINDOWS
+                and cuda_version == constants.CUDAVersion.v11_8
+            ):
+                base_command += [
+                    "--extra-index-url",
+                    "https://download.pytorch.org/whl/cu118",
+                ]
+            result = subprocess.run(
+                base_command,
                 check=False,
             )
         # Beta support for intel arch based on this PR: https://github.com/comfyanonymous/ComfyUI/pull/3439
@@ -149,6 +165,7 @@ def execute(
     skip_manager: bool,
     commit=None,
     gpu: constants.GPU_OPTION = None,
+    cuda_version: constants.CUDAVersion = constants.CUDAVersion.v12_1,
     plat: constants.OS = None,
     skip_torch_or_directml: bool = False,
     skip_requirement: bool = False,
@@ -185,7 +202,7 @@ def execute(
         subprocess.run(["git", "checkout", commit])
 
     install_comfyui_dependencies(
-        repo_dir, gpu, plat, skip_torch_or_directml, skip_requirement
+        repo_dir, gpu, plat, cuda_version, skip_torch_or_directml, skip_requirement
     )
 
     WorkspaceManager().set_recent_workspace(repo_dir)

comfy_cli/constants.py

@@ -55,6 +55,11 @@ class OS(Enum):
 ]
 
 
+class CUDAVersion(str, Enum):
+    v12_1 = "12.1"
+    v11_8 = "11.8"
+
+
 class GPU_OPTION(Enum):
     NVIDIA = "Nvidia"
     AMD = "Amd"
@@ -67,4 +72,4 @@ class GPU_OPTION(Enum):
 # https://github.com/comfyanonymous/ComfyUI/blob/a88b0ebc2d2f933c94e42aa689c42e836eedaf3c/folder_paths.py#L5
 SUPPORTED_PT_EXTENSIONS = (".ckpt", ".pt", ".bin", ".pth", ".safetensors")
 
-NODE_ZIP_FILENAME = "node.zip"
+NODE_ZIP_FILENAME = "node.zip"

pyrightconfig.json

@@ -0,0 +1,3 @@
+{
+    "pythonPlatform": "All", 
+}