From ccd091be823cd913214a8ae6cbd337eeaf050de0 Mon Sep 17 00:00:00 2001 From: malkuthro <166633396+malkuthro@users.noreply.github.com> Date: Sat, 4 Oct 2025 22:12:40 +0200 Subject: [PATCH] adding buttons - A --- .gitignore | 1 + __init__.py | 5 ++++- k_ai_pipeline.py | 44 ++++++++++++++++++++++++++++++++---- web/ai_pipeline.js | 56 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 101 insertions(+), 5 deletions(-) create mode 100644 .gitignore create mode 100644 web/ai_pipeline.js diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1f6e59c --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +_safe* \ No newline at end of file diff --git a/__init__.py b/__init__.py index b73f135..82dda01 100644 --- a/__init__.py +++ b/__init__.py @@ -32,4 +32,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { **wan_display, **resize_display, **pipeline_display -} \ No newline at end of file +} + +WEB_DIRECTORY = "./web" +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] \ No newline at end of file diff --git a/k_ai_pipeline.py b/k_ai_pipeline.py index 146793d..7412aee 100644 --- a/k_ai_pipeline.py +++ b/k_ai_pipeline.py @@ -51,20 +51,39 @@ class EasyAIPipeline: "step": 1, "display": "number" }), - } + "output_directory_display": ("STRING", { + "default": "", + "multiline": True + }), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + "extra_pnginfo": "EXTRA_PNGINFO", + }, } + INPUT_IS_LIST = True + # To customize output names, edit the strings in this tuple (e.g., "output_directory" to "output_xxxxx"). # To change the order from top to bottom in the UI, reorder the elements here. # Ensure RETURN_TYPES and the return statement in generate_pipeline match the order. RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "INT", "INT", "STRING") - RETURN_NAMES = ("OUTPUT - Path+File", "output_name", "version_string", "output_directory", "shot_duration", "seed_value", "shot_name") + RETURN_NAMES = ("filePath_OUTPUT", "output_name", "version_string", "output_directory", "shot_duration", "seed_value", "shot_name") FUNCTION = "generate_pipeline" CATEGORY = "Koolook/VFX" OUTPUT_NODE = True # Marks it as an output node for workflow integration - def generate_pipeline(self, shot_duration, seed_value, job_path, extension, shot_name, ai_method, version): + def generate_pipeline(self, shot_duration, seed_value, job_path, extension, shot_name, ai_method, version, output_directory_display, unique_id=None, extra_pnginfo=None): + shot_duration = shot_duration[0] + seed_value = seed_value[0] + job_path = job_path[0] + extension = extension[0] + shot_name = shot_name[0] + ai_method = ai_method[0] + version = version[0] + output_directory_display = output_directory_display[0] + # Generate version string like 'v001' version_str = f"v{version:03d}" @@ -78,8 +97,25 @@ class EasyAIPipeline: # Final file path: output_directory/name filePath_OUTPUT = os.path.join(output_directory, output_name).replace("\\", "/") + # Update node display if possible + if unique_id is not None and extra_pnginfo is not None: + if isinstance(extra_pnginfo, dict) and "workflow" in extra_pnginfo: + workflow = extra_pnginfo["workflow"] + node = next( + (x for x in workflow["nodes"] if str(x["id"]) == str(unique_id)), + None, + ) + if node and "widgets_values" in node: + # Update the last widget (output_directory_display) + node["widgets_values"][-1] = output_directory + else: + print("Error: extra_pnginfo is not a dict or missing 'workflow' key") + # Return all for chaining in workflows (e.g., connect to savers or prompts) - return (filePath_OUTPUT, output_name, version_str, output_directory, shot_duration, seed_value, shot_name) + return { + "ui": {"text": [output_directory]}, + "result": (filePath_OUTPUT, output_name, version_str, output_directory, shot_duration, seed_value, shot_name) + } # Individual node mappings NODE_CLASS_MAPPINGS = { diff --git a/web/ai_pipeline.js b/web/ai_pipeline.js new file mode 100644 index 0000000..d0a05d2 --- /dev/null +++ b/web/ai_pipeline.js @@ -0,0 +1,56 @@ +import { app } from "../../../scripts/app.js"; + +app.registerExtension({ + name: "koolook.ai_pipeline", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name === "EasyAIPipeline") { + const originalOnNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function() { + if (originalOnNodeCreated) { + originalOnNodeCreated.call(this); + } + + // Add multiline text widget for display (read-only, not serialized) + const displayWidget = this.addWidget( + "text", + "Output Directory Display", + "", + () => {}, + { multiline: true, serialize: false } + ); + displayWidget.inputEl.readOnly = true; + displayWidget.inputEl.style.height = "100px"; // Adjust for visibility + + // Add button to compute and display output_directory + this.addWidget( + "button", + "Update Display", + null, + () => { + // Get current input values from widgets + const job_path = this.widgets.find(w => w.name === "job_path")?.value || ""; + const shot_name = this.widgets.find(w => w.name === "shot_name")?.value || ""; + const ai_method = this.widgets.find(w => w.name === "ai_method")?.value || ""; + const version = this.widgets.find(w => w.name === "version")?.value || 1; + + // Compute version_str + const version_str = `v${version.toString().padStart(3, '0')}`; + + // Compute output_directory (mirror Python logic) + let output_directory = [job_path, shot_name, ai_method, version_str] + .filter(part => part.trim() !== "") + .join("/") + .replace(/\\/g, "/"); + + // Update display + displayWidget.value = output_directory; + + // Redraw canvas + this.setDirtyCanvas(true, true); + }, + { serialize: false } + ); + }; + } + } +}); \ No newline at end of file