From 69d4eadfab591d7d784c76e6c4e4b78913df6601 Mon Sep 17 00:00:00 2001 From: asagi4 <130366179+asagi4@users.noreply.github.com> Date: Thu, 2 Nov 2023 21:24:33 +0200 Subject: [PATCH] Initial commit --- Makefile | 8 ++++ README.md | 31 +++++++++++++ __init__.py | 10 ++++ utils/jinja_render.py | 83 +++++++++++++++++++++++++++++++++ utils/other.py | 50 ++++++++++++++++++++ utils/wildcards.py | 104 ++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 286 insertions(+) create mode 100644 Makefile create mode 100644 README.md create mode 100644 __init__.py create mode 100644 utils/jinja_render.py create mode 100644 utils/other.py create mode 100644 utils/wildcards.py diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..3ba0c43 --- /dev/null +++ b/Makefile @@ -0,0 +1,8 @@ +all: check format + @echo "Done" +check: + pyflakes *.py */*.py +format: + black -l 120 *.py */*.py + +.PHONY: check format all diff --git a/README.md b/README.md new file mode 100644 index 0000000..0ca46cc --- /dev/null +++ b/README.md @@ -0,0 +1,31 @@ +# ComfyUI utility nodes + +A collection of miscellaneous utility nodes for ComfyUI + +# Nodes + +## MU_JinjaRender +You can use this node to evaluate a string as a Jinja2 template. Note, however, that because ComfyUI's frontend uses `{}` for syntax, There are the following modifications to Jinja syntax: + +- `{% %}` becomes `<% %>` +- `{{ }}` becomes `<= =>` +- `{# #}` becomes `<# #>` + +### Functions in Jinja templates + +The following functions and constants are available: + +- `pi` +- `min`, `max`, `clamp(minimum, value, maximum)`, +- `abs`, `round`, `ceil`, `floor` +- `sqrt` `sin`, `cos`, `tan`, `asin`, `acos`, `atan`. These functions are rounded to two decimals + + +In addition, a special `steps` function exists. + +The `steps` function will generate a list of steps for iterating. + +You can call it either as `steps(end)`, `steps(end, step=0.1)` or `steps(start, end, step)`. `step` is an optional parameter that defaults to `0.1`. It'll return steps *inclusive* of start and end as long as step doesn't go past the end. + +The second form is equivalent to `steps(step, end, step)`. i.e. it starts at the first step. + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..c70849a --- /dev/null +++ b/__init__.py @@ -0,0 +1,10 @@ +from .utils import jinja_render as j +from .utils import wildcards as w +from .utils import other as o + +NODE_CLASS_MAPPINGS = { + "MUJinjaRender": j.MUJinjaRender, + "MUSimpleWildcard": w.MUSimpleWildcard, + "MUConditioningCutoff": o.MUConditioningCutoff, + "MUStringConcat": o.MUStringConcat, +} diff --git a/utils/jinja_render.py b/utils/jinja_render.py new file mode 100644 index 0000000..fccf161 --- /dev/null +++ b/utils/jinja_render.py @@ -0,0 +1,83 @@ +import logging +import math +import re +from datetime import datetime + +log = logging.getLogger("comfyui-misc-utils") + + +def steps(start, end=None, step=0.1): + if end is None: + end = start + start = step + while start <= end: + yield start + start += step + start = round(start, 2) + + +def template(template, sequence, *funcs): + funcs = [lambda x: x] + list(*funcs) + res = [] + for item in sequence: + x = template + for i, f in enumerate(funcs): + x = x.replace(f"${i}", str(f(i))) + res.append(x) + + return "".join(res) + + +def clamp(a, b, c): + return max(a, min(b, c)) + + +JINJA_ENV = { + "pi": math.pi, + "floor": math.floor, + "ceil": math.ceil, + "min": min, + "max": max, + "abs": abs, + "clamp": clamp, + "round": round, + "template": template, + "steps": steps, + "datetime": datetime, +} + +for fname in ["sqrt", "sin", "cos", "tan", "asin", "acos", "atan"]: + f = getattr(math, fname) + JINJA_ENV[fname] = lambda x: round(f(x), 2) + + +def render_jinja(text): + from jinja2 import Environment + + jenv = Environment( + block_start_string="<%", + block_end_string="%>", + variable_start_string="<=", + variable_end_string="=>", + comment_start_string="<#", + comment_end_string="#>", + ) + + return jenv.from_string(text, globals=JINJA_ENV).render() + + +class MUJinjaRender: + @classmethod + def INPUT_TYPES(s): + return {"required": {"text": ("STRING", {"default": "", "multiline": True})}} + + RETURN_TYPES = ("STRING",) + + CATEGORY = "promptcontrol/tools" + FUNCTION = "render" + + def render(self, text): + t = render_jinja(text) + if t.strip() != text.strip(): + log.info("Jinja render result: %s", re.sub("\s+", " ", t, flags=re.MULTILINE)) + return (t,) diff --git a/utils/other.py b/utils/other.py new file mode 100644 index 0000000..9bcb886 --- /dev/null +++ b/utils/other.py @@ -0,0 +1,50 @@ +class MUStringConcat: + @classmethod + def INPUT_TYPES(s): + t = ("STRING", {"default": ""}) + return { + "optional": { + "string1": t, + "string2": t, + "string3": t, + "string4": t, + } + } + + RETURN_TYPES = ("STRING",) + + CATEGORY = "miscutils" + FUNCTION = "cat" + + def cat(self, string1="", string2="", string3="", string4=""): + return string1 + string2 + string3 + string4 + + +class MUConditioningCutoff: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "conds": ("CONDITIONING",), + "cutoff": ("FLOAT", {"min": 0.00, "max": 1.00, "default": 0.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("CONDITIONING",) + CATEGORY = "miscutils" + FUNCTION = "apply" + + def apply(self, conds, cutoff): + res = [] + new_start = 1.0 + for c in conds: + end = c[1].get("end_percent", 0.0) + if 1.0 - end < cutoff: + continue + c = [c[0].clone(), c[1].copy()] + c[1]["start_percent"] = new_start + c[1]["end_percent"] = end + new_start = end + res.append(c) + + return (res,) diff --git a/utils/wildcards.py b/utils/wildcards.py new file mode 100644 index 0000000..b38644e --- /dev/null +++ b/utils/wildcards.py @@ -0,0 +1,104 @@ +import logging +from os import environ +import random +import re +from pathlib import Path +from server import PromptServer + +log = logging.getLogger("comfyui-misc-utils") + +CLASS_NAME = "MUSimpleWildcard" + + +def wildcard_prompt_handler(json_data): + log.info("Resolving wildcards...") + 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): + wildcard_info = json_data.get("extra_data", {}).get("extra_pnginfo", {}).get(CLASS_NAME, {}) + n = json_data["prompt"][node_id] + if not (n["inputs"].get("use_pnginfo") and node_id in wildcard_info): + text = MUSimpleWildcard.select(n["inputs"]["text"], n["inputs"]["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 + + +PromptServer.instance.add_on_prompt_handler(wildcard_prompt_handler) + + +def variable_substitution(text): + var_re = re.compile(r"(\$[a-z]+)\s*=([^;\n]*);?") + m = var_re.search(text) + while m: + var = m[1] + sub = m[2] + s, e = m.span() + text = text[:s] + text[e:] + log.info("Substituting %s with '%s'", var, sub) + text = text.replace(var, sub) + m = var_re.search(text) + return text + + +class MUSimpleWildcard: + RAND = random.Random() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "text": ("STRING", {"default": ""}), + "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 = "promptcontrol/tools" + FUNCTION = "doit" + + @classmethod + def read_wildcards(cls, name): + path = environ.get("MU_WILDCARD_BASEDIR", "wildcards") + + 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()] + except: + return [name] + + @classmethod + def select(cls, text, seed): + cls.RAND.seed(seed) + matches = re.findall(r"(\$([A-Za-z0-9_/.-]+)(\+[0-9]+)?\$)", text) + for placeholder, wildcard, offset in matches: + if offset: + offset = int(offset[1:]) + cls.RAND.seed(seed + offset) + w = cls.RAND.choice(cls.read_wildcards(wildcard)) + text = text.replace(placeholder, w, 1) + log.info("Selected wildcard %s for %s", w, placeholder) + if offset: + cls.RAND.seed(seed) + return text + + 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) + text = variable_substitution(text) + return (text,)