diff --git a/README.md b/README.md index f67ed6a..36f37ba 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,8 @@ new 2. Front-end input can be left unconnected. Added back-end output for random seed and reserved value. Back-end output can also be left unconnected. 3. Added a front-end VRAM cleanup toggle, allowing use as a VRAM cleanup node. Option to restore environment variables to default (0.6GB) manually before output. 4. Added maximum reserved value, effective in Auto mode, preventing excessive reservation in certain cases while slightly reducing Auto mode's capability. -image + +![example](image/example.jpg) 2025-10-10新增自动模式，自动模式会检测系统“已使用”的显存数量，再叠加用户设置值进行预留。避免多进程用户因为显存问题卡住运行。 @@ -30,14 +31,6 @@ new 接在排行较前的节点处即可，观察windows任务管理器共享显存溢出多少，就需要设置保留多少（可以略微多一点），填入该数值。运行工作流实时生效，输入单位是GB。 -![N57)EGC5978{(Y36IV~13AL](https://github.com/user-attachments/assets/245e5f11-c16d-403c-a438-567040f12ebf) - -![19~HL`3H{F %%LBE)3~3GPC](https://github.com/user-attachments/assets/fd8b61e4-e2e5-42ca-a516-2ddc1c7d0d8d) - -![_$)59`(5~XN5OH7NM %WHU](https://github.com/user-attachments/assets/bb652d70-805b-452e-a522-f271c8c70bf4) - -![image](https://github.com/user-attachments/assets/48f8ca7f-2a13-4ef5-a5bb-5f6ef9c974e3) - 2026-06-27 DynamicVRAM compatibility update - The node now updates ComfyUI DynamicVRAM / comfy-aimdo simple vram headroom when DynamicVRAM is enabled. diff --git a/image/example.jpg b/image/example.jpg new file mode 100644 index 0000000..b5970ad Binary files /dev/null and b/image/example.jpg differ diff --git a/nodes.py b/nodes.py index 72eb483..f410146 100644 --- a/nodes.py +++ b/nodes.py @@ -3,9 +3,15 @@ from comfy import model_management import random import time import gc -# 尝试导入pynvml库，如果没有安装则提供相应提示 +# 尝试导入 pynvml；在部分 AMD/Intel 环境里，导入阶段本身就可能因为缺少 NVML DLL 直接抛异常。 try: import pynvml +except Exception as e: + pynvml_installed = False + pynvml = None + print("[ReservedVRAM]警告：pynvml不可用，auto选项将不可用。") + print(f"[ReservedVRAM]pynvml导入失败: {e}") +else: try: pynvml.nvmlInit() pynvml_installed = True @@ -14,10 +20,6 @@ try: pynvml = None print("[ReservedVRAM]警告：pynvml可导入但NVML初始化失败，auto选项将不可用。") print(f"[ReservedVRAM]NVML初始化失败: {e}") -except ImportError: - pynvml_installed = False - pynvml = None - print("[ReservedVRAM]警告：未安装pynvml库，auto选项将不可用。") # 初始化随机状态 initial_random_state = random.getstate() @@ -25,11 +27,11 @@ random.seed(time.time()) reserved_vram_random_state = random.getstate() random.setstate(initial_random_state) -def get_gpu_memory_info(): - """获取GPU显存信息""" - if pynvml_installed and pynvml is not None: - try: - handle = pynvml.nvmlDeviceGetHandleByIndex(0) +def get_gpu_memory_info(): + """获取GPU显存信息""" + if pynvml_installed and pynvml is not None: + try: + handle = pynvml.nvmlDeviceGetHandleByIndex(0) memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle) total = memory_info.total / (1024 * 1024 * 1024) used = memory_info.used / (1024 * 1024 * 1024) @@ -51,44 +53,44 @@ def get_gpu_memory_info(): return total_gb, used_gb except Exception as e: print(f"[ReservedVRAM]获取GPU信息出错(torch): {e}") - - return None, None - -def set_reserved_vram(reserved_gb): - reserved_gb = max(0.0, float(reserved_gb)) - reserved_vram = int(reserved_gb * 1024 * 1024 * 1024) - if hasattr(model_management, "set_extra_reserved_vram"): - model_management.set_extra_reserved_vram(reserved_gb) - else: - model_management.EXTRA_RESERVED_VRAM = reserved_vram - sync_dynamic_vram_headroom(reserved_vram) - -def sync_dynamic_vram_headroom(reserved_vram): - try: - import comfy.memory_management as memory_management - - if not getattr(memory_management, "aimdo_enabled", False): - return - - import comfy_aimdo.control as aimdo_control - - if getattr(aimdo_control, "lib", None) is None: - return - - try: - aimdo_control.init(simple_vram_headroom=int(reserved_vram)) - except TypeError: - setter = getattr(aimdo_control.lib, "set_simple_vram_headroom", None) - if setter is not None: - setter(int(reserved_vram)) - except Exception as e: - print(f"[ReservedVRAM]同步DynamicVRAM预留显存失败: {e}") - -def new_random_seed(): - """生成一个新的随机种子""" - global reserved_vram_random_state - prev_random_state = random.getstate() - random.setstate(reserved_vram_random_state) + + return None, None + +def set_reserved_vram(reserved_gb): + reserved_gb = max(0.0, float(reserved_gb)) + reserved_vram = int(reserved_gb * 1024 * 1024 * 1024) + if hasattr(model_management, "set_extra_reserved_vram"): + model_management.set_extra_reserved_vram(reserved_gb) + else: + model_management.EXTRA_RESERVED_VRAM = reserved_vram + sync_dynamic_vram_headroom(reserved_vram) + +def sync_dynamic_vram_headroom(reserved_vram): + try: + import comfy.memory_management as memory_management + + if not getattr(memory_management, "aimdo_enabled", False): + return + + import comfy_aimdo.control as aimdo_control + + if getattr(aimdo_control, "lib", None) is None: + return + + try: + aimdo_control.init(simple_vram_headroom=int(reserved_vram)) + except TypeError: + setter = getattr(aimdo_control.lib, "set_simple_vram_headroom", None) + if setter is not None: + setter(int(reserved_vram)) + except Exception as e: + print(f"[ReservedVRAM]同步DynamicVRAM预留显存失败: {e}") + +def new_random_seed(): + """生成一个新的随机种子""" + global reserved_vram_random_state + prev_random_state = random.getstate() + random.setstate(reserved_vram_random_state) seed = random.randint(1, 1125899906842624) reserved_vram_random_state = random.getstate() random.setstate(prev_random_state) @@ -167,26 +169,28 @@ class ReservedVRAMSetter: if mode == "auto": total, used = get_gpu_memory_info() if total is not None and used is not None: + # 自动计算预留显存 auto_reserved = used + reserved auto_reserved = max(0, auto_reserved) + # 如果设置了最大预留值且大于0，则应用限制 if auto_max_reserved > 0: auto_reserved = min(auto_reserved, auto_max_reserved) print(f'[ReservedVRAM]set EXTRA_RESERVED_VRAM={auto_reserved:.2f}GB (自动模式: 总显存={total:.2f}GB, 已用={used:.2f}GB, 最大限制值{auto_max_reserved:.2f}GB)') - else: - print(f'[ReservedVRAM]set EXTRA_RESERVED_VRAM={auto_reserved:.2f}GB (自动模式: 总显存={total:.2f}GB, 已用={used:.2f}GB)') - set_reserved_vram(auto_reserved) - final_reserved_vram = round(auto_reserved, 2) - else: - manual_reserved = max(0, reserved) - set_reserved_vram(manual_reserved) - print(f'[ReservedVRAM]set EXTRA_RESERVED_VRAM={manual_reserved}GB (自动模式不可用，使用手动值)') - final_reserved_vram = round(manual_reserved, 2) - else: - # 手动模式 - reserved = max(0, reserved) - set_reserved_vram(reserved) - print(f'[ReservedVRAM]set EXTRA_RESERVED_VRAM={reserved}GB (手动模式)，忽略最大限制值') - final_reserved_vram = round(reserved, 2) + else: + print(f'[ReservedVRAM]set EXTRA_RESERVED_VRAM={auto_reserved:.2f}GB (自动模式: 总显存={total:.2f}GB, 已用={used:.2f}GB)') + set_reserved_vram(auto_reserved) + final_reserved_vram = round(auto_reserved, 2) + else: + manual_reserved = max(0, reserved) + set_reserved_vram(manual_reserved) + print(f'[ReservedVRAM]set EXTRA_RESERVED_VRAM={manual_reserved}GB (自动模式不可用，使用手动值)') + final_reserved_vram = round(manual_reserved, 2) + else: + # 手动模式 + reserved = max(0, reserved) + set_reserved_vram(reserved) + print(f'[ReservedVRAM]set EXTRA_RESERVED_VRAM={reserved}GB (手动模式)，忽略最大限制值') + final_reserved_vram = round(reserved, 2) from comfy_execution.graph import ExecutionBlocker output_value = anything if anything is not None else ExecutionBlocker(None)