322 lines
15 KiB
Python
322 lines
15 KiB
Python
import json
|
|
from ..exif.base_format import BaseFormat
|
|
from ....components import utility
|
|
|
|
# comfyui node types
|
|
KSAMPLER_TYPES = ["KSampler", "KSamplerAdvanced", "PrimereKSampler"]
|
|
VAE_ENCODE_TYPE = ["VAEEncode", "VAEEncodeForInpaint"]
|
|
CHECKPOINT_LOADER_TYPE = [
|
|
"CheckpointLoader",
|
|
"CheckpointLoaderSimple",
|
|
"unCLIPCheckpointLoader",
|
|
"Checkpoint Loader (Simple)",
|
|
"PrimereCKPTLoader",
|
|
"PrimereVisualCKPT",
|
|
"PrimereCKPT"
|
|
]
|
|
CLIP_TEXT_ENCODE_TYPE = [
|
|
"CLIPTextEncode",
|
|
"CLIPTextEncodeSDXL",
|
|
"CLIPTextEncodeSDXLRefiner",
|
|
"PrimereCLIPEncoder"
|
|
]
|
|
SAVE_IMAGE_TYPE = ["SaveImage", "Image Save", "PrimerePreviewImage"]
|
|
|
|
class ComfyUI(BaseFormat):
|
|
def __init__(self, info: dict = None, raw: str = ""):
|
|
super().__init__(info, raw)
|
|
self._comfy_png()
|
|
|
|
def _comfy_png(self):
|
|
prompt = self._info.get("prompt") or {}
|
|
gendata_json = {}
|
|
if 'gendata' in self._info:
|
|
gendata = self._info.get("gendata") or {}
|
|
gendata_json = json.loads(gendata)
|
|
workflow = self._info.get("workflow") or {}
|
|
prompt_json = json.loads(prompt)
|
|
|
|
if len(gendata_json) > 0:
|
|
FINAL_DICT = gendata_json
|
|
self._parameter = FINAL_DICT
|
|
else:
|
|
# find end node of each flow
|
|
end_nodes = list(filter( lambda item: item[-1].get("class_type") in ["SaveImage"] + KSAMPLER_TYPES, prompt_json.items(),))
|
|
longest_flow = {}
|
|
longest_nodes = []
|
|
longest_flow_len = 0
|
|
|
|
for end_node in end_nodes:
|
|
flow, nodes = self._comfy_traverse(prompt_json, str(end_node[0]))
|
|
if len(nodes) > longest_flow_len:
|
|
longest_flow = flow
|
|
longest_nodes = nodes
|
|
longest_flow_len = len(nodes)
|
|
|
|
SizeID = None
|
|
ModelID = None
|
|
PositiveID = None
|
|
NegativeID = None
|
|
|
|
try:
|
|
if 'latent_image' in flow:
|
|
SizeID = flow['latent_image'][0]
|
|
except Exception:
|
|
SizeID = None
|
|
|
|
try:
|
|
if 'model' in flow:
|
|
ModelID = flow['model'][0]
|
|
except Exception:
|
|
ModelID = None
|
|
|
|
try:
|
|
if 'positive' in flow:
|
|
PositiveID = flow['positive'][0]
|
|
except Exception:
|
|
PositiveID = None
|
|
|
|
try:
|
|
if 'negative' in flow:
|
|
NegativeID = flow['negative'][0]
|
|
except Exception:
|
|
NegativeID = None
|
|
|
|
FINAL_DICT = {}
|
|
FINAL_DICT['negative'] = ""
|
|
FINAL_DICT['positive'] = ""
|
|
|
|
if PositiveID and NegativeID and 'text_g' in prompt_json[PositiveID]['inputs']:
|
|
FINAL_DICT['positive'] = prompt_json[PositiveID]['inputs']['text_g']
|
|
FINAL_DICT['negative'] = prompt_json[NegativeID]['inputs']['text_g']
|
|
|
|
if PositiveID and NegativeID and 'text' in prompt_json[PositiveID]['inputs']:
|
|
FINAL_DICT['positive'] = prompt_json[PositiveID]['inputs']['text']
|
|
FINAL_DICT['negative'] = prompt_json[NegativeID]['inputs']['text']
|
|
|
|
if PositiveID == None or ('text_g' not in prompt_json[PositiveID]['inputs'] and 'text' not in prompt_json[PositiveID]['inputs']):
|
|
if hasattr(self, '_positive'):
|
|
FINAL_DICT['positive'] = self._positive
|
|
if hasattr(self, '_negative'):
|
|
FINAL_DICT['negative'] = self._negative
|
|
|
|
if 'steps' in flow and type(flow['steps']) == int:
|
|
FINAL_DICT['steps'] = flow['steps']
|
|
if 'sampler_name' in flow and 'scheduler' in flow and type(flow['sampler_name']) == str and type(flow['scheduler']) == str:
|
|
FINAL_DICT['sampler'] = flow['sampler_name'] + ' ' + flow['scheduler']
|
|
if 'seed' in flow and type(flow['seed']) == int:
|
|
FINAL_DICT['seed'] = flow['seed']
|
|
if 'cfg' in flow and (type(flow['cfg']) == int or type(flow['cfg']) == float):
|
|
FINAL_DICT['cfg_scale'] = flow['cfg']
|
|
|
|
if ModelID and 'ckpt_name' in prompt_json[ModelID]['inputs'] and type(prompt_json[ModelID]['inputs']['ckpt_name']) == str:
|
|
FINAL_DICT['model_name'] = prompt_json[ModelID]['inputs']['ckpt_name'] # flow['ckpt_name']
|
|
elif 'ckpt_name' in flow and type(flow['ckpt_name']) == str:
|
|
FINAL_DICT['model_name'] = flow['ckpt_name']
|
|
|
|
if SizeID and 'width' in prompt_json[SizeID]['inputs'] and 'height' in prompt_json[SizeID]['inputs'] and type(prompt_json[SizeID]['inputs']['width']) == int:
|
|
origwidth = str(prompt_json[SizeID]['inputs']['width'])
|
|
origheight = str(prompt_json[SizeID]['inputs']['height'])
|
|
FINAL_DICT['width'] = int(origwidth)
|
|
FINAL_DICT['height'] = int(origheight)
|
|
FINAL_DICT['size_string'] = origwidth + 'x' + origheight
|
|
|
|
self._parameter = FINAL_DICT
|
|
|
|
def _comfy_traverse(self, prompt, end_node):
|
|
flow = {}
|
|
node = [end_node]
|
|
inputs = {}
|
|
try:
|
|
inputs = prompt[end_node]["inputs"]
|
|
except:
|
|
print("node error")
|
|
return flow, node
|
|
match prompt[end_node]["class_type"]:
|
|
case node_type if node_type in SAVE_IMAGE_TYPE:
|
|
try:
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["images"][0])
|
|
flow = utility.merge_dict(flow, last_flow)
|
|
node += last_node
|
|
except:
|
|
print("comfyUI SaveImage error")
|
|
case node_type if node_type in KSAMPLER_TYPES:
|
|
try:
|
|
flow = inputs
|
|
last_flow1, last_node1 = self._comfy_traverse(prompt, inputs["model"][0])
|
|
last_flow2, last_node2 = self._comfy_traverse(prompt, inputs["latent_image"][0])
|
|
positive = self._comfy_traverse(prompt, inputs["positive"][0])
|
|
if isinstance(positive, str):
|
|
self._positive = positive
|
|
elif isinstance(positive, dict):
|
|
self._positive_sdxl.update(positive)
|
|
negative = self._comfy_traverse(prompt, inputs["negative"][0])
|
|
if isinstance(negative, str):
|
|
self._negative = negative
|
|
elif isinstance(negative, dict):
|
|
self._negative_sdxl.update(negative)
|
|
seed = None
|
|
# handle "CR Seed"
|
|
if inputs.get("seed") and isinstance(inputs.get("seed"), list):
|
|
seed = {"seed": self._comfy_traverse(prompt, inputs["seed"][0])}
|
|
elif inputs.get("noise_seed") and isinstance(inputs.get("noise_seed"), list):
|
|
seed = {
|
|
"noise_seed": self._comfy_traverse(prompt, inputs["noise_seed"][0])
|
|
}
|
|
if seed:
|
|
flow.update(seed)
|
|
flow = utility.merge_dict(flow, last_flow1)
|
|
flow = utility.merge_dict(flow, last_flow2)
|
|
node += last_node1 + last_node2
|
|
except:
|
|
print("comfyUI KSampler error")
|
|
case node_type if node_type in CLIP_TEXT_ENCODE_TYPE:
|
|
try:
|
|
match node_type:
|
|
case "CLIPTextEncode":
|
|
# SDXLPromptStyler
|
|
if isinstance(inputs["text"], list):
|
|
text = int(inputs["text"][0])
|
|
prompt_styler = self._comfy_traverse(prompt, str(text))
|
|
self._positive = prompt_styler[0]
|
|
self._negative = prompt_styler[1]
|
|
return
|
|
elif isinstance(inputs["text"], str):
|
|
return inputs.get("text")
|
|
case "CLIPTextEncodeSDXL":
|
|
# SDXLPromptStyler
|
|
self._is_sdxl = True
|
|
if isinstance(inputs["text_g"], list):
|
|
text_g = int(inputs["text_g"][0])
|
|
text_l = int(inputs["text_l"][0])
|
|
prompt_styler_g = self._comfy_traverse(prompt, str(text_g))
|
|
prompt_styler_l = self._comfy_traverse(prompt, str(text_l))
|
|
self._positive_sdxl["Clip G"] = prompt_styler_g[0]
|
|
self._positive_sdxl["Clip L"] = prompt_styler_l[0]
|
|
self._negative_sdxl["Clip G"] = prompt_styler_g[1]
|
|
self._negative_sdxl["Clip L"] = prompt_styler_l[1]
|
|
return
|
|
elif isinstance(inputs["text_g"], str):
|
|
return {
|
|
"Clip G": inputs.get("text_g"),
|
|
"Clip L": inputs.get("text_l"),
|
|
}
|
|
case "CLIPTextEncodeSDXLRefiner":
|
|
self._is_sdxl = True
|
|
if isinstance(inputs["text"], list):
|
|
# SDXLPromptStyler
|
|
text = int(inputs["text"][0])
|
|
prompt_styler = self._comfy_traverse(prompt, str(text))
|
|
self._positive_sdxl["Refiner"] = prompt_styler[0]
|
|
self._negative_sdxl["Refiner"] = prompt_styler[1]
|
|
return
|
|
elif isinstance(inputs["text"], str):
|
|
return {"Refiner": inputs.get("text")}
|
|
except:
|
|
print("comfyUI CLIPText error")
|
|
case "LoraLoader":
|
|
try:
|
|
flow = inputs
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["model"][0])
|
|
flow = utility.merge_dict(flow, last_flow)
|
|
node += last_node
|
|
except:
|
|
print("comfyUI LoraLoader error")
|
|
case node_type if node_type in CHECKPOINT_LOADER_TYPE:
|
|
try:
|
|
return inputs, node
|
|
except:
|
|
print("comfyUI CheckpointLoader error")
|
|
case node_type if node_type in VAE_ENCODE_TYPE:
|
|
try:
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["pixels"][0])
|
|
flow = utility.merge_dict(flow, last_flow)
|
|
node += last_node
|
|
except:
|
|
print("comfyUI VAE error")
|
|
case "ControlNetApplyAdvanced":
|
|
try:
|
|
positive = self._comfy_traverse(prompt, inputs["positive"][0])
|
|
if isinstance(positive, str):
|
|
self._positive = positive
|
|
elif isinstance(positive, dict):
|
|
self._positive_sdxl.update(positive)
|
|
negative = self._comfy_traverse(prompt, inputs["negative"][0])
|
|
if isinstance(negative, str):
|
|
self._negative = negative
|
|
elif isinstance(negative, dict):
|
|
self._negative_sdxl.update(negative)
|
|
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["image"][0])
|
|
flow = utility.merge_dict(flow, last_flow)
|
|
node += last_node
|
|
except:
|
|
print("comfyUI ControlNetApply error")
|
|
case "ImageScale":
|
|
try:
|
|
flow = inputs
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["image"][0])
|
|
flow = utility.merge_dict(flow, last_flow)
|
|
node += last_node
|
|
except:
|
|
print("comfyUI ImageScale error")
|
|
case "UpscaleModelLoader":
|
|
try:
|
|
return {"upscaler": inputs["model_name"]}
|
|
except:
|
|
print("comfyUI UpscaleLoader error")
|
|
case "ImageUpscaleWithModel":
|
|
try:
|
|
flow = inputs
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["image"][0])
|
|
model = self._comfy_traverse(prompt, inputs["upscale_model"][0])
|
|
flow = utility.merge_dict(flow, last_flow)
|
|
flow = utility.merge_dict(flow, model)
|
|
node += last_node
|
|
except:
|
|
print("comfyUI UpscaleModel error")
|
|
case "ConditioningCombine":
|
|
try:
|
|
last_flow1, last_node1 = self._comfy_traverse(prompt, inputs["conditioning_1"][0])
|
|
last_flow2, last_node2 = self._comfy_traverse(prompt, inputs["conditioning_2"][0])
|
|
flow = utility.merge_dict(flow, last_flow1)
|
|
flow = utility.merge_dict(flow, last_flow2)
|
|
node += last_node1 + last_node2
|
|
except:
|
|
print("comfyUI ConditioningCombine error")
|
|
# custom nodes
|
|
case "SDXLPromptStyler":
|
|
try:
|
|
return inputs.get("text_positive"), inputs.get("text_negative")
|
|
except:
|
|
print("comfyUI SDXLPromptStyler error")
|
|
case "CR Seed":
|
|
try:
|
|
return inputs.get("seed")
|
|
except:
|
|
print("comfyUI CR Seed error")
|
|
case _:
|
|
try:
|
|
last_flow = {}
|
|
last_node = []
|
|
if inputs.get("samples"):
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["samples"][0])
|
|
elif inputs.get("image") and isinstance(inputs.get("image"), list):
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["image"][0])
|
|
elif inputs.get("model"):
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["model"][0])
|
|
elif inputs.get("clip"):
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["clip"][0])
|
|
elif inputs.get("samples_from"):
|
|
last_flow, last_node = self._comfy_traverse(prompt, inputs["samples_from"][0])
|
|
elif inputs.get("conditioning"):
|
|
result = self._comfy_traverse(prompt, inputs["conditioning"][0])
|
|
if isinstance(result, str):
|
|
return result
|
|
elif isinstance(result, list):
|
|
last_flow, last_node = result
|
|
flow = utility.merge_dict(flow, last_flow)
|
|
node += last_node
|
|
except:
|
|
print("comfyUI bridging node error")
|
|
return flow, node |