fix import, update README

This commit is contained in:
windecay
2026-07-04 14:05:38 +08:00
parent b9448f58fb
commit be93b9375e
3 changed files with 69 additions and 72 deletions
+2 -9
View File
@@ -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.
<img width="1919" height="1461" alt="image" src="https://github.com/user-attachments/assets/5b3af05d-5051-4fc9-b2e7-fd7cb7cfe719" />
![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.
BIN
View File
Binary file not shown.

After

Width:  |  Height:  |  Size: 165 KiB

+67 -63
View File
@@ -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)