xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-14 10:30:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

feature: Set the Ascend NPU to use a single one

Browse Source
This commit is contained in:
muxuezzz 2025-09-01 16:09:46 +08:00
parent 9b15155972
commit 6fb85e16c8
2 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/cli_args.py
View File

@ -50,6 +50,7 @@ parser.add_argument("--input-directory", type=str, default=None, help="Set the C
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=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use. All other devices will not be visible.")
parser.add_argument("--ascend-device", type=int, default=None, metavar="ASCEND_DEVICE_ID", help="Set the id of the ascend device this instance will use. 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).")

3
main.py
View File

@ -128,6 +128,9 @@ if __name__ == "__main__":
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device)
logging.info("Set cuda device to: {}".format(args.cuda_device))
elif args.ascend_device is not None:
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.ascend_device)
logging.info("Set npu device to: {}".format(args.ascend_device))
if args.oneapi_device_selector is not None:
os.environ['ONEAPI_DEVICE_SELECTOR'] = args.oneapi_device_selector