Adam
591aa08da4
- True runs remember captured text per (tab path, node) with a token. Off runs whose widget hadn't received it use the remembered text and update the embedded metadata. IS_CHANGED includes the carried token, so prompts queued ahead pick it up - Scope pushed prompt updates to the tab the run was queued from, fixing captured text overwriting a same-id node in another open workflow - Nodes fetch missed text when a workflow loads or its tab is re-opened - In-memory only; tokens are compared by equality, so restarts are safe - Saver and Passthrough get hidden stash_tab/stash_token inputs (appended last, no socket or widget, not saved in widgets_values)
253 lines
11 KiB
Python
253 lines
11 KiB
Python
import os
|
|
import json
|
|
from server import PromptServer
|
|
import hashlib
|
|
from .data_utils import init_data_file, update_node_in_workflow, prompt_has_input
|
|
from . import stash_sync
|
|
|
|
class PromptStashSaver:
|
|
def __init__(self):
|
|
self.base_dir = os.path.dirname(os.path.realpath(__file__))
|
|
self.data_file = init_data_file(self.base_dir)
|
|
self.data = self.load_data()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
},
|
|
"optional": {
|
|
"use_input_text": ("BOOLEAN", {"default": False, "label_on": "Use Input", "label_off": "Use Prompt"}),
|
|
"text": ("STRING", {"default": "", "forceInput": True, "tooltip": "Optional input text", "lazy": True}),
|
|
"prompt_text": ("STRING", {"multiline": True, "default": "", "placeholder": "Enter prompt text"}),
|
|
"save_as_key": ("STRING", {"default": "", "placeholder": "Enter key to save as"}),
|
|
# COMBO inputs declared so the frontend creates co-existing input
|
|
# sockets (connectable by Primitive nodes). The actual option lists
|
|
# are populated live on the client from server events; these defaults
|
|
# are only the initial values. Appended last to preserve the existing
|
|
# widget order so widgets_values indices and saved workflows still map.
|
|
"load_saved": (["None"], {"default": "None"}),
|
|
"prompt_lists": (["default"], {"default": "default"}),
|
|
# Filled in by the frontend at queue time (hidden, no socket);
|
|
# used to carry input text forward to later runs. See stash_sync.py.
|
|
"stash_tab": ("STRING", {"default": "", "forceInput": True}),
|
|
"stash_token": ("STRING", {"default": "", "forceInput": True}),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"prompt": "PROMPT",
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("text", "meta")
|
|
OUTPUT_TOOLTIPS = (
|
|
"The prompt text",
|
|
"JSON string with info about the prompt: prompt_name, list_name, source, etc. "
|
|
"Feed it to a JSON-aware node to extract fields (e.g. prompt_name for a filename prefix).",
|
|
)
|
|
FUNCTION = "process"
|
|
CATEGORY = "utils"
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, load_saved="None", prompt_lists="default"):
|
|
# load_saved / prompt_lists are UI-only state used to load saved data
|
|
# outside the run phase; they do not affect process() output. Their live
|
|
# option lists exist only on the client (populated from server events),
|
|
# so the backend's static COMBO validation (against the ["None"] /
|
|
# ["default"] stubs in INPUT_TYPES) would wrongly reject any real
|
|
# selection. Listing these inputs as params makes ComfyUI skip combo
|
|
# validation for them (execution.py validate_inputs, line ~997).
|
|
return True
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, use_input_text=False, text="", prompt_text="", save_as_key="", load_saved="None", prompt_lists="default", stash_tab="", stash_token="", unique_id=None, extra_pnginfo=None, prompt=None):
|
|
m = hashlib.sha256()
|
|
|
|
# Always include these parameters as they affect the output
|
|
m.update(str(use_input_text).encode())
|
|
m.update(str(prompt_text).encode())
|
|
m.update(str(unique_id).encode())
|
|
|
|
# Text carried forward from a True run the frontend hadn't taken in when
|
|
# this prompt was queued. IS_CHANGED runs when the prompt starts
|
|
# executing (not when queued), so prompts queued ahead see it.
|
|
carried = stash_sync.carried_entry(stash_tab, unique_id, stash_token)
|
|
m.update(str(carried["token"] if carried else "").encode())
|
|
|
|
# These feed the "meta" output (prompt_name / list_name)
|
|
m.update(str(save_as_key).encode())
|
|
m.update(str(load_saved).encode())
|
|
m.update(str(prompt_lists).encode())
|
|
|
|
# Only include the text input if use_input_text is True
|
|
if use_input_text and text is not None:
|
|
m.update(str(text).encode())
|
|
|
|
return m.hexdigest()
|
|
|
|
def check_lazy_status(self, use_input_text=False, text="", prompt_text="", save_as_key="", load_saved="None", prompt_lists="default", stash_tab="", stash_token="", unique_id=None, extra_pnginfo=None, prompt=None):
|
|
# Only need the text input if use_input_text is True.
|
|
needed = []
|
|
if use_input_text:
|
|
if not prompt_has_input(prompt, unique_id, "text"):
|
|
raise ValueError(
|
|
f"Prompt Stash Saver (node {unique_id}): 'use_input_text' is on, but nothing "
|
|
"is connected to the 'text' input. Connect a text input or switch to 'Use Prompt'."
|
|
)
|
|
needed.append("text")
|
|
return needed
|
|
|
|
def load_data(self):
|
|
if os.path.exists(self.data_file):
|
|
try:
|
|
with open(self.data_file, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
# Broadcast initial data to all nodes
|
|
PromptServer.instance.send_sync("prompt-stash-update-all", {
|
|
"lists": data.get("lists", {"default": {}})
|
|
})
|
|
return data
|
|
except Exception as e:
|
|
print(f"Error loading prompts: {e}")
|
|
return {"lists": {"default": {}}}
|
|
|
|
def save_data(self):
|
|
try:
|
|
os.makedirs(os.path.dirname(self.data_file), exist_ok=True)
|
|
with open(self.data_file, 'w', encoding='utf-8') as f:
|
|
json.dump(self.data, f, indent=2)
|
|
return True
|
|
except Exception as e:
|
|
print(f"Error saving data: {e}")
|
|
return False
|
|
|
|
def save_prompt(self, save_as_key, prompt, list_name, unique_id):
|
|
save_as_key = save_as_key.strip()
|
|
if not save_as_key or not prompt:
|
|
return False
|
|
|
|
if list_name not in self.data["lists"]:
|
|
list_name = "default"
|
|
|
|
self.data["lists"][list_name][save_as_key] = prompt
|
|
success = self.save_data()
|
|
|
|
if success:
|
|
# Notify all nodes of the update
|
|
PromptServer.instance.send_sync("prompt-stash-update-all", {
|
|
"lists": self.data["lists"]
|
|
})
|
|
return success
|
|
|
|
def delete_prompt(self, save_as_key, list_name, unique_id):
|
|
if list_name not in self.data["lists"]:
|
|
return False
|
|
|
|
if save_as_key in self.data["lists"][list_name]:
|
|
del self.data["lists"][list_name][save_as_key]
|
|
success = self.save_data()
|
|
|
|
if success:
|
|
# Notify all nodes of the update
|
|
PromptServer.instance.send_sync("prompt-stash-update-all", {
|
|
"lists": self.data["lists"]
|
|
})
|
|
return success
|
|
return False
|
|
|
|
@staticmethod
|
|
def build_meta(use_input_text, text, save_as_key, load_saved, prompt_lists, from_input=None):
|
|
"""Build the JSON string for the "meta" output.
|
|
|
|
prompt_name precedence:
|
|
1. "llm input" - text came in via the `text` input, this run or carried
|
|
forward from an earlier run (`from_input`); any
|
|
save_as_key / load_saved on the node is stale
|
|
2. save_as_key - the user typed/loaded a name into "Save Name"
|
|
3. load_saved - a saved prompt is selected but Save Name is empty
|
|
4. "" - manual, unnamed text
|
|
"""
|
|
key = (save_as_key or "").strip()
|
|
loaded = load_saved if load_saved and load_saved != "None" else ""
|
|
|
|
if from_input is None:
|
|
from_input = bool(use_input_text and text is not None)
|
|
|
|
if from_input:
|
|
prompt_name = "llm input"
|
|
source = "input"
|
|
elif key:
|
|
prompt_name = key
|
|
source = "saved" if loaded else "manual"
|
|
elif loaded:
|
|
prompt_name = loaded
|
|
source = "saved"
|
|
else:
|
|
prompt_name = ""
|
|
source = "manual"
|
|
|
|
meta = {
|
|
"prompt_name": prompt_name,
|
|
"list_name": prompt_lists,
|
|
"source": source, # "input" | "saved" | "manual"
|
|
"save_as_key": key,
|
|
"load_saved": loaded,
|
|
"use_input_text": bool(use_input_text),
|
|
}
|
|
return json.dumps(meta, ensure_ascii=False, indent=2)
|
|
|
|
def process(self, use_input_text=False, text="", prompt_text="", save_as_key="", load_saved="None", prompt_lists="default", stash_tab="", stash_token="", unique_id=None, extra_pnginfo=None, prompt=None):
|
|
from_input = bool(use_input_text and text is not None)
|
|
|
|
# Update the prompt text based on use_input_text toggle
|
|
output_text = prompt_text
|
|
if from_input:
|
|
output_text = text
|
|
# Remember it so later runs use it even if the frontend misses the
|
|
# push below, then tell the frontend (only the matching tab applies it).
|
|
token = stash_sync.store_text(stash_tab, unique_id, text)
|
|
stash_sync.push_update(unique_id, text, stash_tab, token)
|
|
else:
|
|
# Off run: if a True run produced text the frontend hadn't taken in
|
|
# when this prompt was queued, use that instead of the stale widget.
|
|
carried = stash_sync.carried_entry(stash_tab, unique_id, stash_token)
|
|
if carried:
|
|
from_input = True
|
|
output_text = carried["text"]
|
|
stash_sync.push_update(unique_id, output_text, stash_tab, carried["token"])
|
|
|
|
# Text that came from the input (this run or carried) counts as "llm input".
|
|
meta = self.build_meta(use_input_text, text, save_as_key, load_saved, prompt_lists, from_input=from_input)
|
|
|
|
if from_input:
|
|
# Handle both list and dict formats of extra_pnginfo
|
|
workflow = None
|
|
if isinstance(extra_pnginfo, list) and len(extra_pnginfo) > 0:
|
|
workflow = extra_pnginfo[0].get("workflow")
|
|
elif isinstance(extra_pnginfo, dict):
|
|
workflow = extra_pnginfo.get("workflow")
|
|
|
|
if workflow:
|
|
|
|
def apply_stash_changes(node):
|
|
if "widgets_values" in node:
|
|
# Set use_input_text to False in metadata (index 0 based on INPUT_TYPES order)
|
|
use_input_text_index = 0 # First widget in optional inputs
|
|
prompt_text_index = 1 # Second non-forceInput widget in optional inputs
|
|
|
|
# Safety check, make sure there are at least 2 elements
|
|
if len(node["widgets_values"]) > prompt_text_index:
|
|
node["widgets_values"][use_input_text_index] = False # Force use_input_text to False in metadata
|
|
node["widgets_values"][prompt_text_index] = output_text # Update the prompt text
|
|
|
|
update_node_in_workflow(workflow, unique_id, apply_stash_changes)
|
|
|
|
if prompt and unique_id is not None:
|
|
node_id_str = str(unique_id)
|
|
if node_id_str in prompt:
|
|
prompt[node_id_str]['inputs']['use_input_text'] = False
|
|
prompt[node_id_str]['inputs']['prompt_text'] = output_text
|
|
|
|
return (output_text, meta) |