diff --git a/js/preview_machine.js b/js/preview_machine.js index 38b1e17..778a1ad 100644 --- a/js/preview_machine.js +++ b/js/preview_machine.js @@ -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; diff --git a/js/prompt_autocomplete.js b/js/prompt_autocomplete.js index ccc93c9..e99bef9 100644 --- a/js/prompt_autocomplete.js +++ b/js/prompt_autocomplete.js @@ -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; diff --git a/js/resolution_machine.js b/js/resolution_machine.js index bbff920..74c9bfa 100644 --- a/js/resolution_machine.js +++ b/js/resolution_machine.js @@ -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 { diff --git a/js/touchpad_pan_zoom.js b/js/touchpad_pan_zoom.js index 18815cb..b788b6a 100644 --- a/js/touchpad_pan_zoom.js +++ b/js/touchpad_pan_zoom.js @@ -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) { diff --git a/nodes/latent_machine.py b/nodes/latent_machine.py index ef08f10..8e50305 100644 --- a/nodes/latent_machine.py +++ b/nodes/latent_machine.py @@ -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) diff --git a/nodes/prompt_autocomplete.py b/nodes/prompt_autocomplete.py index 661a134..5f32986 100644 --- a/nodes/prompt_autocomplete.py +++ b/nodes/prompt_autocomplete.py @@ -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: diff --git a/nodes/save_machine.py b/nodes/save_machine.py index 4b31952..264b9a9 100644 --- a/nodes/save_machine.py +++ b/nodes/save_machine.py @@ -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 diff --git a/nodes/upscale_machine.py b/nodes/upscale_machine.py index fbe376e..07bf417 100644 --- a/nodes/upscale_machine.py +++ b/nodes/upscale_machine.py @@ -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