Merge pull request #772 from ermanitu/comfyUI_improvements

added compatibility with several comfyUI workflows
pull/773/merge
zanllp 2025-01-15 02:39:33 +08:00 committed by GitHub
commit a6925ab492
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 6 deletions

20
scripts/iib/tool.py Normal file → Executable file
View File

@ -394,7 +394,8 @@ def get_img_geninfo_txt_path(path: str):
return txt_path return txt_path
def is_img_created_by_comfyui(img: Image): def is_img_created_by_comfyui(img: Image):
return img.info.get('prompt') and img.info.get('workflow') prompt = img.info.get('prompt')
return prompt and (img.info.get('workflow') or ("class_type" in prompt)) # ermanitu
def is_img_created_by_comfyui_with_webui_gen_info(img: Image): def is_img_created_by_comfyui_with_webui_gen_info(img: Image):
return is_img_created_by_comfyui(img) and img.info.get('parameters') return is_img_created_by_comfyui(img) and img.info.get('parameters')
@ -405,9 +406,8 @@ def get_comfyui_exif_data(img: Image):
return {} return {}
meta_key = '3' meta_key = '3'
data: Dict[str, any] = json.loads(prompt) data: Dict[str, any] = json.loads(prompt)
for i in range(3, 32): for i in data.keys():
try: try:
i = str(i)
if data[i]["class_type"].startswith("KSampler"): if data[i]["class_type"].startswith("KSampler"):
meta_key = i meta_key = i
break break
@ -415,11 +415,13 @@ def get_comfyui_exif_data(img: Image):
pass pass
meta = {} meta = {}
KSampler_entry = data[meta_key]["inputs"] KSampler_entry = data[meta_key]["inputs"]
#print(KSampler_entry) # for testing
# As a workaround to bypass parsing errors in the parser. # As a workaround to bypass parsing errors in the parser.
# https://github.com/jiw0220/stable-diffusion-image-metadata/blob/00b8d42d4d1a536862bba0b07c332bdebb2a0ce5/src/index.ts#L130 # https://github.com/jiw0220/stable-diffusion-image-metadata/blob/00b8d42d4d1a536862bba0b07c332bdebb2a0ce5/src/index.ts#L130
meta["Steps"] = "Unknown" meta["Steps"] = KSampler_entry.get("steps", "Unknown")
meta["Sampler"] = KSampler_entry["sampler_name"] meta["Sampler"] = KSampler_entry["sampler_name"]
meta["Model"] = data[KSampler_entry["model"][0]]["inputs"]["ckpt_name"] meta["Model"] = data[KSampler_entry["model"][0]]["inputs"].get("ckpt_name")
meta["Source Identifier"] = "ComfyUI" meta["Source Identifier"] = "ComfyUI"
def get_text_from_clip(idx: str) : def get_text_from_clip(idx: str) :
inputs = data[idx]["inputs"] inputs = data[idx]["inputs"]
@ -427,7 +429,13 @@ def get_comfyui_exif_data(img: Image):
if isinstance(text, list): # type:CLIPTextEncode (NSP) mode:Wildcards if isinstance(text, list): # type:CLIPTextEncode (NSP) mode:Wildcards
text = data[text[0]]["inputs"]["text"] text = data[text[0]]["inputs"]["text"]
return text.strip() return text.strip()
pos_prompt = get_text_from_clip(KSampler_entry["positive"][0])
in_node = data[str(KSampler_entry["positive"][0])]
if in_node["class_type"] != "FluxGuidance":
pos_prompt = get_text_from_clip(KSampler_entry["positive"][0])
else:
pos_prompt = get_text_from_clip(in_node["inputs"]["conditioning"][0])
neg_prompt = get_text_from_clip(KSampler_entry["negative"][0]) neg_prompt = get_text_from_clip(KSampler_entry["negative"][0])
pos_prompt_arr = unique_by(parse_prompt(pos_prompt)["pos_prompt"]) pos_prompt_arr = unique_by(parse_prompt(pos_prompt)["pos_prompt"])
return { return {