xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-16 01:36:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make cudaMalloc backed opt-in

Browse Source
This commit is contained in:
lspindler 2025-08-11 10:18:50 +02:00
parent 9126c0cfe4
commit 1c7fcdd089
2 changed files with 10 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/cli_args.py
View File

@ -53,7 +53,6 @@ parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID"
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.") 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 = 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("--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.")
fp_group = parser.add_mutually_exclusive_group() fp_group = parser.add_mutually_exclusive_group()

15
cuda_malloc.py
View File

@ -2,6 +2,7 @@ import os
import importlib.util import importlib.util
from comfy.cli_args import args from comfy.cli_args import args
import subprocess import subprocess
import logging
#Can't use pytorch to get the GPU names because the cuda malloc has to be set before the first import. #Can't use pytorch to get the GPU names because the cuda malloc has to be set before the first import.
def get_gpu_names(): def get_gpu_names():
@ -50,7 +51,7 @@ blacklist = {"GeForce GTX TITAN X", "GeForce GTX 980", "GeForce GTX 970", "GeFor
"GeForce GTX 1650", "GeForce GTX 1630", "Tesla M4", "Tesla M6", "Tesla M10", "Tesla M40", "Tesla M60" "GeForce GTX 1650", "GeForce GTX 1630", "Tesla M4", "Tesla M6", "Tesla M10", "Tesla M40", "Tesla M60"
} }
def cuda_malloc_supported(): def device_cuda_malloc_supported():
try: try:
names = get_gpu_names() names = get_gpu_names()
except: except:
@ -62,8 +63,8 @@ def cuda_malloc_supported():
return False return False
return True return True
def cuda_malloc_supported():
if not args.cuda_malloc: software_supported = False
try: try:
version = "" version = ""
torch_spec = importlib.util.find_spec("torch") torch_spec = importlib.util.find_spec("torch")
@ -76,16 +77,20 @@ if not args.cuda_malloc:
version = module.__version__ version = module.__version__
if int(version[0]) >= 2 and "+cu" in version: #enable by default for torch version 2.0 and up only on cuda torch if int(version[0]) >= 2 and "+cu" in version: #enable by default for torch version 2.0 and up only on cuda torch
args.cuda_malloc = cuda_malloc_supported() software_supported = True
except: except:
pass pass
return (software_supported and device_cuda_malloc_supported())
if args.cuda_malloc and not args.disable_cuda_malloc: if args.cuda_malloc:
env_var = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', None) env_var = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', None)
if env_var is None: if env_var is None:
env_var = "backend:cudaMallocAsync" env_var = "backend:cudaMallocAsync"
else: else:
env_var += ",backend:cudaMallocAsync" env_var += ",backend:cudaMallocAsync"
if not cuda_malloc_supported():
logging.warning("WARNING: this card most likely does not support cuda-malloc\n")
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = env_var os.environ['PYTORCH_CUDA_ALLOC_CONF'] = env_var