260 lines
9.8 KiB
Python
260 lines
9.8 KiB
Python
import json
|
|
import logging
|
|
|
|
from comfy_api.latest import io
|
|
|
|
from .macros import expand_macros as macroexpand
|
|
from .macros import expand_segs as segexpand
|
|
from .macros import expand_subs as subexpand
|
|
from .macros import substitute_var
|
|
from .parser import parse_prompt_schedules
|
|
|
|
log = logging.getLogger("comfyui-prompt-control")
|
|
|
|
|
|
class PCSetLogLevel(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="PCSetLogLevel",
|
|
display_name="PC: Configure Logging (for debug)",
|
|
category="promptcontrol/tools",
|
|
description="A debug node to configure Prompt Control logging level. Pass a CLIP through it before you run any PC nodes",
|
|
inputs=[
|
|
io.Clip.Input("clip"),
|
|
io.Combo.Input("level", options=["INFO", "DEBUG", "WARNING", "ERROR"], default="INFO", optional=True),
|
|
],
|
|
outputs=[io.Clip.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, clip, level="INFO") -> io.NodeOutput:
|
|
log.setLevel(getattr(logging, level))
|
|
log.info("Set logging level to %s", level)
|
|
return io.NodeOutput(clip)
|
|
|
|
|
|
class PCAddMaskToCLIP(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="PCAddMaskToCLIP",
|
|
display_name="PC: Attach Mask",
|
|
category="promptcontrol/tools",
|
|
description="Attaches a mask to a CLIP object so that they can be referred to in a prompt using IMASK(). Using this node multiple times adds more masks rather than replacing existing ones.",
|
|
inputs=[
|
|
io.Clip.Input("clip"),
|
|
io.Mask.Input("mask", optional=True),
|
|
],
|
|
outputs=[io.Clip.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, clip, mask=None) -> io.NodeOutput:
|
|
return PCAddMaskToCLIPMany.execute(clip, mask1=mask)
|
|
|
|
|
|
class PCAddMaskToCLIPMany(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="PCAddMaskToCLIPMany",
|
|
display_name="PC: Attach Mask (multi)",
|
|
category="promptcontrol/tools",
|
|
description="Multi-input version of PCAddMaskToCLIP, for convenience",
|
|
inputs=[
|
|
io.Clip.Input("clip"),
|
|
io.Mask.Input("mask1", optional=True),
|
|
io.Mask.Input("mask2", optional=True),
|
|
io.Mask.Input("mask3", optional=True),
|
|
io.Mask.Input("mask4", optional=True),
|
|
],
|
|
outputs=[io.Clip.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, clip, mask1=None, mask2=None, mask3=None, mask4=None) -> io.NodeOutput:
|
|
clip = clip.clone()
|
|
current_masks = clip.patcher.model_options.get("x-promptcontrol.masks", [])
|
|
current_masks.extend(m for m in (mask1, mask2, mask3, mask4) if m is not None)
|
|
clip.patcher.model_options["x-promptcontrol.masks"] = current_masks
|
|
return io.NodeOutput(clip)
|
|
|
|
|
|
class PCSetPCTextEncodeSettings(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="PCSetPCTextEncodeSettings",
|
|
display_name="PC: Configure PCTextEncode",
|
|
category="promptcontrol/tools",
|
|
description="Configures default values for PCTextEncode",
|
|
inputs=[
|
|
io.Clip.Input("clip"),
|
|
io.Int.Input("mask_width", default=512, min=64, max=4096 * 4, optional=True),
|
|
io.Int.Input("mask_height", default=512, min=64, max=4096 * 4, optional=True),
|
|
io.Int.Input("sdxl_width", default=1024, min=0, max=4096 * 4, optional=True),
|
|
io.Int.Input("sdxl_height", default=1024, min=0, max=4096 * 4, optional=True),
|
|
io.Int.Input("sdxl_target_w", default=1024, min=0, max=4096 * 4, optional=True),
|
|
io.Int.Input("sdxl_target_h", default=1024, min=0, max=4096 * 4, optional=True),
|
|
io.Int.Input("sdxl_crop_w", default=0, min=0, max=4096 * 4, optional=True),
|
|
io.Int.Input("sdxl_crop_h", default=0, min=0, max=4096 * 4, optional=True),
|
|
],
|
|
outputs=[io.Clip.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
clip,
|
|
mask_width=512,
|
|
mask_height=512,
|
|
sdxl_width=1024,
|
|
sdxl_height=1024,
|
|
sdxl_target_w=1024,
|
|
sdxl_target_h=1024,
|
|
sdxl_crop_w=0,
|
|
sdxl_crop_h=0,
|
|
) -> io.NodeOutput:
|
|
settings = {
|
|
"mask_width": mask_width,
|
|
"mask_height": mask_height,
|
|
"sdxl_width": sdxl_width,
|
|
"sdxl_height": sdxl_height,
|
|
"sdxl_twidth": sdxl_target_w,
|
|
"sdxl_theight": sdxl_target_h,
|
|
"sdxl_cwidth": sdxl_crop_w,
|
|
"sdxl_cheight": sdxl_crop_h,
|
|
}
|
|
clip = clip.clone()
|
|
clip.patcher.model_options["x-promptcontrol.settings"] = settings
|
|
return io.NodeOutput(clip)
|
|
|
|
|
|
class PCExtractScheduledPrompt(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="PCExtractScheduledPrompt",
|
|
display_name="PC: Show Prompt",
|
|
category="promptcontrol/tools",
|
|
description="Parses the input prompt and returns the prompt scheduled at the specified point",
|
|
inputs=[
|
|
io.String.Input("text", multiline=True),
|
|
io.Float.Input("at", min=0.0, max=1.0, default=1.0, step=0.01),
|
|
io.String.Input("tags", default="", optional=True),
|
|
io.Boolean.Input("expand_segs", default=False, optional=True),
|
|
io.Boolean.Input("expand_subs", default=False, optional=True),
|
|
io.Boolean.Input("expand_macros", default=False, optional=True),
|
|
],
|
|
outputs=[io.String.Output()],
|
|
search_aliases=["extract scheduled prompt"],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, text, at, tags="", expand_segs=False, expand_subs=False, expand_macros=False) -> io.NodeOutput:
|
|
if expand_macros:
|
|
text = macroexpand(text)
|
|
schedule = parse_prompt_schedules(text, filters=tags)
|
|
_, entry = schedule.at_step(at)
|
|
prompt_text = entry.get("prompt", "")
|
|
if expand_segs:
|
|
prompt_text = segexpand(prompt_text, do_subs=expand_subs)
|
|
if expand_subs:
|
|
prompt_text = subexpand(prompt_text)
|
|
return io.NodeOutput(prompt_text)
|
|
|
|
|
|
class PCMacroExpand(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="PCMacroExpand",
|
|
display_name="PC: Expand Macros",
|
|
category="promptcontrol/tools",
|
|
description="Expands DEF macros in a string and returns the result",
|
|
inputs=[
|
|
io.String.Input("text", multiline=True),
|
|
],
|
|
outputs=[io.String.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, text) -> io.NodeOutput:
|
|
return io.NodeOutput(macroexpand(text))
|
|
|
|
|
|
class PCLinkHelper(io.ComfyNode):
|
|
# a-z
|
|
NAMES = [chr(97 + i) for i in range(26)]
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
t1 = io.Autogrow.TemplateNames(io.AnyType.Input("link", raw_link=True), min=0, names=cls.NAMES)
|
|
t2 = io.Autogrow.TemplateNames(
|
|
io.AnyType.Input("value", lazy=True), min=0, names=[f"var{i + 1}" for i in range(50)]
|
|
)
|
|
return io.Schema(
|
|
node_id="PCNODELinkHelper",
|
|
display_name="PC: Extra argument helper for NODE",
|
|
category="promptcontrol/tools",
|
|
description="Takes in arbitrary inputs and renders them as NODE-compatible values, replacing $a -> $z with JSON link values.",
|
|
is_experimental=True,
|
|
inputs=[
|
|
io.Autogrow.Input("links", template=t1),
|
|
io.Autogrow.Input(
|
|
"vars",
|
|
template=t2,
|
|
),
|
|
io.String.Input(
|
|
"template",
|
|
tooltip="The variables $a to $z will be replaced in this text with their corresponding input's JSON link value",
|
|
placeholder="In this text you can refer to the input links as $a, $b etc. and the var inputs as either $var1 or $json1 etc. (the latter will be rendered through Python's json.dumps function which will cause strings to be quoted)",
|
|
multiline=True,
|
|
),
|
|
],
|
|
outputs=[io.String.Output()],
|
|
)
|
|
|
|
# This requires https://github.com/Comfy-Org/ComfyUI/pull/15103 to work properly
|
|
# Without that PR, all inputs will be evaluated non-lazily
|
|
@classmethod
|
|
def check_lazy_status(cls, template, links, vars):
|
|
r = []
|
|
for name, (v, input_name) in vars.items():
|
|
if v is None and f"${name}" in template or v is None and f"$json{name[3:]}" in template:
|
|
r.append(input_name)
|
|
return r
|
|
|
|
@classmethod
|
|
def execute(cls, template, links, vars) -> io.NodeOutput:
|
|
text = template
|
|
for k in cls.NAMES:
|
|
v = "__EMPTY__"
|
|
if k in links:
|
|
# Replace : with \: to avoid breaking scheduling syntax when linking subgraphs. Any function that consumes this should replace \: with :
|
|
v = json.dumps(links[k]).replace(":", r"\:")
|
|
text = substitute_var(text, k, v)
|
|
for i in range(50):
|
|
v = "__EMPTY__"
|
|
k = f"var{i + 1}"
|
|
if k in vars:
|
|
v = vars[k]
|
|
text = substitute_var(text, k, str(v))
|
|
if f"$json{i + 1}" in text:
|
|
v = v if v == "__EMPTY__" else json.dumps(v)
|
|
text = substitute_var(text, f"json{i + 1}", v)
|
|
|
|
return io.NodeOutput(text)
|
|
|
|
|
|
NODES = [
|
|
PCSetPCTextEncodeSettings,
|
|
PCAddMaskToCLIP,
|
|
PCAddMaskToCLIPMany,
|
|
PCSetLogLevel,
|
|
PCExtractScheduledPrompt,
|
|
PCMacroExpand,
|
|
PCLinkHelper,
|
|
]
|