Merge pull request #18 from LuqP2/fix/krea2-prompt-fallback

fix: recover Krea2 prompt metadata

metadata_utils.py

@@ -3,7 +3,7 @@ Compatibility shim for MetaHub Save Node metadata helpers.
 
 The implementation is kept in metadata_utils_impl.py; this module exposes the
 same public API expected by the save nodes while pinning the published node
-version to 1.2.0.
+version to 1.2.1.
 """
 
 try:
@@ -11,7 +11,7 @@ try:
 except ImportError:
     import metadata_utils_impl as _impl
 
-METAHUB_SAVE_NODE_VERSION = "1.2.0"
+METAHUB_SAVE_NODE_VERSION = "1.2.1"
 _impl.METAHUB_SAVE_NODE_VERSION = METAHUB_SAVE_NODE_VERSION
 
 for _name in dir(_impl):

metadata_utils_impl.py

@@ -14,7 +14,7 @@ from xml.sax.saxutils import escape as xml_escape
 import numpy as np
 from PIL import Image, PngImagePlugin
 
-METAHUB_SAVE_NODE_VERSION = "1.2.0"
+METAHUB_SAVE_NODE_VERSION = "1.2.1"
 
 try:
     import piexif

metahub_save_node.py

@@ -308,6 +308,11 @@ class MetaHubSaveNode:
                     return extracted_value
                 return default_value
 
+            def normalize_prompt_override(value):
+                if isinstance(value, bool):
+                    return None
+                return value
+
             def normalize_int(value, default_value):
                 try:
                     return int(value)
@@ -326,8 +331,10 @@ class MetaHubSaveNode:
             sampler_value = resolve_value(sampler_name, extracted.get("sampler_name"), "euler")
             scheduler_value = resolve_value(scheduler, extracted.get("scheduler"), "normal")
             model_name_value = resolve_value(model_name, extracted.get("model_name"), "")
-            positive_value = resolve_value(positive, extracted.get("positive"), "")
-            negative_value = resolve_value(negative, extracted.get("negative"), "")
+            positive_override = normalize_prompt_override(positive)
+            negative_override = normalize_prompt_override(negative)
+            positive_value = resolve_value(positive_override, extracted.get("positive"), "")
+            negative_value = resolve_value(negative_override, extracted.get("negative"), "")
             denoise_value = normalize_float(resolve_value(denoise, extracted.get("denoise"), 1.0), 1.0)
             vae_name_value = resolve_value(vae_name, extracted.get("vae_name"), "")
             metadata_fields = [
@@ -349,8 +356,8 @@ class MetaHubSaveNode:
                 "sampler_name": sampler_name,
                 "scheduler": scheduler,
                 "model_name": model_name,
-                "positive": positive,
-                "negative": negative,
+                "positive": positive_override,
+                "negative": negative_override,
                 "denoise": denoise,
                 "vae_name": vae_name,
             }

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "imagemetahub-comfyui-save"
 description = "Official companion node for Image MetaHub. Saves A1111/Civitai-compatible parameters plus extended Image MetaHub metadata (workflow JSON, SHA256 model hashes, VRAM peak/total, GPU device, generation time, and steps/sec) for reproducibility and benchmarking."
-version = "1.2.0"
+version = "1.2.1"
 license = {file = "LICENSE"} 
 # classifiers = [
 #     # For OS-independent nodes (works on all operating systems)

tests/test_metadata_utils.py

@@ -334,7 +334,7 @@ def test_extract_workflow_attribution_from_node_properties():
 
     assert attribution["token"] == "imhcrt_br_creator_workflow_v1_random"
     assert attribution["source"] == "metahub_save_node"
-    assert attribution["node_version"] == "1.2.0"
+    assert attribution["node_version"] == "1.2.1"
 
 
 def test_build_imh_metadata_includes_attribution_without_a1111_parameters():

tests/test_workflow_extractor.py

@@ -479,3 +479,62 @@ def test_workflow_extractor_resolves_custom_advanced_sampler_chain():
     assert data["model_name"] == "jibMixZIT_v10.safetensors"
     assert data["vae_name"] == "diffusion_pytorch_model.safetensors"
     assert {"seed", "steps", "sampler_name", "scheduler", "denoise", "positive", "model_name", "vae_name"}.isdisjoint(missing)
+
+
+def test_workflow_extractor_follows_krea2_false_switch_to_prompt_text():
+    expected = "A studio portrait with dramatic directional lighting"
+    prompt = {
+        "67": {"class_type": "PrimitiveBoolean", "inputs": {"value": False}},
+        "70": {
+            "class_type": "ComfySwitchNode",
+            "inputs": {"switch": ["67", 0], "on_false": ["71", 0], "on_true": ["61", 0]},
+        },
+        "71": {"class_type": "PrimitiveStringMultiline", "inputs": {"value": expected}},
+        "72": {
+            "class_type": "KSampler",
+            "inputs": {
+                "seed": 123,
+                "steps": 10,
+                "cfg": 1,
+                "sampler_name": "euler",
+                "scheduler": "simple",
+                "positive": ["73", 0],
+            },
+        },
+        "73": {"class_type": "CLIPTextEncode", "inputs": {"text": ["70", 0]}},
+    }
+
+    extracted, missing = WorkflowExtractor(prompt).extract()
+
+    assert extracted["positive"] == expected
+    assert "positive" not in missing
+
+
+def test_workflow_extractor_preserves_empty_selected_switch_branch():
+    prompt = {
+        "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive prompt"}},
+        "2": {"class_type": "PrimitiveStringMultiline", "inputs": {"value": "inactive negative prompt"}},
+        "3": {"class_type": "ConditioningZeroOut", "inputs": {}},
+        "4": {
+            "class_type": "ComfySwitchNode",
+            "inputs": {"switch": False, "on_false": ["3", 0], "on_true": ["2", 0]},
+        },
+        "5": {"class_type": "CLIPTextEncode", "inputs": {"text": ["4", 0]}},
+        "6": {
+            "class_type": "KSampler",
+            "inputs": {
+                "seed": 123,
+                "steps": 10,
+                "cfg": 1,
+                "sampler_name": "euler",
+                "scheduler": "simple",
+                "positive": ["1", 0],
+                "negative": ["5", 0],
+            },
+        },
+    }
+
+    extracted, missing = WorkflowExtractor(prompt).extract()
+
+    assert extracted["negative"] == ""
+    assert "negative" not in missing

workflow_extractor.py

@@ -553,7 +553,7 @@ class WorkflowExtractor:
     def _get_clip_text(self, node: Dict[str, Any]) -> Optional[str]:
         inputs = node.get("inputs", {})
         text = self._get_literal_input(inputs, "text")
-        if text is not None and str(text).strip():
+        if text is not None and not isinstance(text, bool) and str(text).strip():
             return str(text)
         text_from_conn = self._resolve_string_from_connection(inputs.get("text"))
         if text_from_conn is not None:
@@ -592,6 +592,13 @@ class WorkflowExtractor:
         if self._is_prompt_encoder(class_type):
             return self._get_clip_text(node)
 
+        if "switch" in class_lower and "on_false" in inputs and "on_true" in inputs:
+            switch_value = self._resolve_boolean_input(inputs.get("switch"))
+            branch = "on_true" if switch_value is True else "on_false"
+            value = self._resolve_string_from_connection_with_visited(inputs.get(branch), visited)
+            if value is not None:
+                return value
+
         if class_lower == "joinstrings" or "join strings" in class_lower:
             delimiter = self._get_literal_input(inputs, "delimiter")
             joiner = str(delimiter) if delimiter is not None else " "
@@ -606,7 +613,7 @@ class WorkflowExtractor:
 
         for key in ("text", "positive", "negative", "prompt", "string", "value"):
             value = self._get_literal_input(inputs, key)
-            if value is not None and str(value).strip():
+            if value is not None and not isinstance(value, bool) and str(value).strip():
                 return str(value)
 
         for key in ("positive", "text", "string1", "string2", "conditioning"):
@@ -621,6 +628,16 @@ class WorkflowExtractor:
 
         return None
 
+    def _resolve_boolean_input(self, value: Any) -> Optional[bool]:
+        if isinstance(value, bool):
+            return value
+        node_id = self._get_connection_node_id(value)
+        if not node_id:
+            return None
+        node = self._get_node(node_id)
+        literal = node.get("inputs", {}).get("value") if node else None
+        return literal if isinstance(literal, bool) else None
+
     def _resolve_string_from_connection_with_visited(self, conn: Any, visited: Set[str]) -> Optional[str]:
         node_id = self._get_connection_node_id(conn)
         if not node_id: