zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-24 18:10:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Limit Windows multi-GPU visibility (#15737)

Browse Source
This commit is contained in:
rattus
2026-08-20 09:31:47 +10:00
committed by GitHub
parent c67885b145
commit 5ab2f7a2d6
3 changed files with 48 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
main.py
View File

@@ -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