mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-11 14:02:37 +08:00
264 lines
8.0 KiB
Python
264 lines
8.0 KiB
Python
import re
|
|
|
|
from comfy.graph_utils import GraphBuilder
|
|
|
|
class InversionDemoAdvancedPromptNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True}),
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "CONDITIONING")
|
|
FUNCTION = "advanced_prompt"
|
|
|
|
CATEGORY = "InversionDemo Nodes"
|
|
|
|
def parse_prompt(self, prompt):
|
|
# Get all string pieces matching the pattern "<lora:(name):(strength)(:(clip_strength))?>"
|
|
# where name is a string and strength is a float
|
|
# and clip_strength is an optional float
|
|
pattern = r"<lora:([^:]+):([-0-9.]+)(?::([-0-9.]+))?>"
|
|
loras = re.findall(pattern, prompt)
|
|
if len(loras) == 0:
|
|
return prompt, loras
|
|
cleaned_prompt = re.sub(pattern, "", prompt).strip()
|
|
print("Cleaned prompt: '%s'" % cleaned_prompt)
|
|
return cleaned_prompt, loras
|
|
|
|
|
|
def advanced_prompt(self, prompt, clip, model):
|
|
cleaned_prompt, loras = self.parse_prompt(prompt)
|
|
graph = GraphBuilder()
|
|
for lora in loras:
|
|
lora_name = lora[0]
|
|
lora_model_strength = float(lora[1])
|
|
lora_clip_strength = lora_model_strength if lora[2] == "" else float(lora[2])
|
|
|
|
loader = graph.node("LoraLoader", model=model, clip=clip, lora_name = lora_name, strength_model = lora_model_strength, strength_clip = lora_clip_strength)
|
|
model = loader.out(0)
|
|
clip = loader.out(1)
|
|
encoder = graph.node("CLIPTextEncode", clip=clip, text=cleaned_prompt)
|
|
return {
|
|
"result": (model, clip, encoder.out(0)),
|
|
"expand": graph.finalize(),
|
|
}
|
|
|
|
class InversionDemoFakeAdvancedPromptNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True}),
|
|
"clip": ("CLIP",),
|
|
"model": ("MODEL",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "CONDITIONING")
|
|
FUNCTION = "advanced_prompt"
|
|
|
|
CATEGORY = "InversionDemo Nodes"
|
|
|
|
def advanced_prompt(self, prompt, clip, model):
|
|
tokens = clip.tokenize(prompt)
|
|
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
|
return (model, clip, [[cond, {"pooled_output": pooled}]])
|
|
|
|
class InversionDemoLazySwitch:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"switch": ("BOOL",),
|
|
"on_false": ("*", {"lazy": True}),
|
|
"on_true": ("*", {"lazy": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("*",)
|
|
FUNCTION = "switch"
|
|
|
|
CATEGORY = "InversionDemo Nodes"
|
|
|
|
def check_lazy_status(self, switch, on_false = None, on_true = None):
|
|
if switch and on_true is None:
|
|
return ["on_true"]
|
|
if not switch and on_false is None:
|
|
return ["on_false"]
|
|
|
|
def switch(self, switch, on_false = None, on_true = None):
|
|
value = on_true if switch else on_false
|
|
return (value,)
|
|
|
|
NUM_IF_ELSE_NODES = 10
|
|
class InversionDemoLazyConditional:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
args = {
|
|
"value1": ("*", {"lazy": True}),
|
|
"condition1": ("BOOL", {"forceInput": True}),
|
|
}
|
|
|
|
for i in range(1,NUM_IF_ELSE_NODES):
|
|
args["value%d" % (i + 1)] = ("*", {"lazy": True})
|
|
args["condition%d" % (i + 1)] = ("BOOL", {"lazy": True, "forceInput": True})
|
|
|
|
args["else"] = ("*", {"lazy": True})
|
|
|
|
return {
|
|
"required": {},
|
|
"optional": args,
|
|
}
|
|
|
|
RETURN_TYPES = ("*",)
|
|
FUNCTION = "conditional"
|
|
|
|
CATEGORY = "InversionDemo Nodes"
|
|
|
|
def check_lazy_status(self, **kwargs):
|
|
for i in range(0,NUM_IF_ELSE_NODES):
|
|
cond = "condition%d" % (i + 1)
|
|
if cond not in kwargs:
|
|
return [cond]
|
|
if kwargs[cond]:
|
|
val = "value%d" % (i + 1)
|
|
if val not in kwargs:
|
|
return [val]
|
|
else:
|
|
return []
|
|
|
|
if "else" not in kwargs:
|
|
return ["else"]
|
|
|
|
def conditional(self, **kwargs):
|
|
for i in range(0,NUM_IF_ELSE_NODES):
|
|
cond = "condition%d" % (i + 1)
|
|
if cond not in kwargs:
|
|
return [cond]
|
|
if kwargs.get(cond, False):
|
|
val = "value%d" % (i + 1)
|
|
return (kwargs.get(val, None),)
|
|
|
|
return (kwargs.get("else", None),)
|
|
|
|
|
|
class InversionDemoLazyIndexSwitch:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"index": ("INT", {"default": 0, "min": 0, "max": 9, "step": 1}),
|
|
"value0": ("*", {"lazy": True}),
|
|
},
|
|
"optional": {
|
|
"value1": ("*", {"lazy": True}),
|
|
"value2": ("*", {"lazy": True}),
|
|
"value3": ("*", {"lazy": True}),
|
|
"value4": ("*", {"lazy": True}),
|
|
"value5": ("*", {"lazy": True}),
|
|
"value6": ("*", {"lazy": True}),
|
|
"value7": ("*", {"lazy": True}),
|
|
"value8": ("*", {"lazy": True}),
|
|
"value9": ("*", {"lazy": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("*",)
|
|
FUNCTION = "index_switch"
|
|
|
|
CATEGORY = "InversionDemo Nodes"
|
|
|
|
def check_lazy_status(self, index, **kwargs):
|
|
key = "value%d" % index
|
|
if key not in kwargs:
|
|
return [key]
|
|
|
|
def index_switch(self, index, **kwargs):
|
|
key = "value%d" % index
|
|
return kwargs[key]
|
|
|
|
class InversionDemoLazyMixImages:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image1": ("IMAGE",{"lazy": True}),
|
|
"image2": ("IMAGE",{"lazy": True}),
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "mix"
|
|
|
|
CATEGORY = "InversionDemo Nodes"
|
|
|
|
def check_lazy_status(self, mask, image1 = None, image2 = None):
|
|
mask_min = mask.min()
|
|
mask_max = mask.max()
|
|
needed = []
|
|
if image1 is None and (mask_min != 1.0 or mask_max != 1.0):
|
|
needed.append("image1")
|
|
if image2 is None and (mask_min != 0.0 or mask_max != 0.0):
|
|
needed.append("image2")
|
|
return needed
|
|
|
|
# Not trying to handle different batch sizes here just to keep the demo simple
|
|
def mix(self, mask, image1 = None, image2 = None):
|
|
mask_min = mask.min()
|
|
mask_max = mask.max()
|
|
if mask_min == 0.0 and mask_max == 0.0:
|
|
return (image1,)
|
|
elif mask_min == 1.0 and mask_max == 1.0:
|
|
return (image2,)
|
|
|
|
if len(mask.shape) == 2:
|
|
mask = mask.unsqueeze(0)
|
|
if len(mask.shape) == 3:
|
|
mask = mask.unsqueeze(3)
|
|
if mask.shape[3] < image1.shape[3]:
|
|
mask = mask.repeat(1, 1, 1, image1.shape[3])
|
|
|
|
return (image1 * (1. - mask) + image2 * mask,)
|
|
|
|
GENERAL_NODE_CLASS_MAPPINGS = {
|
|
"InversionDemoAdvancedPromptNode": InversionDemoAdvancedPromptNode,
|
|
"InversionDemoFakeAdvancedPromptNode": InversionDemoFakeAdvancedPromptNode,
|
|
"InversionDemoLazySwitch": InversionDemoLazySwitch,
|
|
"InversionDemoLazyIndexSwitch": InversionDemoLazyIndexSwitch,
|
|
"InversionDemoLazyMixImages": InversionDemoLazyMixImages,
|
|
"InversionDemoLazyConditional": InversionDemoLazyConditional,
|
|
}
|
|
|
|
GENERAL_NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"InversionDemoAdvancedPromptNode": "Advanced Prompt",
|
|
"InversionDemoFakeAdvancedPromptNode": "Fake Advanced Prompt",
|
|
"InversionDemoLazySwitch": "Lazy Switch",
|
|
"InversionDemoLazyIndexSwitch": "Lazy Index Switch",
|
|
"InversionDemoLazyMixImages": "Lazy Mix Images",
|
|
"InversionDemoLazyConditional": "Lazy Conditional",
|
|
}
|