From 5f65c5f8eda183b28b8e7940a7027084027290bd Mon Sep 17 00:00:00 2001 From: admin Date: Sat, 10 Feb 2024 03:29:17 +0800 Subject: [PATCH] [feat] upgrade --- comfy/model_management.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index a8dc91b9e..282cfb911 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -4,6 +4,7 @@ from comfy.cli_args import args import comfy.utils import torch import sys +import os.path class VRAMState(Enum): DISABLED = 0 #No vram present: no need to move models to vram @@ -92,8 +93,13 @@ def get_total_memory(dev=None, torch_total_too=False): dev = get_torch_device() if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): - mem_total = psutil.virtual_memory().total - mem_total_torch = mem_total + if os.path.isfile('/sys/fs/cgroup/memory/memory.limit_in_bytes'): + with open('/sys/fs/cgroup/memory/memory.limit_in_bytes', 'r') as f: + mem_total = int(f.read()) + mem_total_torch = mem_total + else: + mem_total = psutil.virtual_memory().total + mem_total_torch = mem_total else: if directml_enabled: mem_total = 1024 * 1024 * 1024 #TODO @@ -650,8 +656,15 @@ def get_free_memory(dev=None, torch_free_too=False): dev = get_torch_device() if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): - mem_free_total = psutil.virtual_memory().available - mem_free_torch = mem_free_total + if os.path.isfile('/sys/fs/cgroup/memory/memory.limit_in_bytes'): + with open('/sys/fs/cgroup/memory/memory.limit_in_bytes', 'r') as f: + mem_used_total = psutil.virtual_memory().used + mem_total = int(f.read()) + 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: if directml_enabled: mem_free_total = 1024 * 1024 * 1024 #TODO