Optimize custom nodes, fix memory leaks, and add V2 settings compatibility

This commit is contained in:
Satadal Dhara
2026-06-18 07:49:44 +05:30
parent 8417902269
commit 7eb8a036c7
8 changed files with 136 additions and 58 deletions
+17 -4
View File
@@ -8,11 +8,22 @@ const SAVE_MACHINE_CLASS = "Save_Machine";
// Track hover state
const hoveredNodeIds = new Set();
// Helper to get settings for both ComfyUI V1 and V2
function getSettingValue(id, defaultValue) {
if (app.extensionManager && app.extensionManager.setting) {
const val = app.extensionManager.setting.get(id);
return val !== undefined ? val : defaultValue;
} else if (app.ui && app.ui.settings) {
return app.ui.settings.getSettingValue(id, defaultValue);
}
return defaultValue;
}
app.registerExtension({
name: "SATA_UtilityNode.Preview_Machine",
async setup() {
app.ui.settings.addSetting({
settings: [
{
id: GLOBAL_HIDE_SETTING_ID,
name: "SATA Utility: Global Hide Previews",
type: "boolean",
@@ -20,8 +31,10 @@ app.registerExtension({
onChange: (value) => {
if (app.graph) app.graph.setDirtyCanvas(true, true);
}
});
}
],
async setup() {
setupGlobalMouseTracker();
},
@@ -46,7 +59,7 @@ function isTargetNode(node) {
}
function shouldHideNode(node) {
const globalHide = app.ui.settings.getSettingValue(GLOBAL_HIDE_SETTING_ID, false);
const globalHide = getSettingValue(GLOBAL_HIDE_SETTING_ID, false);
if (node.comfyClass === PREVIEW_MACHINE_CLASS) {
return true;
+29 -12
View File
@@ -1,30 +1,38 @@
import { app } from "../../../scripts/app.js";
// Helper to get settings for both ComfyUI V1 and V2
function getSettingValue(id, defaultValue) {
if (app.extensionManager && app.extensionManager.setting) {
const val = app.extensionManager.setting.get(id);
return val !== undefined ? val : defaultValue;
} else if (app.ui && app.ui.settings) {
return app.ui.settings.getSettingValue(id, defaultValue);
}
return defaultValue;
}
// Register extension
app.registerExtension({
name: "SATA_UtilityNode.PromptAutocomplete",
async setup() {
// Add settings
app.ui.settings.addSetting({
settings: [
{
id: "SATA_UtilityNode.PromptAutocomplete.Trigger",
name: "Prompt Autocomplete Trigger Character",
type: "text",
defaultValue: "#",
});
app.ui.settings.addSetting({
},
{
id: "SATA_UtilityNode.PromptAutocomplete.Global",
name: "Prompt Autocomplete Global Mode (All Text Widgets)",
type: "boolean",
defaultValue: false,
});
},
}
],
async nodeCreated(node) {
// Check if we should attach to this node
async initAutocompleteForNode(node) {
const isTargetNode = node.comfyClass === "PromptAutocomplete";
const isGlobal = app.ui.settings.getSettingValue("SATA_UtilityNode.PromptAutocomplete.Global", false);
const isGlobal = getSettingValue("SATA_UtilityNode.PromptAutocomplete.Global", false);
if (!isTargetNode && !isGlobal) return;
@@ -42,6 +50,14 @@ app.registerExtension({
});
},
async nodeCreated(node) {
this.initAutocompleteForNode(node);
},
async loadedGraphNode(node) {
this.initAutocompleteForNode(node);
},
snippetsCache: null,
async loadSnippets() {
@@ -137,7 +153,7 @@ app.registerExtension({
triggerChar: "#",
handleInput(e, input, widget) {
const trigger = app.ui.settings.getSettingValue("SATA_UtilityNode.PromptAutocomplete.Trigger", "#");
const trigger = getSettingValue("SATA_UtilityNode.PromptAutocomplete.Trigger", "#");
this.triggerChar = trigger;
const val = input.value;
@@ -403,6 +419,7 @@ app.registerExtension({
},
selectItem(input) {
if (!input) return;
const item = this.filteredItems[this.selectedIndex];
if (!item) return;
+2 -2
View File
@@ -81,8 +81,8 @@ app.registerExtension({
if (resolutionWidget.value === "Custom") {
widthWidget.type = "INT";
heightWidget.type = "INT";
widthWidget.computeSize = LGraphCanvas.prototype.computeSize;
heightWidget.computeSize = LGraphCanvas.prototype.computeSize;
delete widthWidget.computeSize;
delete heightWidget.computeSize;
widthWidget.disabled = false;
heightWidget.disabled = false;
} else {
+4 -1
View File
@@ -97,7 +97,10 @@ LGraphCanvas.prototype.processMouseWheel = function (/** @type {WheelEvent} */ e
if (!this.graph || !this.allow_dragcanvas) return;
const { clientX: x, clientY: y, deltaX, deltaY, ctrlKey, metaKey } = event;
const rect = this.canvas.getBoundingClientRect();
const x = event.clientX - rect.left;
const y = event.clientY - rect.top;
const { deltaX, deltaY, ctrlKey, metaKey } = event;
// Check if the event occurred inside the canvas viewport
if (this.viewport) {
+2 -1
View File
@@ -10,7 +10,7 @@ class Latent_Machine:
"""
def __init__(self):
self.device = comfy.model_management.get_torch_device()
pass
@classmethod
def INPUT_TYPES(s):
@@ -48,6 +48,7 @@ class Latent_Machine:
CATEGORY = "SATA_UtilityNode"
def generate_noise(self, width, height, batch_size, model_type, noise_type, intensity, seed):
self.device = comfy.model_management.get_torch_device()
# Set seed for reproducibility
torch.manual_seed(seed)
random.seed(seed)
+1
View File
@@ -10,6 +10,7 @@ PROMPT_DIR = os.path.join(os.path.dirname(__file__), "..", "prompt")
def list_prompt_files():
try:
if not os.path.exists(PROMPT_DIR):
os.makedirs(PROMPT_DIR, exist_ok=True)
return []
return sorted([f for f in os.listdir(PROMPT_DIR) if f.endswith(".csv") or f.endswith(".json")])
except Exception as e:
+50 -10
View File
@@ -41,7 +41,7 @@ def make_filename(template: str):
return get_timestamp("%Y-%m-%d-%H%M%S") if filename == "" else filename
def resolve_placeholders(template: str, prompt=None) -> str:
def resolve_placeholders(template: str, prompt=None, extra_pnginfo=None) -> str:
"""Resolve placeholders of the form %node_name/field% by looking into extra_pnginfo and prompt.
This is a best-effort resolver: it attempts several lookup strategies and falls back to leaving
@@ -50,18 +50,58 @@ def resolve_placeholders(template: str, prompt=None) -> str:
if not template or '%' not in template:
return template
# Build title/type to node ID mapping from extra_pnginfo
node_title_to_id = {}
node_type_to_ids = {}
if extra_pnginfo and isinstance(extra_pnginfo, dict) and "workflow" in extra_pnginfo:
workflow = extra_pnginfo["workflow"]
if isinstance(workflow, dict) and "nodes" in workflow:
for node in workflow["nodes"]:
if isinstance(node, dict):
node_id = str(node.get("id"))
node_title = node.get("title")
node_type = node.get("type")
if node_title:
node_title_to_id[node_title] = node_id
if node_type:
if node_type not in node_type_to_ids:
node_type_to_ids[node_type] = []
node_type_to_ids[node_type].append(node_id)
def lookup(key: str):
# key is like 'node_name/field' or 'field'
if not key:
return None
# Try prompt mapping
node_ref = None
field = key
if '/' in key:
node_ref, field = key.split('/', 1)
node_id = None
if node_ref:
if prompt and node_ref in prompt:
node_id = node_ref
elif node_ref in node_title_to_id:
node_id = node_title_to_id[node_ref]
elif node_ref in node_type_to_ids:
node_id = node_type_to_ids[node_ref][0]
if node_id and prompt and node_id in prompt:
node_data = prompt[node_id]
if isinstance(node_data, dict) and "inputs" in node_data:
inputs = node_data["inputs"]
if field in inputs:
return inputs[field]
# Fallback: check if the key is directly in prompt or if it is just a field name in some node
if prompt and isinstance(prompt, dict):
if key in prompt:
return prompt[key]
if '/' in key:
_, field = key.split('/', 1)
if field in prompt:
return prompt[field]
for nid, ndata in prompt.items():
if isinstance(ndata, dict) and "inputs" in ndata:
inputs = ndata.get("inputs", {})
if field in inputs:
return inputs[field]
return None
@@ -117,7 +157,7 @@ class Save_Machine:
quality_jpeg_or_webp=100, lossless_webp=True, hide_preview=True, prompt=None, extra_pnginfo=None):
# Resolve placeholders first (uses prompt)
try:
resolved = resolve_placeholders(str(path_and_filename), prompt)
resolved = resolve_placeholders(str(path_and_filename), prompt, extra_pnginfo)
except Exception:
resolved = str(path_and_filename)
@@ -153,8 +193,8 @@ class Save_Machine:
# Resolve any %node/field% placeholders in filename and path using prompt
try:
filename = resolve_placeholders(filename, prompt)
path = resolve_placeholders(path, prompt)
filename = resolve_placeholders(filename, prompt, extra_pnginfo)
path = resolve_placeholders(path, prompt, extra_pnginfo)
except Exception:
# Best-effort: if resolution fails, continue with original templates
pass
+31 -28
View File
@@ -92,37 +92,40 @@ class Upscale_Machine:
device = model_management.get_torch_device()
upscale_model.to(device)
in_img = image_bhwc.movedim(-1, -3).to(device)
try:
in_img = image_bhwc.movedim(-1, -3).to(device)
tile = 512
overlap = 32
oom = True
while oom:
try:
if not pbar:
steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(
in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap
tile = 512
overlap = 32
oom = True
while oom:
try:
if not pbar:
steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(
in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap
)
local_pbar = comfy.utils.ProgressBar(steps)
else:
local_pbar = pbar
s = comfy.utils.tiled_scale(
in_img,
lambda a: upscale_model(a),
tile_x=tile,
tile_y=tile,
overlap=overlap,
upscale_amount=getattr(upscale_model, "scale", 4),
pbar=local_pbar
)
local_pbar = comfy.utils.ProgressBar(steps)
else:
local_pbar = pbar
oom = False
except model_management.OOM_EXCEPTION as e:
model_management.soft_empty_cache()
tile //= 2
if tile < 128:
raise e
finally:
upscale_model.cpu()
s = comfy.utils.tiled_scale(
in_img,
lambda a: upscale_model(a),
tile_x=tile,
tile_y=tile,
overlap=overlap,
upscale_amount=getattr(upscale_model, "scale", 4),
pbar=local_pbar
)
oom = False
except model_management.OOM_EXCEPTION as e:
tile //= 2
if tile < 128:
raise e
upscale_model.cpu()
s = torch.clamp(s, min=0.0, max=1.0)
return s