diff --git a/nodes.py b/nodes.py index 934ea6e..702c8dd 100644 --- a/nodes.py +++ b/nodes.py @@ -1,6 +1,30 @@ from typing import Any as any_type from comfy import model_management +# 尝试导入pynvml库,如果没有安装则提供相应提示 +try: + import pynvml + pynvml_installed = True + pynvml.nvmlInit() +except ImportError: + pynvml_installed = False + print("警告:未安装pynvml库,auto选项将不可用。") + +def get_gpu_memory_info(): + """获取GPU显存信息""" + if not pynvml_installed: + return None, None + + try: + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle) + total = memory_info.total / (1024 * 1024 * 1024) # 转换为GB + used = memory_info.used / (1024 * 1024 * 1024) # 转换为GB + return total, used + except Exception as e: + print(f"获取GPU信息出错: {e}") + return None, None + class AlwaysEqualProxy(str): def __eq__(self, _): return True @@ -8,6 +32,7 @@ class AlwaysEqualProxy(str): def __ne__(self, _): return False any_type = AlwaysEqualProxy("*") + class ReservedVRAMSetter: @classmethod def INPUT_TYPES(s): @@ -18,8 +43,12 @@ class ReservedVRAMSetter: "default": 1.0, "min": 0.6, "step": 0.1, - "display": "reserved (GB)" + "display": "reserved (GB)" }), + "mode": (["manual", "auto"], { + "default": "auto", + "display": "Mode" + }) }, "hidden": {"unique_id": "UNIQUE_ID", "extra_pnginfo": "EXTRA_PNGINFO"} } @@ -29,9 +58,27 @@ class ReservedVRAMSetter: OUTPUT_NODE = True FUNCTION = "set_vram" CATEGORY = "VRAM" - def set_vram(self, anything, reserved, unique_id=None, extra_pnginfo=None): - model_management.EXTRA_RESERVED_VRAM = int(reserved * 1024 * 1024 * 1024) - print(f'set EXTRA_RESERVED_VRAM={reserved}GB') + + def set_vram(self, anything, reserved, mode="auto", unique_id=None, extra_pnginfo=None): + if mode == "auto": + if pynvml_installed: + total, used = get_gpu_memory_info() + if total and used: + # 自动计算预留显存:使用当前已用显存+reserved设置值作为预留值 + auto_reserved = used + reserved + model_management.EXTRA_RESERVED_VRAM = int(auto_reserved * 1024 * 1024 * 1024) + print(f'set EXTRA_RESERVED_VRAM={auto_reserved:.2f}GB (自动模式: 总显存={total:.2f}GB, 已用={used:.2f}GB)') + else: + model_management.EXTRA_RESERVED_VRAM = int(reserved * 1024 * 1024 * 1024) + print(f'set EXTRA_RESERVED_VRAM={reserved}GB (自动模式失败,使用手动值)') + else: + model_management.EXTRA_RESERVED_VRAM = int(reserved * 1024 * 1024 * 1024) + print(f'set EXTRA_RESERVED_VRAM={reserved}GB (pynvml未安装,auto选项不可用)') + else: + # 手动模式 + model_management.EXTRA_RESERVED_VRAM = int(reserved * 1024 * 1024 * 1024) + print(f'set EXTRA_RESERVED_VRAM={reserved}GB (手动模式)') + return (anything,) NODE_CLASS_MAPPINGS = { @@ -39,4 +86,4 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { "ReservedVRAMSetter": "Set Reserved VRAM(GB) ⚙️" -} +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..4365be3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +pynvml \ No newline at end of file