Optimize custom nodes, fix memory leaks, and add V2 settings compatibility
This commit is contained in:
+17
-4
@@ -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
@@ -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;
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user