diff --git a/prompt_control/macros.py b/prompt_control/macros.py new file mode 100644 index 0000000..18a4954 --- /dev/null +++ b/prompt_control/macros.py @@ -0,0 +1,81 @@ +# vim: sw=4 ts=4 +from __future__ import annotations + +import logging +import re + +from .utils import find_closing_paren, get_function + +logging.basicConfig() +log = logging.getLogger("comfyui-prompt-control") + + +def parse_search(search): + arg_start = search.find("(") + args = "" + name = search.strip() + if arg_start > 0: + arg_end = find_closing_paren(search, arg_start + 1) + if arg_end < 0: + arg_end = len(search) + name = search[:arg_start].strip() + args = search[arg_start + 1 : arg_end] + + if not name: + return None + args = args.strip() + # If using the form DEF(F()=$1) then the default value of $1 is the empty string + args = [a.strip() for a in args.split(";")] if arg_start > 0 else [] + return name, args + + +def expand_macros(text): + text, defs = get_function(text, "DEF", defaults=None) + res = text + prevres = text + replacements = [] + for d in defs: + if not d.args: + continue + r = d.args[0].split("=", 1) + search = parse_search(r[0].strip()) + if not search or len(r) != 2: + log.warning("Ignoring invalid DEF(%s)", d) + continue + replacements.append((search, r[1].strip())) + iterations = 0 + while True: + iterations += 1 + if iterations > 10: + raise ValueError("Unable to resolve DEFs, make sure there are no cycles!") + return text + for search, replace in replacements: + res = substitute_defcall(res, search, replace) + if res == prevres: + break + prevres = res + if res.strip() != text.strip(): + res = res.strip() + log.info("DEFs expanded to: %s", res) + return res + + +def substitute_defcall(text, search, replace): + name, default_args = search + text, defns = get_function(text, name, defaults=None, placeholder=f"DEFNCALL{name}", require_args=False) + for i, d in enumerate(defns): + ph = d.placeholder + assert ph is not None, "This is a bug" + parameters = d.args + paramvals = [] + if parameters: + paramvals = [x.strip() for x in parameters[0].split(";")] + r = replace + for i, v in enumerate(paramvals): + r = re.sub(rf"\${i + 1}\b", v, r) + + for i, v in enumerate(default_args): + r = re.sub(rf"\${i + 1}\b", v, r) + + text = text.replace(ph, r) + return text diff --git a/prompt_control/parser.py b/prompt_control/parser.py index ded13ad..676125f 100644 --- a/prompt_control/parser.py +++ b/prompt_control/parser.py @@ -2,14 +2,13 @@ from __future__ import annotations import logging -import re from functools import lru_cache from math import ceil from typing import TypeAlias import lark -from .utils import find_closing_paren, get_function +from .macros import expand_macros logging.basicConfig() log = logging.getLogger("comfyui-prompt-control") @@ -414,77 +413,6 @@ class PromptSchedule: return len(self.parsed_prompt) - 1, self.parsed_prompt[-1] -def parse_search(search): - arg_start = search.find("(") - args = "" - name = search.strip() - if arg_start > 0: - arg_end = find_closing_paren(search, arg_start + 1) - if arg_end < 0: - arg_end = len(search) - name = search[:arg_start].strip() - args = search[arg_start + 1 : arg_end] - - if not name: - return None - args = args.strip() - # If using the form DEF(F()=$1) then the default value of $1 is the empty string - args = [a.strip() for a in args.split(";")] if arg_start > 0 else [] - return name, args - - -def expand_macros(text): - text, defs = get_function(text, "DEF", defaults=None) - res = text - prevres = text - replacements = [] - for d in defs: - if not d.args: - continue - r = d.args[0].split("=", 1) - search = parse_search(r[0].strip()) - if not search or len(r) != 2: - log.warning("Ignoring invalid DEF(%s)", d) - continue - replacements.append((search, r[1].strip())) - iterations = 0 - while True: - iterations += 1 - if iterations > 10: - raise ValueError("Unable to resolve DEFs, make sure there are no cycles!") - return text - for search, replace in replacements: - res = substitute_defcall(res, search, replace) - if res == prevres: - break - prevres = res - if res.strip() != text.strip(): - res = res.strip() - log.info("DEFs expanded to: %s", res) - return res - - -def substitute_defcall(text, search, replace): - name, default_args = search - text, defns = get_function(text, name, defaults=None, placeholder=f"DEFNCALL{name}", require_args=False) - for i, d in enumerate(defns): - ph = d.placeholder - assert ph is not None, "This is a bug" - parameters = d.args - paramvals = [] - if parameters: - paramvals = [x.strip() for x in parameters[0].split(";")] - r = replace - for i, v in enumerate(paramvals): - r = re.sub(rf"\${i + 1}\b", v, r) - - for i, v in enumerate(default_args): - r = re.sub(rf"\${i + 1}\b", v, r) - - text = text.replace(ph, r) - return text - - @lru_cache def parse_prompt_schedules(prompt, **kwargs): prompt = expand_macros(prompt)