Fix ComfyUI prompt extraction for ImpactWildcardProcessor

scripts/iib/tool.py

@@ -487,48 +487,196 @@ def is_img_created_by_comfyui(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')
 
+
 def get_comfyui_exif_data(img: Image):
+    prompt = None
     if img.format == "PNG":
         prompt = img.info.get('prompt')
     elif img.format == "WEBP":
         exif = img.info.get("exif")
-        split = [x.decode("utf-8", errors="ignore") for x in exif.split(b"\x00")]
+        if exif:
-        prompt_str = find(split, lambda x: x.lower().startswith("prompt:"))
+            split = [x.decode("utf-8", errors="ignore") for x in exif.split(b"\x00")]
-        if prompt_str:
+            prompt_str = find(split, lambda x: x.lower().startswith("prompt:"))
-            prompt = prompt_str.split(":", 1)[1] if prompt_str else None
+            if prompt_str:
+                prompt = prompt_str.split(":", 1)[1]
+
     if not prompt:
         return {}
+
+    data: Dict[str, Any] = json.loads(prompt)
+
     meta_key = '3'
-    data: Dict[str, any] = json.loads(prompt)
     for i in data.keys():
         try:
             if data[i]["class_type"].startswith("KSampler"):
                 meta_key = i
                 break
-        except:
+        except Exception:
             pass
+
+    if meta_key not in data:
+        return {}
+
     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"] = KSampler_entry.get("steps", "Unknown")
-    meta["Sampler"] = KSampler_entry["sampler_name"]
+    meta["Sampler"] = KSampler_entry.get("sampler_name", "Unknown")
-    meta["Model"] = data[KSampler_entry["model"][0]]["inputs"].get("ckpt_name")
+    try:
+        meta["Model"] = data[KSampler_entry["model"][0]]["inputs"].get("ckpt_name")
+    except Exception:
+        meta["Model"] = None
     meta["Source Identifier"] = "ComfyUI"
-    def get_text_from_clip(idx: str) :
+
-        try:
+    text_key_priority = [
-            inputs = data[idx]["inputs"]
+        "populated_text",   # ImpactWildcardProcessor の最終展開結果
-            if "text" in inputs:
+        "text",
-                text = inputs["text"]
+        "prompt",
-            elif "t5xxl" in inputs:
+        "positive",
-                text = inputs["t5xxl"]
+        "negative",
+        "string",
+        "value",
+        "t5xxl",
+    ]
+    text_key_blacklist = {
+        "wildcard_text",            # 展開前テンプレート
+        "Select to add Wildcard",   # UI 用
+        "select_to_add_wildcard",
+        "template",
+        "pattern",
+    }
+
+    wildcard_patterns = [
+        re.compile(r"__[^_\n]+__"),
+        re.compile(r"\{[^{}\n]*\|[^{}\n]*\}"),
+    ]
+
+    def normalize_text(value):
+        if not isinstance(value, str):
+            return None
+        value = value.strip()
+        return value if value else None
+
+    def looks_unexpanded_wildcard(text: str) -> bool:
+        return any(p.search(text) for p in wildcard_patterns)
+
+    def get_node(node_id):
+        key = str(node_id)
+        return data.get(key) or data.get(node_id)
+
+    def get_best_text_from_inputs(inputs: dict):
+        if not isinstance(inputs, dict):
+            return ""
+
+        clean_candidates = []
+        wildcard_candidates = []
+
+        def add_candidate(value):
+            text = normalize_text(value)
+            if not text:
+                return
+            if looks_unexpanded_wildcard(text):
+                wildcard_candidates.append(text)
             else:
-                return ""
+                clean_candidates.append(text)
-            if isinstance(text, list): # type:CLIPTextEncode (NSP) mode:Wildcards
+
-                text = data[text[0]]["inputs"]["text"]
+        for key in text_key_priority:
-            return text.strip()
+            if key in text_key_blacklist:
+                continue
+            add_candidate(inputs.get(key))
+
+        if clean_candidates:
+            return clean_candidates[0]
+        if wildcard_candidates:
+            return wildcard_candidates[0]
+
+        for key, value in inputs.items():
+            if key in text_key_blacklist:
+                continue
+            add_candidate(value)
+
+        if clean_candidates:
+            return clean_candidates[0]
+        if wildcard_candidates:
+            return wildcard_candidates[0]
+        return ""
+
+    def extract_text_from_node(node: dict):
+        if not isinstance(node, dict):
+            return ""
+
+        class_type = node.get("class_type", "")
+        inputs = node.get("inputs", {}) or {}
+
+        # 明示的な特例: ImpactWildcardProcessor は populated_text を最優先
+        if class_type == "ImpactWildcardProcessor":
+            populated = normalize_text(inputs.get("populated_text"))
+            if populated:
+                return populated
+            # wildcard_text は意図的に返さない
+            return ""
+
+        return get_best_text_from_inputs(inputs)
+
+    def resolve_text_from_ref(ref, visited=None):
+        if visited is None:
+            visited = set()
+
+        node_id = ref[0] if isinstance(ref, list) and len(ref) >= 1 else ref
+        node_key = str(node_id)
+        if node_key in visited:
+            return ""
+        visited.add(node_key)
+
+        node = get_node(node_id)
+        if not isinstance(node, dict):
+            return ""
+
+        direct_text = extract_text_from_node(node)
+        if direct_text:
+            return direct_text
+
+        inputs = node.get("inputs", {}) or {}
+
+        # FluxGuidance の場合は conditioning を優先して辿る
+        if node.get("class_type") == "FluxGuidance":
+            conditioning = inputs.get("conditioning")
+            if isinstance(conditioning, list) and len(conditioning) >= 1:
+                resolved = resolve_text_from_ref(conditioning, visited)
+                if resolved:
+                    return resolved
+
+        # よく使う接続キーを優先
+        preferred_link_keys = [
+            "text",
+            "conditioning",
+            "positive",
+            "negative",
+            "prompt",
+            "string",
+            "value",
+        ]
+        for key in preferred_link_keys:
+            value = inputs.get(key)
+            if isinstance(value, list) and len(value) >= 1:
+                resolved = resolve_text_from_ref(value, visited)
+                if resolved:
+                    return resolved
+
+        # fallback: 全 input を走査
+        for value in inputs.values():
+            if isinstance(value, list) and len(value) >= 1:
+                resolved = resolve_text_from_ref(value, visited)
+                if resolved:
+                    return resolved
+
+        return ""
+
+    def get_text_from_clip(idx):
+        try:
+            return resolve_text_from_ref(idx)
         except Exception as e:
             print(e)
             return ""
@@ -542,36 +690,30 @@ def get_comfyui_exif_data(img: Image):
         for node_id, node_data in data.items():
             try:
                 class_type = node_data.get("class_type", "")
-                inputs = node_data.get("inputs", {})
+                if "CLIPTextEncode" in class_type or class_type == "ImpactWildcardProcessor":
-
+                    text = resolve_text_from_ref(node_id)
-                if "CLIPTextEncode" in class_type:
-                    text = inputs.get("text", "")
-                    if isinstance(text, list):
-                        text = data[text[0]]["inputs"].get("text", "")
                     if text:
                         all_prompts.append(text.strip())
             except Exception as e:
                 print(e)
                 pass
 
-        all_prompts_str = "\nBREAK\n".join(all_prompts) if all_prompts else ""
+        all_prompts_str = "\nBREAK\n".join(unique_by(all_prompts)) if all_prompts else ""
         pos_prompt = all_prompts_str
         neg_prompt = ""
     else:
-        in_node = data[str(KSampler_entry["positive"][0])]
+        positive_ref = KSampler_entry.get("positive")
-        if in_node["class_type"] != "FluxGuidance":
+        negative_ref = KSampler_entry.get("negative")
-            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 = get_text_from_clip(positive_ref) if positive_ref else ""
+        neg_prompt = get_text_from_clip(negative_ref) if negative_ref else ""
 
     pos_prompt_arr = unique_by(parse_prompt(pos_prompt)["pos_prompt"])
     return {
         "meta": meta,
         "pos_prompt": pos_prompt_arr,
         "pos_prompt_raw": pos_prompt,
-        "neg_prompt_raw" : neg_prompt
+        "neg_prompt_raw": neg_prompt
     }
 
 def comfyui_exif_data_to_str(data):