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

Merge 0b5953d15d35a369fe78fe48a76ea76b7125f3cc into a00e1489d2249ebf619a698fcac7c21c7669a50a

Browse Source
This commit is contained in:
Xuexue 2025-01-19 17:25:19 +00:00 committed by GitHub
commit 4e5e32d4c7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 31 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
comfy/model_management.py
View File

@ -22,10 +22,12 @@ from enum import Enum
from comfy.cli_args import args from comfy.cli_args import args
import torch import torch
import sys import sys
import os.path
import platform import platform
import weakref import weakref
import gc import gc
class VRAMState(Enum): class VRAMState(Enum):
DISABLED = 0 #No vram present: no need to move models to vram DISABLED = 0 #No vram present: no need to move models to vram
NO_VRAM = 1 #Very low vram: enable all the options to save vram NO_VRAM = 1 #Very low vram: enable all the options to save vram
@ -128,14 +130,32 @@ def get_torch_device():
else: else:
return torch.device(torch.cuda.current_device()) return torch.device(torch.cuda.current_device())
def get_containerd_memory_limit():
cgroup_memory_limit = '/sys/fs/cgroup/memory/memory.limit_in_bytes'
if os.path.isfile(cgroup_memory_limit):
with open(cgroup_memory_limit, 'r') as f:
return int(f.read())
return 0
def get_containerd_memory_used():
cgroup_memory_used = '/sys/fs/cgroup/memory/memory.usage_in_bytes'
if os.path.isfile(cgroup_memory_used):
with open(cgroup_memory_used, 'r') as f:
return int(f.read())
return 0
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
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 = get_containerd_memory_limit()
mem_total_torch = mem_total if mem_total > 0:
mem_total_torch = mem_total
else:
mem_total = psutil.virtual_memory().total
mem_total_torch = mem_total
else: else:
if directml_enabled: if directml_enabled:
mem_total = 1024 * 1024 * 1024 #TODO mem_total = 1024 * 1024 * 1024 #TODO
@ -937,8 +957,14 @@ def get_free_memory(dev=None, torch_free_too=False):
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_free_total = psutil.virtual_memory().available mem_total = get_containerd_memory_limit()
mem_free_torch = mem_free_total if mem_total > 0:
mem_used_total = get_containerd_memory_used()
mem_free_total = mem_total - mem_used_total
mem_free_torch = mem_free_total
else:
mem_free_total = psutil.virtual_memory().available
mem_free_torch = mem_free_total
else: else:
if directml_enabled: if directml_enabled:
mem_free_total = 1024 * 1024 * 1024 #TODO mem_free_total = 1024 * 1024 * 1024 #TODO
@ -993,6 +1019,7 @@ def is_device_mps(device):
def is_device_cuda(device): def is_device_cuda(device):
return is_device_type(device, 'cuda') return is_device_type(device, 'cuda')
def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False): def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
global directml_enabled global directml_enabled