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

Merge 250ae5db22 into 8fd07170f1

Browse Source
This commit is contained in:
j2gg0s 2025-12-29 11:55:21 +08:00 committed by GitHub
commit 7f03c2b750
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 31 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
comfy/model_management.py
View File

@ -188,13 +188,27 @@ def get_torch_device():
else: else:
return torch.device(torch.cuda.current_device()) return torch.device(torch.cuda.current_device())
cgroup_path_memory_max = None
if os.path.isfile('/sys/fs/cgroup/memory.max'):
cgroup_path_memory_max = '/sys/fs/cgroup/memory.max'
elif os.path.isfile('/sys/fs/cgroup/memory/memory.limit_in_bytes'):
cgroup_path_memory_max = '/sys/fs/cgroup/memory/memory.limit_in_bytes'
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 = -1
if cgroup_path_memory_max is not None:
with open(cgroup_path_memory_max, 'r') as f:
raw = f.read()
if raw.isdigit():
# NOTE: maybe max or empty
mem_total = int(raw)
if mem_total < 0:
mem_total = psutil.virtual_memory().total
mem_total_torch = mem_total mem_total_torch = mem_total
else: else:
if directml_enabled: if directml_enabled:
@ -1259,13 +1273,28 @@ def force_upcast_attention_dtype():
else: else:
return None return None
cgroup_path_memory_used = None
if os.path.isfile('/sys/fs/cgroup/memory.current'):
cgroup_path_memory_used = '/sys/fs/cgroup/memory.current'
elif os.path.isfile('/sys/fs/cgroup/memory/memory.usage_in_bytes'):
cgroup_path_memory_used = '/sys/fs/cgroup/memory/memory.usage_in_bytes'
def get_free_memory(dev=None, torch_free_too=False): def get_free_memory(dev=None, torch_free_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_free_total = psutil.virtual_memory().available mem_used = -1
if cgroup_path_memory_used is not None:
with open(cgroup_path_memory_used, 'r') as f:
raw = f.read()
if raw.isdigit():
mem_used = int(f.read())
if mem_used > 0:
mem_free_total = get_total_memory(dev) - mem_used
else:
mem_free_total = psutil.virtual_memory().available
mem_free_torch = mem_free_total mem_free_torch = mem_free_total
else: else:
if directml_enabled: if directml_enabled: