add args: npu-device

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

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:

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