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

Enable the --cuda-device parameter to support both CUDA and Ascend NPUs simultaneously.

Browse Source
This commit is contained in:
muxuezzz 2025-09-17 20:26:50 +08:00
parent 6fb85e16c8
commit c791d33ad4
2 changed files with 14 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/cli_args.py
View File

@ -50,7 +50,6 @@ 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).")

24
main.py
View File

@ -103,18 +103,23 @@ def execute_prestartup_script():
apply_custom_paths()
execute_prestartup_script()
if 'torch' in sys.modules:
logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
# Main code
import asyncio
import shutil
import threading
import gc
import comfy.model_management
if os.name == "nt":
os.environ['MIMALLOC_PURGE_DELAY'] = '0'
if __name__ == "__main__":
device = comfy.model_management.get_torch_device()
device_name = comfy.model_management.get_torch_device_name(device)
if args.default_device is not None:
default_dev = args.default_device
devices = list(range(32))
@ -125,12 +130,13 @@ if __name__ == "__main__":
os.environ['HIP_VISIBLE_DEVICES'] = str(devices)
if args.cuda_device is not None:
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 device_name != 'npu':
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))
else:
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device)
logging.info("Set npu device to: {}".format(args.cuda_device))
if args.oneapi_device_selector is not None:
os.environ['ONEAPI_DEVICE_SELECTOR'] = args.oneapi_device_selector
@ -142,8 +148,7 @@ if __name__ == "__main__":
import cuda_malloc
if 'torch' in sys.modules:
logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
import comfy.utils
@ -151,7 +156,6 @@ import execution
import server
from protocol import BinaryEventTypes
import nodes
import comfy.model_management
import comfyui_version
import app.logger
import hook_breaker_ac10a0