298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-11 02:15:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[feat] get correct memory used in container

Browse Source
This commit is contained in:
xuexuexue 2024-07-10 00:59:00 +08:00
parent 0215f8013f
commit 55f9f76a56
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
comfy/model_management.py
View File

@ -94,6 +94,13 @@ def get_containerd_memory_limit():
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):
global directml_enabled
if dev is None:
@ -664,7 +671,7 @@ def get_free_memory(dev=None, torch_free_too=False):
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
mem_total = get_containerd_memory_limit()
if mem_total > 0:
mem_used_total = psutil.virtual_memory().used
mem_used_total = get_containerd_memory_used()
mem_free_total = mem_total - mem_used_total
mem_free_torch = mem_free_total
else: