From 6fa8778ac4ccf1cec1fc8def961397a305316387 Mon Sep 17 00:00:00 2001
From: Windecay <4453636@qq.com>
Date: Sun, 29 Jun 2025 17:09:00 +0800
Subject: [PATCH] Update nodes.py

---
 nodes.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/nodes.py b/nodes.py
index fe3cde0..3fa1e7f 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1,5 +1,5 @@
-from comfy.model_management import set_extra_reserved_vram
 from typing import Any as any_type
+from comfy import model_management
 
 class AlwaysEqualProxy(str):
     def __eq__(self, _):
@@ -29,15 +29,14 @@ class ReservedVRAMSetter:
     OUTPUT_NODE = True
     FUNCTION = "set_vram"
     CATEGORY = "VRAM"
-
     def set_vram(self, anything, reserved, unique_id=None, extra_pnginfo=None):
-        set_extra_reserved_vram(reserved)
+        model_management.EXTRA_RESERVED_VRAM = int(reserved * 1024 * 1024 * 1024)
+        print(f'set EXTRA_RESERVED_VRAM={reserved}GB')
         return (anything,)
 
 NODE_CLASS_MAPPINGS = {
     "ReservedVRAMSetter": ReservedVRAMSetter
 }
-
 NODE_DISPLAY_NAME_MAPPINGS = {
     "ReservedVRAMSetter": "Set Reserved VRAM(GB) ⚙️"
 }