fix import, update README
This commit is contained in:
@@ -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" />
|
||||
|
||||

|
||||
|
||||
2025-10-10新增自动模式,自动模式会检测系统“已使用”的显存数量,再叠加用户设置值进行预留。避免多进程用户因为显存问题卡住运行。
|
||||
|
||||
@@ -30,14 +31,6 @@ new
|
||||
|
||||
接在排行较前的节点处即可,观察windows任务管理器共享显存溢出多少,就需要设置保留多少(可以略微多一点),填入该数值。运行工作流实时生效,输入单位是GB。
|
||||
|
||||

|
||||
|
||||

|
||||
|
||||

|
||||
|
||||

|
||||
|
||||
2026-06-27 DynamicVRAM compatibility update
|
||||
|
||||
- The node now updates ComfyUI DynamicVRAM / comfy-aimdo simple vram headroom when DynamicVRAM is enabled.
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 165 KiB |
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user