修改为纯后端清理，同步执行，使用off_load会清理缓存并重新加载模型

2025-11-04 02:34:26 +08:00
parent 7ba3265c18
commit 34278a55d6
3 changed files with 45 additions and 80 deletions
@@ -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,
@@ -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]
@@ -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);
-                });
-            }
-        });
-    }
-});