diff --git a/scripts/iib/tool.py b/scripts/iib/tool.py old mode 100644 new mode 100755 index 171610e..1779b56 --- a/scripts/iib/tool.py +++ b/scripts/iib/tool.py @@ -394,7 +394,8 @@ def get_img_geninfo_txt_path(path: str): return txt_path 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): return is_img_created_by_comfyui(img) and img.info.get('parameters') @@ -405,9 +406,8 @@ def get_comfyui_exif_data(img: Image): return {} meta_key = '3' data: Dict[str, any] = json.loads(prompt) - for i in range(3, 32): + for i in data.keys(): try: - i = str(i) if data[i]["class_type"].startswith("KSampler"): meta_key = i break @@ -415,11 +415,13 @@ def get_comfyui_exif_data(img: Image): pass meta = {} KSampler_entry = data[meta_key]["inputs"] + #print(KSampler_entry) # for testing + # As a workaround to bypass parsing errors in the parser. # 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["Model"] = data[KSampler_entry["model"][0]]["inputs"]["ckpt_name"] + meta["Model"] = data[KSampler_entry["model"][0]]["inputs"].get("ckpt_name") meta["Source Identifier"] = "ComfyUI" def get_text_from_clip(idx: str) : inputs = data[idx]["inputs"] @@ -427,7 +429,13 @@ def get_comfyui_exif_data(img: Image): if isinstance(text, list): # type:CLIPTextEncode (NSP) mode:Wildcards text = data[text[0]]["inputs"]["text"] 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]) pos_prompt_arr = unique_by(parse_prompt(pos_prompt)["pos_prompt"]) return {