feat: integrate general_manage_vram into MemoryManager node (v2.4.0)

Author: ussoewwin

README.md

@@ -4,7 +4,7 @@
   <img src="https://raw.githubusercontent.com/ussoewwin/ComfyUI-DistorchMemoryManager/main/icon.png" width="128">
 </p>
 
-**ComfyUI-VRAM-Manager** (formerly ComfyUI-DistorchMemoryManager) is an independent memory management custom node for ComfyUI. Provides Distorch memory management functionality for efficient GPU/CPU memory handling. Supports purging of SeedVR2, Qwen3-VL, and Nunchaku models (FLUX/Z-Image/Qwen-Image). Includes Model Patch Memory Cleaner for ModelPatchLoader workflows.
+**ComfyUI-VRAM-Manager** (formerly ComfyUI-DistorchMemoryManager) is an independent memory management custom node for ComfyUI. Provides Distorch memory management functionality for efficient GPU/CPU memory handling. Supports purging of SeedVR2, Qwen3-VL, and Nunchaku models (FLUX/Z-Image/Qwen-Image). Includes Model Patch Memory Cleaner for ModelPatchLoader workflows. Auto-detects non-PyTorch VRAM usage via NVML to prevent OOM errors in multi-process environments.
 
 ## Overview
 
@@ -97,7 +97,7 @@ This is a completely original implementation designed specifically for Distorch
 </p>
 
 * **Description**: Comprehensive memory management node (for advanced users)
-* **Features**: Detailed memory management with UI corruption protection
+* **Features**: Detailed memory management with UI corruption protection and general VRAM management
 * **Input**: Any data type (ANY)
 * **Output**: Any data type (ANY)
 * **Options**:  
@@ -106,6 +106,7 @@ This is a completely original implementation designed specifically for Distorch
    * `force_gc`: Force garbage collection  
    * `reset_virtual_memory`: Reset virtual memory  
    * `restore_original_functions`: Restore original functions
+   * `general_manage_vram`: General VRAM management (v2.4.0) – Auto-detects non-PyTorch VRAM usage (browsers, Discord, OBS, etc.) via NVML and adjusts ComfyUI memory management accordingly. Prevents OOM errors caused by external GPU processes that PyTorch cannot see. Requires `nvidia-ml-py`.
 
 #### Patch Sage Attention DM (New in v2.3.0)
 
@@ -295,6 +296,7 @@ pip install -r requirements.txt
 * For Qwen3-VL workflows: Use DisTorchPurgeVRAMV2 with `purge_qwen3vl_models: True` after Qwen3-VL model usage to prevent OOM. The node automatically handles device_map="auto" case for models distributed across multiple devices.
 * For Nunchaku workflows (FLUX/Z-Image/Qwen-Image/SDXL): Use DisTorchPurgeVRAMV2 with `purge_nunchaku_models: True` after Nunchaku model usage to prevent OOM. The node automatically disables CPU offload and clears models from all detection locations (sys.modules, ComfyUI model management, and gc.get_objects()). For Nunchaku SDXL models (v2.2.0), the node now includes cache clearing functionality that can release approximately 2.5GB of VRAM.
 * For SageAttention workflows (v2.3.0): Use Patch Sage Attention DM node to replace ComfyUI's attention mechanism with SageAttention for improved memory efficiency and performance. The node supports multiple SageAttention implementations and automatically patches attention on each model execution. To disable SageAttention, run the node again with `sage_attention` set to `disabled`.
+* For multi-process environments (v2.4.0): Enable `general_manage_vram` in the Memory Manager node. This uses NVML to detect VRAM consumed by non-PyTorch processes (browsers, Discord, OBS, etc.) and automatically informs ComfyUI to account for it, preventing OOM errors that occur when ComfyUI overestimates available VRAM.
 
 ## License
 

changelog/changelog.md

@@ -1,5 +1,6 @@
 # Release History
 
