fix: limit default GPU management

Avoid querying every visible CUDA device unless multi-GPU operation is explicitly requested.

Refs #15255

comfy/cli_args.py

@@ -74,7 +74,7 @@ parser.add_argument("--temp-directory", type=str, default=None, help="Set the Co
 parser.add_argument("--input-directory", type=str, default=None, help="Set the ComfyUI input directory. Overrides --base-directory.")
 parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.")
 parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.")
-parser.add_argument("--cuda-device", type=str, default=None, metavar="DEVICE_ID", help="Set the ids of cuda devices this instance will use, as a comma-separated list (e.g. '0' or '0,1'). All other devices will not be visible.")
+parser.add_argument("--cuda-device", type=str, default=None, metavar="DEVICE_ID", help="Set the ids of cuda devices this instance will use, as a comma-separated list (e.g. '0' or '0,1'). Passing multiple ids enables core multi-GPU support. All other devices will not be visible.")
 parser.add_argument("--default-device", type=int, default=None, metavar="DEFAULT_DEVICE_ID", help="Set the id of the default device, all other devices will stay visible.")
 cm_group = parser.add_mutually_exclusive_group()
 cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).")

comfy/model_management.py

@@ -243,6 +243,23 @@ def get_all_torch_devices(exclude_current=False):
             devices.remove(current)
     return devices
 
+def get_core_torch_devices(exclude_current=False):
+    """Return devices ComfyUI core is allowed to manage.
+
+    Other visible devices remain available to custom nodes. Core multi-GPU
+    management is enabled only by explicitly passing multiple devices through
+    --cuda-device.
+    """
+    if args.cuda_device is not None and "," in args.cuda_device:
+        devices = get_all_torch_devices()
+    else:
+        devices = [get_torch_device()]
+    if exclude_current:
+        current = get_torch_device()
+        if current in devices:
+            devices.remove(current)
+    return devices
+
 def get_gpu_device_options():
     """Return list of device option strings for node widgets.
 
@@ -250,7 +267,7 @@ def get_gpu_device_options():
     adds "gpu:0", "gpu:1", etc. (vendor-agnostic labels).
     """
     options = ["default", "cpu"]
-    devices = get_all_torch_devices()
+    devices = get_core_torch_devices()
     if len(devices) > 1:
         for i in range(len(devices)):
             options.append(f"gpu:{i}")
@@ -283,7 +300,7 @@ def resolve_gpu_device_option(option: str):
             idx = int(option[4:])
         except ValueError:
             return None
-        devices = get_all_torch_devices()
+        devices = get_core_torch_devices()
         if 0 <= idx < len(devices):
             return devices[idx]
     return None
@@ -603,7 +620,7 @@ try:
 except:
     logging.warning("Could not pick default device.")
 try:
-    for device in get_all_torch_devices(exclude_current=True):
+    for device in get_core_torch_devices(exclude_current=True):
         logging.info("Device: {}".format(get_torch_device_name(device)))
 except:
     pass
@@ -2041,7 +2058,7 @@ def soft_empty_cache(force=False):
         torch.cuda.ipc_collect()
 
 def unload_all_models():
-    for device in get_all_torch_devices():
+    for device in get_core_torch_devices():
         free_memory(1e30, device)
 
 def unload_model_and_clones(model: ModelPatcher, unload_additional_models=True, all_devices=False):
