Initial project release v1.0.0
This commit is contained in:
mkim87404
@@ -0,0 +1,2 @@
|
||||
# Auto detect text files and perform LF normalization
|
||||
* text=auto
|
||||
@@ -0,0 +1,72 @@
|
||||
# ComfyUI-ControlOrder-FreeMemory
|
||||
|
||||
[](https://github.com/comfyanonymous/ComfyUI)
|
||||
[](https://www.python.org/)
|
||||
[](https://pytorch.org/)
|
||||
[](LICENSE)
|
||||
|
||||
`🔷 Control Order & Free Memory` is a dependency-free ComfyUI custom node for controlling the execution order of nodes by routing any type/number of data through it, ensuring all input connections resolve first before being passed to outputs. Optionally free VRAM & RAM at any point in a workflow using device-agnostic & robust memory management utilities managed by ComfyUI that safely unload all models (except any models passing though this node) and release as much VRAM & RAM as possible, while preserving all connected passthrough data (e.g. latents, conditioning, images, masks, models, etc.) through to the next node. This node will also print a descriptive message on the ComfyUI session terminal showing how much VRAM & RAM has been freed.
|
||||
|
||||
All input & output slots are AnyType (*). They can hook onto any node, including loader nodes like "Load Model", "Load VAE", "Load CLIP", etc. that usually don't have any connectable input slots, making linear connection/chaining of all node types possible - enforcing a single, sequential & deterministic flow of execution. This node adaptively expands its input & output slots to support infinitely many passthrough data (persist_any_N), passing every input to their corresponding "N"th output slot unchanged.
|
||||
|
||||
The memory clean up + linear passthrough of data facilitated by this node can help to **speed up workflows** involving multiple models (e.g. CLIP, VAE, Diffusion Model, etc.) that compete for the same VRAM/RAM but don't all need to be loaded at the same time, by separating each model's operation into its own segment and linearly chaining them together with multiple instances of this node, where only the necessary outputs of each segment (e.g. latents, conditioning, images, models, etc.) are passed on to the next segment while unloading all/some models that are no longer needed in the workflow.
|
||||
|
||||
This becomes especially effective in freeing significant amount of VRAM/RAM between multiple KSamplers (e.g. Wan 2.2 High & Low KSamplers), or before & after VAE Encode / VAE Decode / Load CLIP / Load Model / etc. that allow more of the models' layers to be loaded onto the GPU VRAM and less offloaded to the slower System RAM, resulting in faster overall inference times and letting the user push to higher resolution media outputs before hitting OOM. I have personally seen great reductions in total execution time of my workflows using this node.
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
1. Clone the repo into your `ComfyUI/custom_nodes` folder:
|
||||
```bash
|
||||
git clone https://github.com/mkim87404/ComfyUI-ControlOrder-FreeMemory.git
|
||||
```
|
||||
2. Restart the ComfyUI server, double click anywhere in the workflow and search for "**Control Order & Free Memory**" or alternatively, [Right click > Add Node > Control Order & Free Memory]
|
||||
|
||||
## 🔄 Update
|
||||
|
||||
1. To get the latest updates, run `git pull` inside the node's folder:
|
||||
|
||||
```bash
|
||||
cd custom_nodes/ComfyUI-ControlOrder-FreeMemory
|
||||
git pull
|
||||
```
|
||||
|
||||
## 🎨 Example Workflow
|
||||
|
||||
In the `example_workflows/` folder of this repo, I have included example Wan 2.2 T2V & I2V ComfyUI workflows using this node. Simply drag and drop the .json files into the ComfyUI browser tab to start using them. These workflows demonstrate how this node can be used to achieve a single sequential flow of execution, loading & unloading only the necessary model at each section of the workflow before progressing to the next.
|
||||
|
||||
## 🚀 Usage & Tips
|
||||
|
||||
- **Control execution order & free memory:** The `🔷 Control Order & Free Memory` node ensures that all of the input-connected nodes finish executing first before the output-connected nodes start executing. Hence to control the execution order of nodes, you can split your workflow into multiple groups of nodes with differing execution priorities and connect all nodes that should execute first to the input slots of this node, and all nodes that should execute later to the output slots of this node, ensuring that the correct data is being passed from the "persist_any_N" input slot to its corresponding "N"th output slot. You can then chain additional groups of nodes with as many `🔷 Control Order & Free Memory` nodes as you need, until all groups of nodes are connected together to form a single, sequential flow of execution from start to finish. Optionally, toggle "free_memory" ON on this node to unload all models and free as much VRAM & RAM as possible at the node's position, while preserving all the passthrough data. You can connect as many passthrough data as you need through this node, and each of them can be a different data type (e.g. latents, conditioning, images, models, etc.). If any of the passthrough data are models, they will stay loaded and survive the "free_memory" operation. Also, when "free_memory" is ON, this node will print a descriptive message on the ComfyUI session terminal where you can check how much VRAM & RAM has been freed.
|
||||
|
||||
- **Connecting this node into loader-type nodes:** Because all input & output slots are AnyType (*), they can hook onto any node, including loader nodes like "Load Model", "Load VAE", "Load CLIP", etc. And when hooking the output of this node into those loader-type nodes, you will often notice that the only argument slot available for connection are "filename" combo/dropdown fields (STRING) for selecting the model to load. To ensure that we feed the correct value into this field, I have included a simple helper node called `📁 Filename Selector` (already included in this repo), which you can search & drop into the workflow and simply select a filaname from your native OS file explorer, and use/edit that filename string to feed it into the input slot of the `🔷 Control Order & Free Memory` node that corresponds to the output slot that is connected to the loader node's "filename" field.
|
||||
|
||||
- **Using this node at the start of the workflow:** This node can be placed anywhere in the workflow (e.g. at the very start, in the middle, at the very end, etc.), as all input & output connections are optional. And if you want this node to be at the start of the workflow e.g. to run the memory cleanup at the start of every run, you can connect ComfyUI's default `Int (Comfy Core)` node from [utils > primitive] into any input slot of `🔷 Control Order & Free Memory` node, and set the integer value to change (increment/decrement) on each run. This is to alter the input being fed into the node to trigger its memory cleanup logic at the start of each run, because ComfyUI is deigned to skip executing nodes whose inputs didn't change (which is not an issue when this node is used in the middle or at the end of a workflow). And to re-run the workflow and resume executing from the final state & outputs of the previous run, simply restore the integer value of the `Int (Comfy Core)` node back to the previous run's value, and this will ensure that running the workflow again picks up from the latest run's outputs without discarding them or incurring a re-execution of all the nodes from scratch.
|
||||
|
||||
- **For any uncertainties,** please refer to my included example workflow to see a working example of these nodes.
|
||||
|
||||
## 🍃 Features Summary
|
||||
|
||||
- Controls the execution order of nodes in any ComfyUI workflow to enforce a single, deterministic sequence of ordered logic.
|
||||
- Securely unloads models and frees VRAM & RAM at any point in a workflow, using ComfyUI-managed & device-agnostic memory management utilities.
|
||||
- All nodes are dependency-free & platform/device-agnostic.
|
||||
- Adaptively expands input & output slots to support infinitely many dynamic passthrough data of any type (*).
|
||||
- Linear passthrough design that focuses on one group of related nodes at a time and routes only necessary outputs to the next group of nodes helps with "cable management", tidying up the connections that can get messy as the workflow grows large, making them easier to understand and maintain.
|
||||
|
||||
## ℹ️ Notes
|
||||
|
||||
- Please note that this node is built upon the current ComfyUI implementation of the "comfy.model_management.py" module for memory management, and ComfyUI extension hooks & prototype hijacking patterns for making dynamic updates to the node's UI elements. Hence, it is possible that some features of the node may stop working in future updates of ComfyUI. For more information, you may find these Docs useful:
|
||||
- [ComfyUI model management](https://github.com/Comfy-Org/ComfyUI/blob/master/comfy/model_management.py)
|
||||
- [ComfyUI custom extension hooks](https://docs.comfy.org/custom-nodes/js/javascript_hooks)
|
||||
- [LiteGraph.js](https://tamats.com/projects/litegraph/doc/classes/LGraphNode.html)
|
||||
- You may notice that system RAM usage often does not drop as significantly as GPU VRAM does after this node's "free_memory" operation runs. This is expected behavior in the current implementation of ComfyUI, as ComfyUI intentionally caches the output of every executed node in RAM (e.g. latents, images, conditionings, tensors, etc.) with strong references that survive any py garbage collection invoked from inside the custom node's logic, and only releases those references after the entire "prompt execution" finishes (see [execution.py](https://github.com/Comfy-Org/ComfyUI/blob/master/execution.py), [comfy_execution/caching.py](https://github.com/Comfy-Org/ComfyUI/blob/master/comfy_execution/caching.py)). This is an intentional optimization design by ComfyUI so that downstream nodes can reuse upstream outputs instantly without re-computation. The "free_memory" operation of this custom node already does the maximum possible cleanup with (model unloading + explicit gc.collect() + double device cache clearing before and after the garbage collection to collect as much lingering tensors as possible, etc.). Hence, any RAM that remains allocated after this node runs is ComfyUI's deliberate output caching, which is best to leave up to ComfyUI to manage & optimize behind the scenes across multiple workflow runs.
|
||||
|
||||
## ❤️ Inspired By
|
||||
|
||||
- [willblaschko](https://github.com/willblaschko/ComfyUI-Unload-Models) - Patterns using the "comfy.model_management" module
|
||||
- [Trung0246](https://github.com/Trung0246/ComfyUI-0246) - "ByPassTypeTuple" RETURN_TYPES trick
|
||||
|
||||
- There were several other custom nodes solving parts of the same problem of (memory management + routing + execution order control) but I couldn't find one that combines all tricks together into a single device-agnostic node with a simple design that can take infinitely many passthrough data of any type, which encouraged me to build this custom node. I'll likely be using this node quite regularly myself, so please keep an eye out for more updates!
|
||||
|
||||
## 📜 License
|
||||
|
||||
[MIT License](https://github.com/mkim87404/ComfyUI-ControlOrder-FreeMemory/blob/main/LICENSE) – feel free to use in any personal or commercial project, fork, or open issues/PRs – contributions and feedback all welcome!
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
from .nodes import ControlOrderFreeMemory, FileNameSelector
|
||||
|
||||
WEB_DIRECTORY = "./js"
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ControlOrderFreeMemory": ControlOrderFreeMemory,
|
||||
"FileNameSelector": FileNameSelector,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ControlOrderFreeMemory": "🔷 Control Order & Free Memory",
|
||||
"FileNameSelector": "📁 Filename Selector",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
__version__ = "1.0.0"
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,38 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "comfyui.adaptive.infinite.io",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
||||
if (nodeData.name !== "ControlOrderFreeMemory") return;
|
||||
|
||||
// Keep original trigger logic
|
||||
const originalOnConnectionsChange = nodeType.prototype.onConnectionsChange;
|
||||
|
||||
nodeType.prototype.onConnectionsChange = function (type, index, isConnected, linkInfo) {
|
||||
// Execute original trigger logic & keep its return value
|
||||
const originalReturnValue = originalOnConnectionsChange?.apply(this, arguments);
|
||||
|
||||
// New extended logic - Dynamically increase Input & Output slots on the "ControlOrderFreeMemory" node
|
||||
if (type === 1 && isConnected && linkInfo != null && this.inputs[index].name.startsWith("persist_any_") && (this.inputs.length === 2 || index === this.inputs.length - 1)) {
|
||||
// Only trigger on input connections (type === 1 && isConnected)
|
||||
// Do not trigger when the node is being copied over (linkInfo != null)
|
||||
// Only trigger slot expansion when the last input slot is connected (index === this.inputs.length - 1).
|
||||
|
||||
// Define & Add new Input & Output slots
|
||||
const newSlotName = `persist_any_${this.inputs.length}`
|
||||
const newSlotType = this.inputs[index].type || "*";
|
||||
|
||||
// All connected input slots will correctly send their data to **kwargs in Python backend
|
||||
this.addInput(newSlotName, newSlotType); // Currently, there is no supported way to set a tooltip on dynamically created slots.
|
||||
this.addOutput(newSlotName, newSlotType);
|
||||
|
||||
// Force UI refresh
|
||||
this.setDirtyCanvas(true, true);
|
||||
}
|
||||
// To-Do/Good-To-Have: Remove slots on disconnect? - Functionally redundant for now
|
||||
|
||||
// Return the original function return value, to not break anything.
|
||||
return originalReturnValue;
|
||||
};
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,49 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
app.registerExtension({
|
||||
name: "comfyui.filename.selector",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
||||
if (nodeData.name !== "FileNameSelector") return;
|
||||
|
||||
const originalOnNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
|
||||
nodeType.prototype.onNodeCreated = function () {
|
||||
const originalReturnValue = originalOnNodeCreated?.apply(this, arguments);
|
||||
|
||||
// Find the "string" input widget on the "FileNameSelector" node
|
||||
const stringInputWidget = this.widgets.find(w => w.name === "string");
|
||||
if (!stringInputWidget) return originalReturnValue;
|
||||
|
||||
// Add the "Select Filename" button widget on the "FileNameSelector" node
|
||||
this.addWidget(
|
||||
"button",
|
||||
"Select Filename",
|
||||
null,
|
||||
() => {
|
||||
// Create hidden native file input element without attaching it to DOM (i.e. the input element will be GC'd as soon as the function ends.)
|
||||
const input = document.createElement("input");
|
||||
input.type = "file";
|
||||
input.style.display = "none";
|
||||
|
||||
input.onchange = (e) => {
|
||||
const filename = e.target.files?.[0]?.name;
|
||||
if (filename) {
|
||||
stringInputWidget.value = filename;
|
||||
this.setDirtyCanvas(true, true);
|
||||
}
|
||||
|
||||
// Explicit cleanup (force immediate release & let the variables go out of scope available for GC)
|
||||
e.target.value = ""; // reset the input so the browser drops the FileList
|
||||
input.onchange = null; // break any reference cycle
|
||||
};
|
||||
|
||||
// Trigger the native OS file explorer
|
||||
input.click();
|
||||
},
|
||||
{ serialize: false } // the "string" input widget value on the "FileNameSelector" node will not be saved into the workflow JSON
|
||||
);
|
||||
|
||||
return originalReturnValue;
|
||||
};
|
||||
},
|
||||
});
|
||||
@@ -0,0 +1,199 @@
|
||||
import comfy.model_management as model_management
|
||||
import gc
|
||||
import torch
|
||||
import psutil
|
||||
|
||||
class AnyType(str):
|
||||
def __ne__(self, __value: object) -> bool: # Make this class instance be equal to everything, to bypass ComfyUI's type validation check.
|
||||
return False
|
||||
|
||||
anyType = AnyType("*")
|
||||
|
||||
class MatryoshkaTuple(tuple):
|
||||
def __getitem__(self, index): # Override the tuple "__getitem__" to always return the 1st element. This allows the node's return tuple to grow dynamically while ComfyUI believes the tuple length is still 1.
|
||||
if index > 0:
|
||||
index = 0
|
||||
return super().__getitem__(index)
|
||||
|
||||
class ControlOrderFreeMemory:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"free_memory": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Unload all models and release as much VRAM & RAM as possible while routing & preserving all 'persist_any' passthrough data. Any models passed into 'persist_any' will stay loaded. Prints how much VRAM & RAM has been freed on the ComfyUI session terminal.",
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"persist_any_1": (anyType, {
|
||||
"tooltip": "Persist any type of data through to the next node e.g. models, latents, conditioning, images, masks, etc. This data survives the 'free_memory' operation. I/O slots expand adaptively",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def VALIDATE_INPUTS(cls, **kwargs):
|
||||
return True
|
||||
|
||||
OUTPUT_NODE = True # Let this be a hybrid node that can both be in the middle or at the end of a workflow (with no outputs connected).
|
||||
RETURN_TYPES = MatryoshkaTuple((anyType, )) # Let this node's output slot cardinality grow dynamically from JS hooks, while ComfyUI type validation check believes this is just a tuple of 1 "AnyType" element.
|
||||
RETURN_NAMES = MatryoshkaTuple(("persist_any_1", )) # Same trick
|
||||
OUTPUT_TOOLTIPS = MatryoshkaTuple(("Persist any type of data through to the next node e.g. models, latents, conditioning, images, masks, etc. This data survives the 'free_memory' operation. I/O slots expand adaptively", ))
|
||||
FUNCTION = "passthrough"
|
||||
CATEGORY = "Control Order & Free Memory"
|
||||
DESCRIPTION = """Control the execution order of nodes by routing any data (as many as you need - I/O slots expand adaptively) through this node. Ensures all input-connected nodes finish executing first before the output-connected nodes start executing.
|
||||
|
||||
All input & output slots are AnyType (*). They can hook onto any node, including loader-type nodes like "Load Model", "Load VAE", "Load CLIP", etc. For connecting into loader-type nodes, you can use the "📁 Filename Selector" helper node (already installed with this custom node) to select & feed the filename into the input "persist_any_N" slot that corresponds to the output "persist_any_N" slot that is connecting into the loader-type node.
|
||||
|
||||
Optionally unload all models and release as much VRAM & RAM as possible while routing & preserving all 'persist_any' passthrough data. Any models passed into 'persist_any' will stay loaded. This node will also print how much VRAM & RAM has been freed on the ComfyUI session terminal.
|
||||
"""
|
||||
|
||||
def passthrough(self, **kwargs):
|
||||
# Collect every connected input (only connected slots appear in kwargs when the prompt is sent)
|
||||
input_keys = [key for key in kwargs if key.startswith("persist_any_")]
|
||||
|
||||
# Secure the input values in the exact order they were connected - to persist them throughout the potential memory cleanup operation
|
||||
if input_keys:
|
||||
# Sort the connected inputs in the exact order they were added
|
||||
input_keys.sort(key=lambda input_key: int(input_key[12:]))
|
||||
output_values = [None] * int(input_keys[-1][12:]) # Only need to define output up to the highest seen persist_any_N, and let ComfyUI output None for any undefined outputs beyond N.
|
||||
for input_key in input_keys:
|
||||
output_values[int(input_key[12:]) - 1] = kwargs[input_key]
|
||||
else:
|
||||
output_values = [None]
|
||||
|
||||
# Run only if "free_memory" was toggled ON
|
||||
if kwargs.get("free_memory"):
|
||||
print("🔷 Control Order & Free Memory • Full VRAM & RAM cleanup 🔷")
|
||||
|
||||
try:
|
||||
# Take snapshot of VRAM before cleanup
|
||||
if torch.cuda.is_available(): # NVIDIA
|
||||
initial_vram = torch.cuda.memory_allocated()
|
||||
elif hasattr(torch, 'mps') and torch.mps.is_available(): # Apple Silicon
|
||||
initial_vram = torch.mps.current_allocated_memory()
|
||||
elif hasattr(torch, 'xpu') and torch.xpu.is_available(): # Intel XPU
|
||||
initial_vram = torch.xpu.memory.memory_allocated()
|
||||
elif hasattr(torch, 'npu') and torch.npu.is_available(): # Ascend NPU
|
||||
initial_vram = torch.npu.memory_allocated()
|
||||
elif hasattr(torch, 'mlu') and torch.mlu.is_available(): # Cambricon MLU
|
||||
initial_vram = torch.mlu.memory_allocated()
|
||||
else: # CPU / fallback
|
||||
initial_vram = None
|
||||
|
||||
# Take snapshot of RAM before cleanup
|
||||
initial_ram = psutil.virtual_memory().used # This is fully cross-platform & included in all ComfyUI installations.
|
||||
|
||||
# If any of the inputs were models, keep them loaded
|
||||
loaded_models = model_management.loaded_models() # returns a list of the actual model objects currently tracked by ComfyUI in current_loaded_models
|
||||
keep_loaded = []
|
||||
for key in input_keys:
|
||||
if kwargs[key] in loaded_models: # if any of the inputs was a model that's currently loaded, keep it loaded.
|
||||
keep_loaded.append(kwargs[key])
|
||||
|
||||
# Unload everything else
|
||||
if keep_loaded:
|
||||
model_management.free_memory(1e30, model_management.get_torch_device(), keep_loaded)
|
||||
# This is the core "free as much VRAM as possible" function, where "1e30" means "free this large amount of memory (i.e. everything)". It walks the ComfyUI internal list "current_loaded_models", skips unloading anything in keep_loaded, calls model_unload() on the rest (which does the model unload/detach, unpatch weights, set real_model=None, and for some models a controlled partially_unload() to the offload device - no forced CPU offload unless the model itself decides to partially unload), pops them from current_loaded_models, then calls cleanup_models_gc() (which conditionally runs gc.collect() + soft_empty_cache() if any momdel in current_loaded_models is_dead() i.e. memory leak is suspected) and soft_empty_cache() once (if it unloaded at least 1 model).
|
||||
# This is the official ComfyUI maintained method that safely manages its internal model load states through current_loaded_models entries and other internal memory accounting.
|
||||
print(" - All models unloaded (except models connected into persist_any)")
|
||||
else:
|
||||
model_management.unload_all_models()
|
||||
# This just calls free_memory(1e30, get_torch_device()) with no keep_loaded models list, defaulting it to [] (i.e. it unloads all models).
|
||||
# This is the official ComfyUI maintained method that safely manages its internal model load states through current_loaded_models entries and other internal memory accounting.
|
||||
print(" - All models unloaded")
|
||||
|
||||
print(" - Synchronizing hardware accelerator")
|
||||
model_management.soft_empty_cache(True)
|
||||
# This is a device-agnostic wrapper that does:
|
||||
# → CUDA: torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.ipc_collect()
|
||||
# → MPS / XPU / NPU / MLU: runs the equivalent empty_cache for that backend, and the force param is ignored in current ComfyUI (legacy).
|
||||
except Exception as e:
|
||||
print(f" - Non-fatal error during unload: {e}")
|
||||
finally:
|
||||
# First pass (Device-agnostic empty_cache)
|
||||
# DESIGN: While model_management.soft_empty_cache(True) is a convenient one-liner for cross-device empty_cache(), I've extracted all the cross-device operations to extend & optimize the cleanup sequence per device, while skipping any extra calls to torch.cuda.synchronize() which is now redundant after model_management.soft_empty_cache(True)
|
||||
print(" - Clearing VRAM")
|
||||
if torch.cuda.is_available(): # NVIDIA
|
||||
torch.cuda.empty_cache() # releases the cached VRAM and available memory held by the allocator but not currently in use
|
||||
elif hasattr(torch, 'mps') and torch.mps.is_available(): # Apple Silicon
|
||||
torch.mps.empty_cache()
|
||||
elif hasattr(torch, 'xpu') and torch.xpu.is_available(): # Intel XPU
|
||||
torch.xpu.empty_cache()
|
||||
elif hasattr(torch, 'npu') and torch.npu.is_available(): # Ascend NPU
|
||||
torch.npu.empty_cache()
|
||||
elif hasattr(torch, 'mlu') and torch.mlu.is_available(): # Cambricon MLU
|
||||
torch.mlu.empty_cache()
|
||||
else: # CPU / fallback
|
||||
print(" - No GPU accelerator detected.")
|
||||
|
||||
print(" - Clearing RAM")
|
||||
gc.collect() # release objects from memory that no longer have active references. Critical for freeing CPU RAM + any Python object references after the tensors are gone.
|
||||
|
||||
# Second pass (catches anything GC just released + IPC on CUDA)
|
||||
if torch.cuda.is_available(): # NVIDIA
|
||||
torch.cuda.empty_cache() # double empty_cache to catch lingering tensors after GC
|
||||
torch.cuda.ipc_collect() # frees any lingering CUDA IPC / shared-memory handles that empty_cache() sometimes misses. Useful when models were loaded with certain GGUF/quantized loaders or in multi-process scenarios. Harmless in a normal single-process ComfyUI server.
|
||||
model_management.cleanup_models() # lightweight & harmless "remove dead/stale model wrappers" helper that scans current_loaded_models and removes entries where real_model() is None (i.e. dead wrappers that free_memory may have left behind in some edge cases), pops them from current_loaded_models, and deletes the wrapper. This is only meaningful after the model tensors have already been unloaded + garbage-collected such that every dead model wrapper whose real_model() just became None can be caught.
|
||||
print(" - Clearing CUDA stats")
|
||||
try:
|
||||
torch.cuda.reset_peak_memory_stats() # Optional stats reset. try catch because this can raise in edge cases (no active CUDA context, older PyTorch, or after certain errors)
|
||||
except:
|
||||
pass
|
||||
final_vram = torch.cuda.memory_allocated() # Take snapshot of VRAM after cleanup
|
||||
elif hasattr(torch, 'mps') and torch.mps.is_available(): # Apple Silicon
|
||||
torch.mps.empty_cache()
|
||||
model_management.cleanup_models()
|
||||
final_vram = torch.mps.current_allocated_memory() # Take snapshot of VRAM after cleanup
|
||||
elif hasattr(torch, 'xpu') and torch.xpu.is_available(): # Intel XPU
|
||||
torch.xpu.empty_cache()
|
||||
model_management.cleanup_models()
|
||||
final_vram = torch.xpu.memory.memory_allocated() # Take snapshot of VRAM after cleanup
|
||||
elif hasattr(torch, 'npu') and torch.npu.is_available(): # Ascend NPU
|
||||
torch.npu.empty_cache()
|
||||
model_management.cleanup_models()
|
||||
final_vram = torch.npu.memory_allocated() # Take snapshot of VRAM after cleanup
|
||||
elif hasattr(torch, 'mlu') and torch.mlu.is_available(): # Cambricon MLU
|
||||
torch.mlu.empty_cache()
|
||||
model_management.cleanup_models()
|
||||
final_vram = torch.mlu.memory_allocated() # Take snapshot of VRAM after cleanup
|
||||
else: # CPU / fallback
|
||||
final_vram = None
|
||||
|
||||
# Take snapshot of RAM after cleanup
|
||||
final_ram = psutil.virtual_memory().used
|
||||
|
||||
# Print memory analytics
|
||||
print(" - VRAM & RAM cleanup complete")
|
||||
if initial_vram is None:
|
||||
print(" - GPU VRAM: No GPU accelerator detected.")
|
||||
else:
|
||||
print(f" - GPU VRAM: Initial usage: {initial_vram/1073741824:.2f} GB, Final usage: {final_vram/1073741824:.2f} GB, Freed: {(initial_vram - final_vram)/1073741824:.2f} GB") # 1073741824 == (1024 ** 3)
|
||||
print(f" - System RAM: Initial usage: {initial_ram/1073741824:.2f} GB, Final usage: {final_ram/1073741824:.2f} GB, Freed: {(initial_ram - final_ram)/1073741824:.2f} GB")
|
||||
|
||||
return tuple(output_values)
|
||||
|
||||
class FileNameSelector:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"string": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "Select a filename and further edit the string here."
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("STRING",)
|
||||
OUTPUT_TOOLTIPS = ("Final string for downstream use", )
|
||||
FUNCTION = "select_filename"
|
||||
CATEGORY = "Control Order & Free Memory"
|
||||
DESCRIPTION = "Select a filename from your native OS file picker and optionally edit the string for downstream use."
|
||||
|
||||
def select_filename(self, string):
|
||||
return (string,)
|
||||
Reference in New Issue
Block a user