From 82e0a0a4ffde000955717a602cd4a2f9b742b309 Mon Sep 17 00:00:00 2001
From: Francisco Serradilla <ermanitu@yahoo.es>
Date: Tue, 14 Jan 2025 19:31:49 +0100
Subject: [PATCH] improved searching for KSampler node

---
 scripts/iib/tool.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/scripts/iib/tool.py b/scripts/iib/tool.py
index 7c21de0..1d4f741 100755
--- a/scripts/iib/tool.py
+++ b/scripts/iib/tool.py
@@ -394,7 +394,7 @@ 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') # ermanitu
+    return img.info.get('prompt') #and img.info.get('workflow')
 
 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 +405,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,13 +414,13 @@ def get_comfyui_exif_data(img: Image):
             pass
     meta = {}
     KSampler_entry = data[meta_key]["inputs"]
-    print(KSampler_entry)
+    #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"] = KSampler_entry["steps"]
     meta["Sampler"] = KSampler_entry["sampler_name"]
-    meta["Model"] = data[KSampler_entry["model"][0]]["inputs"].get("ckpt_name") # ermanitu
+    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"]
@@ -430,7 +429,6 @@ def get_comfyui_exif_data(img: Image):
             text = data[text[0]]["inputs"]["text"]
         return text.strip()
     
-    # added by ermanitu
     in_node = data[str(KSampler_entry["positive"][0])]
     if in_node["class_type"] != "FluxGuidance":
         pos_prompt = get_text_from_clip(KSampler_entry["positive"][0])