diff --git a/comfy/cli_args.py b/comfy/cli_args.py index de3e85c08..993173a53 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -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).") diff --git a/main.py b/main.py index c33f0e17b..0ab33d603 100644 --- a/main.py +++ b/main.py @@ -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