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

[feat] upgrade

Browse Source
This commit is contained in:
admin 2024-02-10 03:29:17 +08:00
parent 25a4805e51
commit 5f65c5f8ed
1 changed files with 17 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
comfy/model_management.py
View File

@ -4,6 +4,7 @@ from comfy.cli_args import args
import comfy.utils import comfy.utils
import torch import torch
import sys import sys
import os.path
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
@ -92,6 +93,11 @@ def get_total_memory(dev=None, torch_total_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'):
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 = psutil.virtual_memory().total
mem_total_torch = mem_total mem_total_torch = mem_total
else: else:
@ -650,6 +656,13 @@ 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'):
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_total = psutil.virtual_memory().available
mem_free_torch = mem_free_total mem_free_torch = mem_free_total
else: else: