From 34278a55d6affc488b6d7038b66ca0ee4b308e67 Mon Sep 17 00:00:00 2001 From: LAOGOU-666 <441059767@qq.com> Date: Tue, 4 Nov 2025 02:34:26 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E4=B8=BA=E7=BA=AF=E5=90=8E?= =?UTF-8?q?=E7=AB=AF=E6=B8=85=E7=90=86=EF=BC=8C=E5=90=8C=E6=AD=A5=E6=89=A7?= =?UTF-8?q?=E8=A1=8C=EF=BC=8C=E4=BD=BF=E7=94=A8off=5Fload=E4=BC=9A?= =?UTF-8?q?=E6=B8=85=E7=90=86=E7=BC=93=E5=AD=98=E5=B9=B6=E9=87=8D=E6=96=B0?= =?UTF-8?q?=E5=8A=A0=E8=BD=BD=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- __init__.py | 94 +++++++++++++++++++++++--------------------------- pyproject.toml | 2 +- web/vram.js | 29 ---------------- 3 files changed, 45 insertions(+), 80 deletions(-) delete mode 100644 web/vram.js diff --git a/__init__.py b/__init__.py index 16915c8..af1c31e 100644 --- a/__init__.py +++ b/__init__.py @@ -6,7 +6,9 @@ import os import platform import subprocess import tempfile +import gc from server import PromptServer +import comfy.model_management class AnyType(str): """用于表示任意类型的特殊类，在类型比较时总是返回相等""" @@ -23,11 +25,12 @@ class VRAMCleanup: def INPUT_TYPES(s): return { "required": { - "anything": (any, {}), "offload_model": ("BOOLEAN", {"default": True}), "offload_cache": ("BOOLEAN", {"default": True}), }, - "optional": {}, + "optional": { + "anything": (any, {}), + }, "hidden": { "unique_id": "UNIQUE_ID", "extra_pnginfo": "EXTRA_PNGINFO", @@ -39,23 +42,27 @@ class VRAMCleanup: OUTPUT_NODE = True FUNCTION = "empty_cache" CATEGORY = "Memory Management" + + @classmethod + def IS_CHANGED(s, **kwargs): + # 返回当前时间戳，确保每次都执行 + return float(time.time()) - def empty_cache(self, anything, offload_model, offload_cache, unique_id=None, extra_pnginfo=None): + def empty_cache(self, offload_model, offload_cache, anything=None, unique_id=None, extra_pnginfo=None): try: - # 发送信号到前端 - PromptServer.instance.send_sync("memory_cleanup", { - "type": "cleanup_request", - "data": { - "unload_models": offload_model, - "free_memory": offload_cache - } - }) - print("已发送内存清理信号") + if offload_model: + comfy.model_management.unload_all_models() + + if offload_cache: + gc.collect() + comfy.model_management.soft_empty_cache() + PromptServer.instance.prompt_queue.set_flag("free_memory", True) + + print(f"VRAM清理完成 [卸载模型: {offload_model}, 清空缓存: {offload_cache}]") + except Exception as e: - print(f"发送内存清理信号失败: {str(e)}") - import traceback - print(traceback.format_exc()) - time.sleep(1) + print(f"VRAM清理失败: {str(e)}") + return (anything,) @@ -64,7 +71,6 @@ class RAMCleanup: def INPUT_TYPES(s): return { "required": { - "anything": (any, {}), "clean_file_cache": ("BOOLEAN", {"default": True, "label": "清理文件缓存"}), "clean_processes": ("BOOLEAN", {"default": True, "label": "清理进程内存"}), "clean_dlls": ("BOOLEAN", {"default": True, "label": "清理未使用DLL"}), @@ -76,7 +82,9 @@ class RAMCleanup: "label": "重试次数" }), }, - "optional": {}, + "optional": { + "anything": (any, {}), + }, "hidden": { "unique_id": "UNIQUE_ID", "extra_pnginfo": "EXTRA_PNGINFO", @@ -88,42 +96,33 @@ class RAMCleanup: OUTPUT_NODE = True FUNCTION = "clean_ram" CATEGORY = "Memory Management" + + @classmethod + def IS_CHANGED(s, **kwargs): + # 返回当前时间戳，确保每次都执行 + return float(time.time()) def get_ram_usage(self): memory = psutil.virtual_memory() return memory.percent, memory.available / (1024 * 1024) - def clean_ram(self, anything, clean_file_cache, clean_processes, clean_dlls, retry_times, unique_id=None, extra_pnginfo=None): + def clean_ram(self, clean_file_cache, clean_processes, clean_dlls, retry_times, anything=None, unique_id=None, extra_pnginfo=None): try: - current_usage, available_mb = self.get_ram_usage() - print(f"开始清理RAM - 当前使用率: {current_usage:.1f}%, 可用: {available_mb:.1f}MB") - + before_usage, before_available = self.get_ram_usage() system = platform.system() + for attempt in range(retry_times): - if clean_file_cache: try: if system == "Windows": ctypes.windll.kernel32.SetSystemFileCacheSize(-1, -1, 0) elif system == "Linux": - try: - subprocess.run(["sudo", "sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"], - check=False, stderr=subprocess.PIPE, stdout=subprocess.PIPE) - print("使用sudo清理缓存成功") - except Exception as sudo_e: - print(f"使用sudo清理缓存失败: {str(sudo_e)}") - try: - subprocess.run(["sudo", "sysctl", "vm.drop_caches=3"], - check=False, stderr=subprocess.PIPE, stdout=subprocess.PIPE) - print("使用sysctl清理缓存成功") - except Exception as sysctl_e: - print(f"使用sysctl清理缓存失败: {str(sysctl_e)}") - print("请尝试在终端执行: 'sudo sh -c \"echo 3 > /proc/sys/vm/drop_caches\"'") - except Exception as e: - print(f"清理文件缓存失败: {str(e)}") + subprocess.run(["sudo", "sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"], + check=False, stderr=subprocess.PIPE, stdout=subprocess.PIPE) + except: + pass if clean_processes: - cleaned_processes = 0 if system == "Windows": for process in psutil.process_iter(['pid', 'name']): try: @@ -134,35 +133,30 @@ class RAMCleanup: ) ctypes.windll.psapi.EmptyWorkingSet(handle) ctypes.windll.kernel32.CloseHandle(handle) - cleaned_processes += 1 except: continue - elif system == "Linux": - cleaned_processes = 0 if clean_dlls: try: if system == "Windows": ctypes.windll.kernel32.SetProcessWorkingSetSize(-1, -1, -1) elif system == "Linux": - subprocess.run(["sync"], check=True) - except Exception as e: - print(f"释放内存资源失败: {str(e)}") + except: + pass time.sleep(1) - current_usage, available_mb = self.get_ram_usage() - print(f"清理后内存使用率: {current_usage:.1f}%, 可用: {available_mb:.1f}MB") - print(f"清理完成 - 最终内存使用率: {current_usage:.1f}%, 可用: {available_mb:.1f}MB") + after_usage, after_available = self.get_ram_usage() + freed_mb = after_available - before_available + print(f"RAM清理完成 [{before_usage:.1f}% → {after_usage:.1f}%, 释放: {freed_mb:.0f}MB]") except Exception as e: - print(f"RAM清理过程出错: {str(e)}") + print(f"RAM清理失败: {str(e)}") return (anything,) -WEB_DIRECTORY = "web" NODE_CLASS_MAPPINGS = { "VRAMCleanup": VRAMCleanup, "RAMCleanup": RAMCleanup, diff --git a/pyproject.toml b/pyproject.toml index bec680e..a6dd0a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_memory_cleanup" description = "A ComfyUI extension that provides nodes for memory cleanup, including VRAM and RAM cleanup functions to optimize ComfyUI performance during long running workflows." -version = "1.1.1" +version = "1.1.2" dependencies = ["psutil"] [project.urls] diff --git a/web/vram.js b/web/vram.js deleted file mode 100644 index 6c9bdc7..0000000 --- a/web/vram.js +++ /dev/null @@ -1,29 +0,0 @@ -import { app } from "../../scripts/app.js"; -import { api } from "../../scripts/api.js"; -app.registerExtension({ - name: "memory.cleanup", - init() { - api.addEventListener("memory_cleanup", ({ detail }) => { - if (detail.type === "cleanup_request") { - console.log("收到内存清理请求"); - fetch("/free", { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify(detail.data) - }) - .then(response => { - if (response.ok) { - console.log("内存清理请求已发送"); - } else { - console.error("内存清理请求失败"); - } - }) - .catch(error => { - console.error("发送内存清理请求出错:", error); - }); - } - }); - } -});