diff --git a/README.md b/README.md index 9f7f2fd..5012f7f 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,7 @@ Tips: base64格式字符串比较长,会导致界面卡顿,接口请求带 | × | ConvertToJsonStr | 序列化为json字符串 | | × | SaveTextToLocalFile | 保存文本到本地文件 | | × | ReadTextFromLocalFile | 从本地文件读取文本 | +| × | TryFreeMemory | 回收内(显)存 | ### 示例 ![save api extended](docs/example_note.png) diff --git a/easyapi/UtilNode.py b/easyapi/UtilNode.py index 670f99c..6b98864 100644 --- a/easyapi/UtilNode.py +++ b/easyapi/UtilNode.py @@ -1,3 +1,4 @@ +import gc import mimetypes import os import shutil @@ -9,6 +10,7 @@ import torch import folder_paths from comfy.model_patcher import ModelPatcher import comfy.model_base +import comfy.model_management as mm from .util import tensor_to_pil, hex_to_rgba, any_type @@ -1030,6 +1032,30 @@ class CopyAndRenameFiles: return (save_directory, ) +class TryFreeMemory: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "a": (any_type, {"forceInput": True}), + "do_gc": ("BOOLEAN", {"default": False}), + "unload_models": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = (any_type,) + RETURN_NAMES = ("a",) + FUNCTION = "execute" + CATEGORY = "EasyApi/Utils" + def execute(self, a, do_gc, unload_models): + if unload_models: + mm.unload_all_models() + if do_gc: + # 因为comfyui有缓存机制,还取决于自定义节点是否及时删除引用,效果不大好 + gc.collect() + return (a,) + + NODE_CLASS_MAPPINGS = { "GetImageBatchSize": GetImageBatchSize, "JoinList": JoinList, @@ -1065,6 +1091,7 @@ NODE_CLASS_MAPPINGS = { "CopyAndRenameFiles": CopyAndRenameFiles, "SaveTextToLocalFile": SaveTextToLocalFile, "ReadTextFromLocalFile": ReadTextFromLocalFile, + "TryFreeMemory": TryFreeMemory, } # A dictionary that contains the friendly/humanly readable titles for the nodes @@ -1103,4 +1130,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CopyAndRenameFiles": "CopyAndRenameFiles", "SaveTextToLocalFile": "SaveTextToLocalFile", "ReadTextFromLocalFile": "ReadTextFromLocalFile", + "TryFreeMemory": "TryFreeMemory", }