diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..ed8d3d5 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,27 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + paths: + - "pyproject.toml" + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'LAOGOU-666' }} + steps: + - name: Check out code + uses: actions/checkout@v4 + with: + submodules: true + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@v1 + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/1.png b/1.png new file mode 100644 index 0000000..a7280b5 Binary files /dev/null and b/1.png differ diff --git a/2.png b/2.png new file mode 100644 index 0000000..73fa988 Binary files /dev/null and b/2.png differ diff --git a/README.md b/README.md index 73eb01c..0aabf22 100644 Binary files a/README.md and b/README.md differ diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..f02e91e --- /dev/null +++ b/__init__.py @@ -0,0 +1,143 @@ +import psutil +import ctypes +from ctypes import wintypes +import time + +class AnyType(str): + """用于表示任意类型的特殊类,在类型比较时总是返回相等""" + def __eq__(self, _) -> bool: + return True + + def __ne__(self, __value: object) -> bool: + return False + +any = AnyType("*") + +class VRAMCleanup: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "anything": (any, {}), + "offload_model": ("BOOLEAN", {"default": True}), + "offload_cache": ("BOOLEAN", {"default": True}), + }, + "optional": {}, + "hidden": { + "unique_id": "UNIQUE_ID", + "extra_pnginfo": "EXTRA_PNGINFO", + } + } + + RETURN_TYPES = (any,) + RETURN_NAMES = ("output",) + OUTPUT_NODE = True + FUNCTION = "empty_cache" + CATEGORY = "Memory Management" + + def empty_cache(self, anything, offload_model, offload_cache, unique_id=None, extra_pnginfo=None): + try: + if offload_model: + import comfy.model_management + comfy.model_management.unload_all_models() + if offload_cache: + import torch + import gc + torch.cuda.empty_cache() + gc.collect() + except Exception as e: + print(f"内存清理出错: {str(e)}") + import traceback + print(traceback.format_exc()) + + return (anything,) + + +class RAMCleanup: + @classmethod + 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"}), + "retry_times": ("INT", { + "default": 3, + "min": 1, + "max": 10, + "step": 1, + "label": "重试次数" + }), + }, + "optional": {}, + "hidden": { + "unique_id": "UNIQUE_ID", + "extra_pnginfo": "EXTRA_PNGINFO", + } + } + + RETURN_TYPES = (any,) + RETURN_NAMES = ("output",) + OUTPUT_NODE = True + FUNCTION = "clean_ram" + CATEGORY = "Memory Management" + + 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): + try: + current_usage, available_mb = self.get_ram_usage() + print(f"开始清理RAM - 当前使用率: {current_usage:.1f}%, 可用: {available_mb:.1f}MB") + + for attempt in range(retry_times): + + if clean_file_cache: + try: + ctypes.windll.kernel32.SetSystemFileCacheSize(-1, -1, 0) + except Exception as e: + print(f"清理文件缓存失败: {str(e)}") + + if clean_processes: + cleaned_processes = 0 + for process in psutil.process_iter(['pid', 'name']): + try: + handle = ctypes.windll.kernel32.OpenProcess( + wintypes.DWORD(0x001F0FFF), + wintypes.BOOL(False), + wintypes.DWORD(process.info['pid']) + ) + ctypes.windll.psapi.EmptyWorkingSet(handle) + ctypes.windll.kernel32.CloseHandle(handle) + cleaned_processes += 1 + except: + continue + + if clean_dlls: + try: + ctypes.windll.kernel32.SetProcessWorkingSetSize(-1, -1, -1) + except Exception as e: + print(f"释放DLL失败: {str(e)}") + + 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") + + except Exception as e: + print(f"RAM清理过程出错: {str(e)}") + + return (anything,) + +NODE_CLASS_MAPPINGS = { + "VRAMCleanup": VRAMCleanup, + "RAMCleanup": RAMCleanup, + +} +NODE_DISPLAY_NAME_MAPPINGS = { + "VRAMCleanup": "🎈VRAM-Cleanup", + "RAMCleanup": "🎈RAM-Cleanup", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..689189a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[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.0.0" +dependencies = ["psutil"] + +[project.urls] +Repository = "https://github.com/LAOGOU-666/Comfyui-Memory_Cleanup" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "laogou666" +DisplayName = "Comfyui-Memory_Cleanup" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..0ba111c --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +psutil \ No newline at end of file