@@ -2067,7 +2084,7 @@ def unload_model_and_clones(model: ModelPatcher, unload_additional_models=True,
     if not all_devices:
         free_memory(1e30, get_torch_device(), keep_loaded)
     else:
-        for device in get_all_torch_devices():
+        for device in get_core_torch_devices():
             free_memory(1e30, device, keep_loaded)
 
 def debug_memory_summary():

comfy/multigpu.py

@@ -136,7 +136,7 @@ def create_multigpu_deepclones(model: ModelPatcher, max_gpus: int, gpu_options:
     # Exclude the primary model's actual device, not the global current device:
     # after SelectModelDevice(gpu:N) the primary may not live on the process's
     # current CUDA device, and excluding the wrong device picks bad extras.
-    all_devices = comfy.model_management.get_all_torch_devices(exclude_current=False)
+    all_devices = comfy.model_management.get_core_torch_devices(exclude_current=False)
     full_extra_devices = [d for d in all_devices if d != model.load_device]
     limit_extra_devices = full_extra_devices[:max_gpus-1]
     extra_devices = limit_extra_devices.copy()

main.py

@@ -253,10 +253,10 @@ if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_managemen
         logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows")
     else:
         try:
-            aimdo_initialized = comfy_aimdo.control.init_devices((d.index, int(args.vram_headroom * 1024 ** 3)) for d in comfy.model_management.get_all_torch_devices())
+            aimdo_initialized = comfy_aimdo.control.init_devices((d.index, int(args.vram_headroom * 1024 ** 3)) for d in comfy.model_management.get_core_torch_devices())
         except TypeError:
             # comfy-aimdo 0.4.9 protocol.
-            aimdo_initialized = comfy_aimdo.control.init_devices(d.index for d in comfy.model_management.get_all_torch_devices())
+            aimdo_initialized = comfy_aimdo.control.init_devices(d.index for d in comfy.model_management.get_core_torch_devices())
 
         if aimdo_initialized:
             if console_log_level == 'DEBUG':

server.py

@@ -694,9 +694,9 @@ class PromptServer():
             required_templates_version = FrontendManager.get_required_templates_version()
             comfy_package_versions = FrontendManager.get_comfy_package_versions()
 
-            # Report every torch device visible to multigpu, with the primary
+            # Report every device managed by ComfyUI core, with the primary
             # device first so existing clients that read devices[0] keep working.
-            torch_devices = comfy.model_management.get_all_torch_devices()
+            torch_devices = comfy.model_management.get_core_torch_devices()
             if primary_device in torch_devices:
                 torch_devices = [primary_device] + [d for d in torch_devices if d != primary_device]
             else:

tests/test_model_management_devices.py

@@ -0,0 +1,69 @@
+import pytest
+import torch
+
+from comfy.cli_args import args
+import comfy.model_management as model_management
+
+
+@pytest.mark.parametrize("cuda_device", [None, "0", "GPU-example"])
+def test_core_devices_default_to_current_without_enumerating(monkeypatch, cuda_device):
+    current = torch.device("cuda", 0)
+    monkeypatch.setattr(args, "cuda_device", cuda_device)
+    monkeypatch.setattr(args, "default_device", 1)
+    monkeypatch.setattr(model_management, "get_torch_device", lambda: current)
+    monkeypatch.setattr(
+        model_management,
+        "get_all_torch_devices",
+        lambda: pytest.fail("secondary devices must not be enumerated"),
+    )
+
+    assert model_management.get_core_torch_devices() == [current]
+
+
+@pytest.mark.parametrize("cuda_device", ["0,1", "GPU-primary,GPU-secondary"])
+def test_core_devices_use_explicit_multi_device_selection(monkeypatch, cuda_device):
+    current = torch.device("cuda", 0)
+    devices = [current, torch.device("cuda", 1)]
+    monkeypatch.setattr(args, "cuda_device", cuda_device)
+    monkeypatch.setattr(model_management, "get_torch_device", lambda: current)
+    monkeypatch.setattr(model_management, "get_all_torch_devices", lambda: devices.copy())
+
+    assert model_management.get_core_torch_devices() == devices
+    assert model_management.get_core_torch_devices(exclude_current=True) == devices[1:]
+
+
+def test_all_devices_remain_discoverable(monkeypatch):
+    monkeypatch.setattr(model_management, "cpu_state", model_management.CPUState.GPU)
+    monkeypatch.setattr(model_management, "is_nvidia", lambda: True)
+    monkeypatch.setattr(model_management.torch.cuda, "device_count", lambda: 2)
+
+    assert model_management.get_all_torch_devices() == [
+        torch.device("cuda", 0),
+        torch.device("cuda", 1),
+    ]
+
+
+def test_device_options_only_offer_core_devices(monkeypatch):
+    current = torch.device("cuda", 0)
+    secondary = torch.device("cuda", 1)
+    monkeypatch.setattr(args, "cuda_device", None)
+    monkeypatch.setattr(model_management, "get_torch_device", lambda: current)
+    monkeypatch.setattr(model_management, "get_all_torch_devices", lambda: [current, secondary])
+
+    assert model_management.get_gpu_device_options() == ["default", "cpu"]
+    assert model_management.resolve_gpu_device_option("gpu:1") is None
+
+    monkeypatch.setattr(args, "cuda_device", "0,1")
+    assert model_management.get_gpu_device_options() == ["default", "cpu", "gpu:0", "gpu:1"]
+    assert model_management.resolve_gpu_device_option("gpu:1") == secondary
+
+
+def test_unload_all_models_only_uses_core_devices(monkeypatch):
+    current = torch.device("cuda", 0)
+    freed = []
+    monkeypatch.setattr(model_management, "get_core_torch_devices", lambda: [current])
+    monkeypatch.setattr(model_management, "free_memory", lambda amount, device: freed.append((amount, device)))
+
+    model_management.unload_all_models()
+
+    assert freed == [(1e30, current)]