From 5ab2f7a2d676c1fb7b410c22e82e2ed8f217b56c Mon Sep 17 00:00:00 2001 From: rattus <46076784+rattus128@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:31:47 +1000 Subject: [PATCH] Limit Windows multi-GPU visibility (#15737) --- comfy/cli_args.py | 2 +- cuda_malloc.py | 8 ++++---- main.py | 47 +++++++++++++++++++++++++++++++++++++++++++---- 3 files changed, 48 insertions(+), 9 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index c6660846d06..659e772edd2 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -74,7 +74,7 @@ def get_file_log_outputs(outputs): 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'), or 'all' to leave all currently visible devices available. 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).") diff --git a/cuda_malloc.py b/cuda_malloc.py index 8c4422db82e..022c9ae881b 100644 --- a/cuda_malloc.py +++ b/cuda_malloc.py @@ -28,19 +28,19 @@ def enum_display_devices(): device_info = DISPLAY_DEVICEA() device_info.cb = ctypes.sizeof(device_info) device_index = 0 - gpu_names = set() + gpu_names = [] while user32.EnumDisplayDevicesA(None, device_index, ctypes.byref(device_info), 0): device_index += 1 - gpu_names.add(device_info.DeviceString.decode('utf-8')) + gpu_names.append(device_info.DeviceString.decode('utf-8')) return gpu_names return enum_display_devices() else: - gpu_names = set() + gpu_names = [] out = subprocess.check_output(['nvidia-smi', '-L']) for l in out.split(b'\n'): if len(l) > 0: - gpu_names.add(l.decode('utf-8').split(' (UUID')[0]) + gpu_names.append(l.decode('utf-8').split(' (UUID')[0]) return gpu_names blacklist = {"GeForce GTX TITAN X", "GeForce GTX 980", "GeForce GTX 970", "GeForce GTX 960", "GeForce GTX 950", "GeForce 945M", diff --git a/main.py b/main.py index b6f7d6cef73..9fadc25b92b 100644 --- a/main.py +++ b/main.py @@ -17,7 +17,7 @@ import folder_paths import time from comfy.cli_args import enables_dynamic_vram -from app.logger import setup_logger +from app.logger import setup_logger, log_startup_warning console_log_level = get_console_log_level(args.verbose) file_log_outputs = get_file_log_outputs(args.verbose) setup_logger(log_level=console_log_level, file_outputs=file_log_outputs, use_stdout=args.log_stdout) @@ -41,6 +41,44 @@ os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1' os.environ['DO_NOT_TRACK'] = '1' + import cuda_malloc + + if os.name == "nt": + cuda_visibility = os.environ.get("CUDA_VISIBLE_DEVICES") + device_selection = args.cuda_device + + try: + gpu_count = sum("NVIDIA" in name.upper() for name in cuda_malloc.get_gpu_names()) + except OSError: + gpu_count = 0 + + if gpu_count > 1: + warning = None + multiple_visible = False + if device_selection is None and args.default_device is None and cuda_visibility is None: + os.environ["CUDA_VISIBLE_DEVICES"] = "0" + warning = "Multiple NVIDIA GPUs detected. ComfyUI will use GPU 0 only on Windows by default. To restore all GPUs, pass --cuda-device all --disable-pinned-memory." + elif device_selection == "all": + multiple_visible = cuda_visibility is None or "," in cuda_visibility + elif device_selection is not None: + multiple_visible = "," in device_selection + elif args.default_device is not None: + multiple_visible = True + else: + multiple_visible = "," in cuda_visibility + + if multiple_visible and not args.disable_pinned_memory: + warning = "Multiple NVIDIA GPUs are visible on Windows with pinned memory enabled. Restart with --disable-pinned-memory to avoid CUDA host-transfer failures." + + if warning: + log_startup_warning(f""" +________________________________________________________________________ +WARNING WARNING WARNING WARNING WARNING + +{warning} +________________________________________________________________________ +""".strip()) + faulthandler.enable(file=sys.stderr, all_threads=args.debug_hang) if __name__ == "__main__" and args.debug_hang: dumping_traceback = False @@ -74,7 +112,7 @@ def dump_traceback_on_sigint(signum, frame): if __name__ == "__main__": os.environ['TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL'] = '1' - if args.default_device is not None: + if args.default_device is not None and args.cuda_device != "all": default_dev = args.default_device devices = list(range(32)) devices.remove(default_dev) @@ -83,7 +121,9 @@ def dump_traceback_on_sigint(signum, frame): os.environ['CUDA_VISIBLE_DEVICES'] = str(devices) os.environ['HIP_VISIBLE_DEVICES'] = str(devices) - if args.cuda_device is not None: + if args.cuda_device == "all": + logging.info("Set cuda devices to all") + elif args.cuda_device is not None: os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device) os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device) os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device) @@ -97,7 +137,6 @@ def dump_traceback_on_sigint(signum, frame): if 'CUBLAS_WORKSPACE_CONFIG' not in os.environ: os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8" - import cuda_malloc if "rocm" in cuda_malloc.get_torch_version_noimport(): os.environ['OCL_SET_SVM_SIZE'] = '262144' # set at the request of AMD