128 lines
4.5 KiB
Python
128 lines
4.5 KiB
Python
# ======================================================================
|
|
# Node Pack: artyclaw-comfy
|
|
# Author: artyclaw
|
|
# Version: 1.0.0
|
|
# Description: A power-user toolkit for prompt logic, IO, and RNG.
|
|
# ======================================================================
|
|
|
|
# prompt_block_loader.py
|
|
# ComfyUI Custom Node: Prompt Block Loader
|
|
# Loads two prompt blocks (<name>1 and <name>2) from a config text file.
|
|
# - '#' lines are comments
|
|
# - Blocks are introduced by headers `!<blockname>1` or `!<blockname>2`
|
|
# - Block content runs until the next '!' header or EOF
|
|
# - Automatically reloads when the file changes (via IS_CHANGED)
|
|
# - Optional force_reload toggle to manually bypass cache
|
|
|
|
import os
|
|
import re
|
|
import time
|
|
|
|
HEADER_RE = re.compile(r"^!([A-Za-z0-9_\-]+)([12])\s*$")
|
|
|
|
class PromptBlockLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"config_path": ("STRING", {
|
|
"multiline": False,
|
|
"default": "prompt_blocks.txt",
|
|
"tooltip": "Path to the config .txt file (relative or absolute)."
|
|
}),
|
|
"blockname": ("STRING", {
|
|
"multiline": False,
|
|
"default": "myblock",
|
|
"tooltip": "Base block name to fetch: will output <blockname>1 and <blockname>2."
|
|
}),
|
|
},
|
|
"optional": {
|
|
"force_reload": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Force reload even if file timestamp hasn't changed."
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("block1", "block2")
|
|
FUNCTION = "load_blocks"
|
|
CATEGORY = "Prompting/IO"
|
|
|
|
# --- Comfy cache control ---
|
|
@classmethod
|
|
def IS_CHANGED(cls, config_path, blockname, force_reload=False):
|
|
"""
|
|
ComfyUI calls this to check whether cached outputs are still valid.
|
|
We return the file modification time (and force_reload flag) so it
|
|
re-runs automatically when the file changes or when manually forced.
|
|
"""
|
|
try:
|
|
path = os.path.abspath(config_path)
|
|
mtime = os.path.getmtime(path)
|
|
except Exception:
|
|
# If file missing, return current time to retry next run
|
|
mtime = time.time()
|
|
# Including blockname and force_reload ensures cache invalidation
|
|
return (mtime, blockname, bool(force_reload))
|
|
|
|
# --- helpers ---
|
|
def _read_file(self, path: str) -> list[str]:
|
|
if not os.path.isabs(path):
|
|
path = os.path.abspath(path)
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"Config file not found: {path}")
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return [line.rstrip("\r\n") for line in f.readlines()]
|
|
|
|
def _parse_blocks(self, lines: list[str]) -> dict[tuple[str, str], str]:
|
|
"""
|
|
Parse the file into a dict with keys (blockname, '1' or '2') -> content string.
|
|
Comment lines starting with '#' are ignored.
|
|
"""
|
|
blocks: dict[tuple[str, str], list[str]] = {}
|
|
current_key: tuple[str, str] | None = None
|
|
|
|
for raw in lines:
|
|
if not raw:
|
|
if current_key is not None:
|
|
blocks.setdefault(current_key, []).append("")
|
|
continue
|
|
if raw.startswith("#"):
|
|
continue
|
|
|
|
m = HEADER_RE.match(raw)
|
|
if m:
|
|
name, idx = m.group(1), m.group(2)
|
|
current_key = (name, idx)
|
|
blocks.setdefault(current_key, [])
|
|
continue
|
|
|
|
if current_key is not None:
|
|
blocks[current_key].append(raw)
|
|
|
|
return {k: "\n".join(v).strip("\n") for k, v in blocks.items()}
|
|
|
|
# --- main ---
|
|
def load_blocks(self, config_path: str, blockname: str, force_reload=False):
|
|
try:
|
|
lines = self._read_file(config_path)
|
|
parsed = self._parse_blocks(lines)
|
|
except Exception as e:
|
|
print(f"[PromptBlockLoader] Error: {e}")
|
|
return ("", "")
|
|
|
|
b1 = parsed.get((blockname, "1"), "")
|
|
b2 = parsed.get((blockname, "2"), "")
|
|
return (b1, b2)
|
|
|
|
|
|
# --- registration ---
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PromptBlockLoader": PromptBlockLoader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PromptBlockLoader": "Prompt Block Loader",
|
|
}
|