263 lines
8.1 KiB
Python
263 lines
8.1 KiB
Python
import logging
|
|
from os import environ
|
|
import random
|
|
import re
|
|
from pathlib import Path
|
|
import json
|
|
import mmap
|
|
|
|
from .parse import parse
|
|
|
|
# from .parse import parse
|
|
|
|
log = logging.getLogger("comfyui-misc-utils")
|
|
log.setLevel(logging.INFO)
|
|
|
|
CLASS_NAME = "MUSimpleWildcard"
|
|
|
|
|
|
def get_lora_tags(lora_name):
|
|
filenames = [Path(f) for f in folder_paths.get_filename_list("loras")]
|
|
found = None
|
|
for f in filenames:
|
|
if f.stem == lora_name:
|
|
found = folder_paths.get_full_path("loras", str(f))
|
|
break
|
|
if not found:
|
|
return []
|
|
meta = load_lora_meta(found)
|
|
tags = json.loads(meta.get("ss_tag_frequency", "{}"))
|
|
alltags = []
|
|
tagcounts = {}
|
|
for k in tags:
|
|
for tag, count in tags[k].items():
|
|
t = tag.strip().replace("_", " ")
|
|
tagcounts[t] = tagcounts.get(t, 0) + count
|
|
alltags = [(v, k) for k, v in tagcounts.items()]
|
|
return sorted(alltags, reverse=True)
|
|
|
|
|
|
def load_lora_meta(filename):
|
|
try:
|
|
with open(filename, "r", encoding="utf8") as f:
|
|
with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as m:
|
|
header = m.read(8)
|
|
n = int.from_bytes(header, "little")
|
|
metadata_bytes = m.read(n)
|
|
return json.loads(metadata_bytes).get("__metadata__", {})
|
|
except Exception as e:
|
|
log.error("Metadata load failed for %s: %s", filename, e)
|
|
return {}
|
|
|
|
|
|
def replace_lora_tags(text, seed):
|
|
rand = random.Random(seed)
|
|
func = re.compile(r"\bTAG<(?P<args>.*?)>")
|
|
matches, text = find_and_remove(func, text, placeholder="WILDCARDLORA")
|
|
for placeholder, value in matches.items():
|
|
args = value["args"].strip().split(",")
|
|
replacement = []
|
|
count = 1
|
|
mode = "t"
|
|
try:
|
|
loraname = args[0]
|
|
if len(args) > 1:
|
|
count = int(args[1])
|
|
if len(args) > 2:
|
|
mode = args[2].strip()
|
|
tags = [t[1] for t in get_lora_tags(loraname)]
|
|
if len(tags) <= count:
|
|
replacement = tags
|
|
elif mode == "r":
|
|
# rand
|
|
replacement = rand.sample(population=tags, k=count)
|
|
else:
|
|
replacement = tags[:count]
|
|
except Exception as e:
|
|
print("Error selecting tags:", e)
|
|
pass
|
|
|
|
r = ", ".join(replacement)
|
|
text = text.replace(placeholder, r)
|
|
|
|
return text
|
|
|
|
|
|
global_ctx = None
|
|
|
|
|
|
def wildcard_prompt_handler(json_data):
|
|
for node_id in json_data["prompt"].keys():
|
|
if json_data["prompt"][node_id]["class_type"] == CLASS_NAME:
|
|
handle_wildcard_node(json_data, node_id)
|
|
return json_data
|
|
|
|
|
|
def handle_wildcard_node(json_data, node_id):
|
|
global global_ctx
|
|
wildcard_info = json_data.get("extra_data", {}).get("extra_pnginfo", {}).get(CLASS_NAME, {})
|
|
n = json_data["prompt"][node_id]
|
|
seed = n["inputs"]["seed"]
|
|
if not (n["inputs"].get("use_pnginfo") and node_id in wildcard_info):
|
|
RAND.seed(seed)
|
|
text, delayed = replace_wildcards(n["inputs"]["text"])
|
|
if not global_ctx or "$debugwc" in text:
|
|
global_ctx = read_preamble()
|
|
global_ctx.set("seed", lambda: str(seed))
|
|
text, _ = parse(text, global_ctx)
|
|
text, _ = replace_wildcards(text, delayed)
|
|
text = replace_lora_tags(text, seed)
|
|
|
|
if text.strip() != n["inputs"]["text"].strip():
|
|
json_data["prompt"][node_id]["inputs"]["use_pnginfo"] = True
|
|
wildcard_info[node_id] = text
|
|
json_data["extra_data"]["extra_pnginfo"][CLASS_NAME] = wildcard_info
|
|
return json_data
|
|
|
|
|
|
def read_preamble():
|
|
includes = environ.get("MU_WILDCARD_INCLUDE", "defaults")
|
|
includes = [x.strip() for x in includes.split(";") if x.strip()]
|
|
ctx = None
|
|
for x in includes:
|
|
file = Path(x)
|
|
if x == "defaults":
|
|
curfile = Path(__file__)
|
|
file = curfile.parent / "default_functions.txt"
|
|
try:
|
|
with open(file, "r") as f:
|
|
log.info("Reading functions from %s", file)
|
|
ctx = parse(f.read(), ctx=ctx)[1]
|
|
except Exception as e:
|
|
log.error("Error reading definitions from %s: %s. Ignoring...", file, e)
|
|
return ctx
|
|
|
|
|
|
def find_and_remove(regexp, text, placeholder=""):
|
|
m = regexp.search(text)
|
|
res = {}
|
|
index = 0
|
|
|
|
while m:
|
|
index += 1
|
|
ph = ""
|
|
if placeholder:
|
|
ph = "__" + placeholder + str(index) + "__"
|
|
res[ph] = m.groupdict()
|
|
else:
|
|
res[m["name"]] = m.groupdict()
|
|
s, e = m.span()
|
|
text = text[:s] + ph + text[e:]
|
|
m = regexp.search(text)
|
|
return res, text
|
|
|
|
|
|
RAND = random.Random()
|
|
|
|
|
|
def read_wildcards(name):
|
|
path = environ.get("MU_WILDCARD_BASEDIR", "wildcards")
|
|
r = name.split(":")
|
|
name = r[0]
|
|
filters = r[1:]
|
|
|
|
def matches(x):
|
|
for f in filters:
|
|
if f.startswith("!") and f[1:] in x:
|
|
return False
|
|
elif not f.startswith("!") and f not in x:
|
|
return False
|
|
return True
|
|
|
|
if name == "LORA":
|
|
return [l for l in folder_paths.get_filename_list("loras") if matches(l)]
|
|
|
|
f = (Path(path) / Path(name)).with_suffix(".txt")
|
|
try:
|
|
with open(f, "r") as file:
|
|
return [l.strip() for l in file.readlines() if l.strip() and matches(l)]
|
|
except:
|
|
log.warning("Wildcard file not found for %s", name)
|
|
return [name]
|
|
|
|
|
|
def replace_wildcards(text, delayed_matches=None):
|
|
if delayed_matches is not None:
|
|
# replace delayed match placeholders with names again to properly select them
|
|
for ph, v in delayed_matches.items():
|
|
offset = v["offset"] or 0
|
|
name = v["name"]
|
|
text = text.replace(ph, f"${name}+{offset}$")
|
|
wildcard_re = re.compile(r"\$(?P<name>[A-Za-z0-9_/.!:?-]+)(\+(?P<offset>[0-9]+))?\$")
|
|
matches, text = find_and_remove(wildcard_re, text, placeholder="MU_WILDCARD")
|
|
delayed = {}
|
|
for placeholder, value in matches.items():
|
|
state = None
|
|
name = value["name"]
|
|
if name.startswith("?") and delayed_matches is None:
|
|
# Handle wildcard after macro expansion
|
|
value["name"] = name[1:]
|
|
delayed[placeholder] = value
|
|
continue
|
|
ws = read_wildcards(name)
|
|
offset = int(value["offset"] or 0)
|
|
if ws and offset:
|
|
# advance the state once and store it so that the next non-offset result stays deterministic
|
|
w = RAND.choice(ws)
|
|
state = RAND.getstate()
|
|
# advance the state until offset,
|
|
for _ in range(offset):
|
|
w = RAND.choice(ws)
|
|
elif ws:
|
|
w = RAND.choice(ws)
|
|
else:
|
|
log.warning("No wildcards found for %s", value["name"])
|
|
w = ""
|
|
text = text.replace(placeholder, w)
|
|
log.info("Replaced wildcard %s with '%s'", value["name"], w)
|
|
if state:
|
|
RAND.setstate(state)
|
|
|
|
return text.strip(), delayed
|
|
|
|
|
|
class MUSimpleWildcard:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"default": "", "multiline": True}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
},
|
|
"optional": {"use_pnginfo": ("BOOLEAN", {"default": False})},
|
|
"hidden": {
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
CATEGORY = "misc-utils"
|
|
FUNCTION = "doit"
|
|
|
|
def doit(self, text, seed, extra_pnginfo, unique_id, use_pnginfo=False):
|
|
if use_pnginfo and unique_id in extra_pnginfo.get(CLASS_NAME, {}):
|
|
text = extra_pnginfo[CLASS_NAME][unique_id]
|
|
log.info("MUSimpleWildcard using prompt: %s", text)
|
|
return (text,)
|
|
|
|
|
|
try:
|
|
from server import PromptServer
|
|
import folder_paths
|
|
|
|
PromptServer.instance.add_on_prompt_handler(wildcard_prompt_handler)
|
|
except ImportError:
|
|
print("Could not install wildcard prompt handler, node won't work")
|
|
|
|
if __name__ == "__main__":
|
|
print("Start parse")
|
|
ctx = read_preamble()
|
|
print("End parse")
|