From 8b5ee46a2c8be8f88081de5d6f97e17d10a57c6a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E9=9B=AA=E5=B3=B0?= <lldacing@163.com>
Date: Sat, 21 Dec 2024 10:44:32 +0800
Subject: [PATCH] fix TryFreeMemory

---
 easyapi/UtilNode.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/easyapi/UtilNode.py b/easyapi/UtilNode.py
index 6b98864..ef39c36 100644
--- a/easyapi/UtilNode.py
+++ b/easyapi/UtilNode.py
@@ -11,6 +11,7 @@ import folder_paths
 from comfy.model_patcher import ModelPatcher
 import comfy.model_base
 import comfy.model_management as mm
+from server import PromptServer
 from .util import tensor_to_pil, hex_to_rgba, any_type
 
 
@@ -1038,8 +1039,8 @@ class TryFreeMemory:
         return {
             "required": {
                 "a": (any_type, {"forceInput": True}),
-                "do_gc": ("BOOLEAN", {"default": False}),
-                "unload_models": ("BOOLEAN", {"default": False}),
+                "do_gc": ("BOOLEAN", {"default": False, "tooltip": "设置垃圾回收标志，不会立即执行垃圾回收，在本次流运行结束时ComfyUI主体会根据标识自动执行一次垃圾回收（相当于不会缓存结果）"}),
+                "unload_models": ("BOOLEAN", {"default": False, "tooltip": "释放显存中已经无效的模型，会立即执行一次"}),
             }
         }
 
@@ -1050,9 +1051,10 @@ class TryFreeMemory:
     def execute(self, a, do_gc, unload_models):
         if unload_models:
             mm.unload_all_models()
-        if do_gc:
-            # 因为comfyui有缓存机制，还取决于自定义节点是否及时删除引用，效果不大好
-            gc.collect()
+
+        PromptServer.instance.prompt_queue.set_flag("unload_models", unload_models)
+        PromptServer.instance.prompt_queue.set_flag("free_memory", do_gc)
+
         return (a,)