+* **v2.4.0** – Added `general_manage_vram` parameter to Memory Manager node. Auto-detects non-PyTorch VRAM usage (browsers, Discord, OBS, etc.) via NVML and adjusts ComfyUI memory management accordingly. This prevents OOM errors caused by external GPU processes that PyTorch cannot detect. Added `nvidia-ml-py` dependency. Based on [ComfyUI-ReservedVRAM](https://github.com/Windecay/ComfyUI-ReservedVRAM) by Windecay (Apache-2.0).
 * **v2.3.8** – Technical documentation describing Flash Attention-2 issues when using ComfyUI on PyTorch 2.11.0, including failure symptoms and environment-related constraints. See [Release Notes v2.3.8](https://github.com/ussoewwin/ComfyUI-DistorchMemoryManager/releases/tag/v2.3.8) for details.
 * **v2.3.7** – Added external runtime SageAttention noise guard behavior for the known `Unsupported head_dim: 160` fallback path, reduced repeated error-log spam, and aligned patch documentation to repository-scoped changes only. See [Release Notes v2.3.7](https://github.com/ussoewwin/ComfyUI-DistorchMemoryManager/releases/tag/v2.3.7) for details.
 * **v2.3.6** – Enhanced SageAttention3 (SA3) integration in Patch Sage Attention DM node. Added SA3-specific version detection function (`get_sage_attention3_info()`) with Blackwell support detection. Improved SA3 implementation with tensor layout conversion (NHD to HND) and constraint handling (headdim >= 256, attention mask support). Added automatic fallback to PyTorch SDPA when SA3 constraints are not met. Fixed SA2 version logging to skip when SA3 modes are selected. Supports both `sageattn3` and `sageattn3_per_block_mean` modes with proper per-block mean processing. See [Release Notes v2.3.6](https://github.com/ussoewwin/ComfyUI-DistorchMemoryManager/releases/tag/v2.3.6) for details.

nodes/memory_manager.py

@@ -1,9 +1,54 @@
 """
 Memory Manager nodes for DistorchMemoryManager
+
+Reserved VRAM auto-detection based on ComfyUI-ReservedVRAM by Windecay
+https://github.com/Windecay/ComfyUI-ReservedVRAM
+Original code licensed under Apache License 2.0
 """
 import torch
 import gc
 
+# Safe import for pynvml (NVIDIA Management Library)
+try:
+    import pynvml
+    try:
+        pynvml.nvmlInit()
+        _pynvml_available = True
+    except Exception as e:
+        _pynvml_available = False
+        pynvml = None
+        print(f"[ComfyUI-VRAM-Manager] WARNING: pynvml imported but NVML init failed: {e}")
+except ImportError:
+    _pynvml_available = False
+    pynvml = None
+    print("[ComfyUI-VRAM-Manager] INFO: pynvml (nvidia-ml-py) not installed. general_manage_vram will be unavailable.")
+
+
+def get_non_torch_vram_usage_bytes():
+    """
+    Returns the amount of VRAM (in bytes) consumed by non-PyTorch processes.
+    This is the difference between system-wide GPU usage (via NVML) and
+    PyTorch's own reported usage.
+    Returns None if detection is not possible.
+    """
+    if not _pynvml_available or pynvml is None:
+        return None
+    if not torch.cuda.is_available():
+        return None
+    try:
+        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        nvml_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        system_used = nvml_info.used  # bytes, all processes
+
+        torch_free, torch_total = torch.cuda.mem_get_info()
+        torch_used = torch_total - torch_free  # bytes, PyTorch only
+
+        non_torch = system_used - torch_used
+        return max(0, non_torch)
+    except Exception as e:
+        print(f"[ComfyUI-VRAM-Manager] Error detecting non-PyTorch VRAM usage: {e}")
+        return None
+
 # AnyType mirrors the behavior of the original Purge VRAM node
 class AnyType(str):
     """A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
@@ -38,14 +83,15 @@ def INPUT_TYPES(s):
             "force_gc": ("BOOLEAN", {"default": True}),
             "reset_virtual_memory": ("BOOLEAN", {"default": True}),
             "restore_original_functions": ("BOOLEAN", {"default": False, "tooltip": "Restore original model_management functions"}),
+            "general_manage_vram": ("BOOLEAN", {"default": False, "tooltip": "General VRAM management: auto-detect non-PyTorch VRAM usage (browsers, other apps) via NVML and adjust ComfyUI memory management accordingly. Requires nvidia-ml-py."}),
         }}
 
     RETURN_TYPES = (any,)
     RETURN_NAMES = ("any",)
     FUNCTION = "manage_memory"
     CATEGORY = "Memory"
 
-    def manage_memory(self, anything, clean_gpu, clean_cpu, force_gc, reset_virtual_memory, restore_original_functions):
+    def manage_memory(self, anything, clean_gpu, clean_cpu, force_gc, reset_virtual_memory, restore_original_functions, general_manage_vram=False):
         try:
             if clean_gpu and torch.cuda.is_available():
                 torch.cuda.empty_cache()
@@ -81,6 +127,17 @@ def manage_memory(self, anything, clean_gpu, clean_cpu, force_gc, reset_virtual_
                 except Exception as e:
                     print(f"Function restoration failed: {e}")
 
+            # Auto-detect non-PyTorch VRAM usage and adjust ComfyUI memory management
+            if general_manage_vram:
+                non_torch = get_non_torch_vram_usage_bytes()
+                if non_torch is not None:
+                    import comfy.model_management as mm
+                    mm.EXTRA_RESERVED_VRAM = non_torch
+                    non_torch_gb = non_torch / (1024 * 1024 * 1024)
+                    print(f"[ComfyUI-VRAM-Manager] Detected {non_torch_gb:.2f} GB non-PyTorch VRAM usage; adjusted memory management accordingly")
+                else:
+                    print("[ComfyUI-VRAM-Manager] general_manage_vram: Could not detect non-PyTorch VRAM usage (pynvml unavailable or error)")
+            
             print("Comprehensive memory management completed")
 
         except Exception as e:

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-DistorchMemoryManager"
-version = "2.3.8"
-description = "An independent memory management custom node for ComfyUI. Provides Distorch memory management functionality for efficient GPU/CPU memory handling. Supports purging of SeedVR2, Qwen3-VL, and Nunchaku models (FLUX/Z-Image/Qwen-Image). Includes Model Patch Memory Cleaner for ModelPatchLoader workflows."
+version = "2.4.0"
+description = "An independent memory management custom node for ComfyUI. Provides Distorch memory management functionality for efficient GPU/CPU memory handling. Supports purging of SeedVR2, Qwen3-VL, and Nunchaku models (FLUX/Z-Image/Qwen-Image). Includes Model Patch Memory Cleaner for ModelPatchLoader workflows. Auto-detects non-PyTorch VRAM usage via NVML to prevent OOM errors in multi-process environments."
 authors = [{ name = "ussoewwin" }]
 license = { file = "LICENSE" }
 requires-python = ">=3.10, <3.13"

requirements.txt

@@ -1,2 +1,3 @@
 torch>=1.12.0
-psutil>=5.8.0 
+psutil>=5.8.0
+nvidia-ml-py