wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-12 01:07:30 +08:00
Code Issues Actions 18 Packages Projects Releases Wiki Activity

Always enable cuda malloc on cu130 and higher. (#14381)

Browse Source
This commit is contained in:
comfyanonymous 2026-06-09 18:39:24 -07:00 committed by GitHub
parent 5ece24e73c
commit 5fcf7a4a0f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
cuda_malloc.py
View File

@ -2,6 +2,7 @@ import os
import importlib.util
from comfy.cli_args import args, PerformanceFeature
import subprocess
import re
#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():
@ -77,11 +78,24 @@ try:
except:
pass
def get_raw_cuda_version(version_str):
match = re.search(r'\+cu(\d+)', version_str)
if match:
try:
return int(match.group(1))
except:
pass
return None
if not args.cuda_malloc:
try:
if int(version[0]) >= 2 and "+cu" in version: # enable by default for torch version 2.0 and up only on cuda torch
if PerformanceFeature.AutoTune not in args.fast: # Autotune has issues with cuda malloc
args.cuda_malloc = cuda_malloc_supported()
cuda_version = get_raw_cuda_version(version)
if cuda_version is not None and cuda_version >= 130:
args.cuda_malloc = True
else:
args.cuda_malloc = cuda_malloc_supported()
except:
pass