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

Recognize CPU state, low VRAM & high VRAM when no flags are set

Browse Source 
When flags like args.cpu, args.highvram or args.gpu_only aren't set, model_management should detect the correct state automatically
This commit is contained in:
John Slegers 2024-08-12 03:01:55 +02:00
parent e9589d6d92
commit 8bdaffe3b8
1 changed files with 25 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
comfy/model_management.py
View File

@ -81,6 +81,10 @@ except:
if args.cpu: if args.cpu:
cpu_state = CPUState.CPU cpu_state = CPUState.CPU
def is_nvidia():
global cpu_state
return cpu_state == CPUState.GPU and torch.cuda.is_available()
def is_intel_xpu(): def is_intel_xpu():
global cpu_state global cpu_state
global xpu_available global xpu_available
@ -89,6 +93,10 @@ def is_intel_xpu():
return True return True
return False return False
if args.cpu or (not is_nvidia() and not is_intel_xpu()):
cpu_state = CPUState.CPU
def get_torch_device(): def get_torch_device():
global directml_enabled global directml_enabled
global cpu_state global cpu_state
@ -107,11 +115,12 @@ def get_torch_device():
def get_total_memory(dev=None, torch_total_too=False): def get_total_memory(dev=None, torch_total_too=False):
global directml_enabled global directml_enabled
global total_ram
if dev is None: if dev is None:
dev = get_torch_device() dev = get_torch_device()
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
mem_total = psutil.virtual_memory().total mem_total = total_ram
mem_total_torch = mem_total mem_total_torch = mem_total
else: else:
if directml_enabled: if directml_enabled:
@ -134,9 +143,21 @@ def get_total_memory(dev=None, torch_total_too=False):
else: else:
return mem_total return mem_total
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024) total_ram = psutil.virtual_memory().total
total_ram = psutil.virtual_memory().total / (1024 * 1024) total_vram = get_total_memory(get_torch_device())
logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram)) logging.info(
"Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(
total_vram / (1024 * 1024),
total_ram / (1024 * 1024)
)
)
if cpu_state != CPUState.CPU:
if total_ram < 6 * (1024 * 1024 * 1024) :
set_vram_to = VRAMState.LOW_VRAM
lowvram_available = True
if total_ram > 12 * (1024 * 1024 * 1024) :
vram_state = VRAMState.HIGH_VRAM
try: try:
logging.info("pytorch version: {}".format(torch.version.__version__)) logging.info("pytorch version: {}".format(torch.version.__version__))
@ -173,13 +194,6 @@ else:
except: except:
XFORMERS_IS_AVAILABLE = False XFORMERS_IS_AVAILABLE = False
def is_nvidia():
global cpu_state
if cpu_state == CPUState.GPU:
if torch.version.cuda:
return True
return False
ENABLE_PYTORCH_ATTENTION = False ENABLE_PYTORCH_ATTENTION = False
if args.use_pytorch_cross_attention: if args.use_pytorch_cross_attention:
ENABLE_PYTORCH_ATTENTION = True ENABLE_PYTORCH_ATTENTION = True