Limit Windows multi-GPU visibility (#15737)

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'), 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).")

cuda_malloc.py

@@ -28,19 +28,19 @@ def get_gpu_names():
             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",

main.py

@@ -17,7 +17,7 @@ import importlib.metadata
 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 @@ if __name__ == "__main__":
     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 @@ if os.name == "nt":
 
 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 @@ if __name__ == "__main__":
         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 @@ if __name__ == "__main__":
         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