diff --git a/comfy/cli_args.py b/comfy/cli_args.py index a864205be..cc1ea1779 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.") +parser.add_argument("--npu-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the Ascend npu device this instance will use.") 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).") cm_group.add_argument("--disable-cuda-malloc", action="store_true", help="Disable cudaMallocAsync.") diff --git a/comfy/model_management.py b/comfy/model_management.py index b6f4e2d19..2a68cf20b 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -125,9 +125,20 @@ def is_mlu(): return True return False +def set_npu_device(device_id): + """Set the NPU device to use. + + Args: + device_id (int): The id of the NPU device to use. + """ + if device_id is not None: + torch.npu.set_device(device_id) + def get_torch_device(): global directml_enabled global cpu_state + if is_ascend_npu(): + return torch.device("npu", torch.npu.current_device()) if directml_enabled: global directml_device return directml_device @@ -138,8 +149,6 @@ def get_torch_device(): else: if is_intel_xpu(): return torch.device("xpu", torch.xpu.current_device()) - elif is_ascend_npu(): - return torch.device("npu", torch.npu.current_device()) elif is_mlu(): return torch.device("mlu", torch.mlu.current_device()) else: diff --git a/main.py b/main.py index 1b100fa8a..b490a9a99 100644 --- a/main.py +++ b/main.py @@ -6,6 +6,7 @@ import importlib.util import folder_paths import time from comfy.cli_args import args +import comfy.model_management from app.logger import setup_logger import itertools import utils.extra_config @@ -114,6 +115,10 @@ if __name__ == "__main__": os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device) logging.info("Set cuda device to: {}".format(args.cuda_device)) + if args.npu_device is not None: + comfy.model_management.set_npu_device(args.npu_device) + logging.info("Set npu device to: {}".format(args.npu_device)) + if args.oneapi_device_selector is not None: os.environ['ONEAPI_DEVICE_SELECTOR'] = args.oneapi_device_selector logging.info("Set oneapi device selector to: {}".format(args.oneapi_device_selector)) @@ -137,7 +142,6 @@ import execution import server from server import BinaryEventTypes import nodes -import comfy.model_management import comfyui_version import app.logger