diff --git a/comfy/model_management.py b/comfy/model_management.py index f2e23c446e59..ac50fbd15f99 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -22,6 +22,7 @@ from comfy.cli_args import args, PerformanceFeature import torch import sys +import os import importlib import platform import weakref @@ -81,6 +82,56 @@ def get_supported_float8_types(): FLOAT8_TYPES = get_supported_float8_types() +def get_docker_memory_limit(): + """ + Get Docker container memory limit from cgroup. + Returns None if not in a container or no limit is set. + """ + try: + # Try cgroup v2 first + cgroup_v2_path = '/sys/fs/cgroup/memory.max' + if os.path.exists(cgroup_v2_path): + with open(cgroup_v2_path, 'r') as f: + limit = f.read().strip() + if limit != 'max': # 'max' means no limit + return int(limit) + + # Try cgroup v1 + cgroup_v1_path = '/sys/fs/cgroup/memory/memory.limit_in_bytes' + if os.path.exists(cgroup_v1_path): + with open(cgroup_v1_path, 'r') as f: + limit = int(f.read().strip()) + # Check if it's the default "no limit" value (very large number) + # Typical values: 9223372036854771712 or similar + if limit < (1 << 60): # Less than ~1 exabyte, likely a real limit + return limit + except Exception as e: + logging.debug(f"Could not read cgroup memory limit: {e}") + + return None + +def get_docker_memory_usage(): + """ + Get current memory usage from cgroup. + Returns None if not available. + """ + try: + # Try cgroup v2 + cgroup_v2_path = '/sys/fs/cgroup/memory.current' + if os.path.exists(cgroup_v2_path): + with open(cgroup_v2_path, 'r') as f: + return int(f.read().strip()) + + # Try cgroup v1 + cgroup_v1_path = '/sys/fs/cgroup/memory/memory.usage_in_bytes' + if os.path.exists(cgroup_v1_path): + with open(cgroup_v1_path, 'r') as f: + return int(f.read().strip()) + except Exception as e: + logging.debug(f"Could not read cgroup memory usage: {e}") + + return None + xpu_available = False torch_version = "" try: @@ -200,7 +251,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 + # Check if running in Docker with memory limit + docker_limit = get_docker_memory_limit() + if docker_limit is not None: + mem_total = docker_limit + logging.debug(f"Docker total memory limit: {docker_limit/(1024**3):.2f}GB") + else: + mem_total = psutil.virtual_memory().total mem_total_torch = mem_total else: if directml_enabled: @@ -1214,7 +1271,18 @@ 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 + # Check if running in Docker with memory limit + docker_limit = get_docker_memory_limit() + docker_usage = get_docker_memory_usage() + + if docker_limit is not None and docker_usage is not None: + # Running in Docker with memory limit + mem_free_total = docker_limit - docker_usage + logging.debug(f"Docker memory: limit={docker_limit/(1024**3):.2f}GB, usage={docker_usage/(1024**3):.2f}GB, free={mem_free_total/(1024**3):.2f}GB") + else: + # Not in Docker or no limit set, use system memory + mem_free_total = psutil.virtual_memory().available + mem_free_torch = mem_free_total else: if directml_enabled: