From b9448f58fb77c966645ba6bfef714142075f004c Mon Sep 17 00:00:00 2001 From: windecay <4453636@qq.com> Date: Sat, 27 Jun 2026 00:32:18 +0800 Subject: [PATCH] =?UTF-8?q?feat(reserved=5Fvram):=20=E6=B7=BB=E5=8A=A0Dyna?= =?UTF-8?q?micVRAM=E5=85=BC=E5=AE=B9=E6=80=A7=E6=94=AF=E6=8C=81=EF=BC=8C?= =?UTF-8?q?=E7=BB=9F=E4=B8=80=E6=98=BE=E5=AD=98=E9=A2=84=E7=95=99=E5=A4=84?= =?UTF-8?q?=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增工具函数统一管理显存预留和DynamicVRAM同步逻辑 - 适配ComfyUI官方set_extra_reserved_vram API，兼容旧版本构建 - 同步更新comfy-aimdo的simple vram headroom配置 --- README.md | 6 ++++ nodes.py | 86 +++++++++++++++++++++++++++++++++++++------------------ 2 files changed, 64 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 74f2165..f67ed6a 100644 --- a/README.md +++ b/README.md @@ -37,3 +37,9 @@ new ![_$)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. +- On ComfyUI builds that provide `model_management.set_extra_reserved_vram()`, the node uses that runtime API. On official builds without that API, it still updates `EXTRA_RESERVED_VRAM` and then syncs DynamicVRAM headroom from the node itself. +- Auto mode can use torch CUDA memory info when NVML is unavailable or fails to initialize. diff --git a/nodes.py b/nodes.py index f5450c4..72eb483 100644 --- a/nodes.py +++ b/nodes.py @@ -25,11 +25,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,14 +51,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 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) @@ -142,21 +172,21 @@ class ReservedVRAMSetter: 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)') - model_management.EXTRA_RESERVED_VRAM = int(auto_reserved * 1024 * 1024 * 1024) - final_reserved_vram = round(auto_reserved, 2) - else: - manual_reserved = max(0, reserved) - model_management.EXTRA_RESERVED_VRAM = int(manual_reserved * 1024 * 1024 * 1024) - print(f'[ReservedVRAM]set EXTRA_RESERVED_VRAM={manual_reserved}GB (自动模式不可用，使用手动值)') - final_reserved_vram = round(manual_reserved, 2) - else: - # 手动模式 - reserved = max(0, reserved) - model_management.EXTRA_RESERVED_VRAM = int(reserved * 1024 * 1024 * 1024) - 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)